This repository has been archived by the owner on Jun 12, 2021. It is now read-only.
-
Notifications
You must be signed in to change notification settings - Fork 0
/
aggregation.py
67 lines (59 loc) · 1.49 KB
/
aggregation.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
#!/usr/bin/env python3
import chess
import chess.pgn
import pandas as pd
from typing import List, Tuple
"""
This module pulls the games from pgn files, compiles these games into a
dataframe with meta data.
"""
COLUMNS = [
# 'Event',
# 'Site',
'Date',
# 'Round',
'White',
'Black',
# 'Result',
'BlackElo',
# 'BlackRatingDiff',
# 'ECO',
# 'Termination',
# 'TimeControl',
# 'UTCDate',
# 'UTCTime',
# 'Variant',
'WhiteElo',
# 'WhiteRatingDiff',
]
def load_pgn(filename: str) -> List[chess.pgn.Game]:
"""
Read pgn into list, breaks after hitting EOF
"""
with open(filename) as pgn:
games = []
while True:
game = chess.pgn.read_game(pgn)
if game is not None:
games.append(game)
else:
break
return games
def aggregate(filename: str) -> Tuple[pd.DataFrame, List[chess.pgn.Game]]:
with open(filename) as pgn:
rows = []
games = []
while True:
game = chess.pgn.read_game(pgn)
if game is None:
break
row = [dict(game.headers).get(x) for x in COLUMNS]
rows.append(row)
games.append(game)
# rows.append(list(dict(game.headers).get(x) for x in COLUMNS))
df = pd.DataFrame(rows, columns=COLUMNS)
return df, games
if __name__ == "__main__":
from pandasgui import show
df, games = aggregate("./pgn/Dae.pgn")
# show(df)