diff --git a/main.py b/main.py index 12a7c7c..29ad6ab 100755 --- a/main.py +++ b/main.py @@ -44,6 +44,21 @@ _ch.setFormatter(LogFormatter()) LOG.addHandler(_ch) +TEAMS = {} +_logged_teams = [] + + +def team(player: str) -> str: + """Maps a username to a defined team.""" + if player in TEAMS: + return TEAMS[player] + else: + if not player in _logged_teams: + LOG.warning(f"missing team mapping for {player}") + _logged_teams.append(player) + return player + + class safelist(list): def get(self, index, default=None): try: @@ -133,7 +148,7 @@ def parse_log(game: str, log: str, into: sqlite3.Connection): VALUES (?, ?, ?, ?, ?, ?) ON CONFLICT DO NOTHING """, - (game, turn, player, move, user, target), + (game, turn, team(player), move, user, target), ) case ["switch", name, specie, *rest]: player, name = resolve_mon(name) @@ -143,7 +158,7 @@ def parse_log(game: str, log: str, into: sqlite3.Connection): VALUES (?, ?, ?, ?) ON CONFLICT DO NOTHING """, - (game, turn, player, name), + (game, turn, team(player), name), ) conn.execute( """ @@ -151,7 +166,7 @@ def parse_log(game: str, log: str, into: sqlite3.Connection): VALUES(?, ?, ?, ?) ON CONFLICT DO NOTHING """, - (game, player, name, specie.split(", ")[0]), + (game, team(player), name, specie.split(", ")[0]), ) case ["faint", mon]: player, mon = resolve_mon(mon) @@ -161,7 +176,7 @@ def parse_log(game: str, log: str, into: sqlite3.Connection): VALUES(?, ?, ?, ?) ON CONFLICT DO NOTHING """, - (game, turn, player, mon), + (game, turn, team(player), mon), ) case ["win", player]: conn.execute( @@ -170,7 +185,7 @@ def parse_log(game: str, log: str, into: sqlite3.Connection): SET winner = ? WHERE id = ? """, - (player, game), + (team(player), game), ) case ["-sidestart", side, env]: if not last_move: @@ -223,7 +238,15 @@ def parse_log(game: str, log: str, into: sqlite3.Connection): VALUES(?, ?, ?, ?, ?, ?, ?) ON CONFLICT DO NOTHING """, - (game, turn, player, mon, source, source_user, source_player), + ( + game, + turn, + team(player), + mon, + source, + source_user, + team(source_player), + ), ) case _: # LOG.debug(f"unhandled message {chunks[0]}") @@ -281,6 +304,14 @@ def main(args): action="store_true", help="fetch replays instead of using cache", ) + parser.add_argument( + "-t", + "--teams", + action="store", + metavar="FILE", + default="teams.json", + help="JSON file defining players to teams", + ) parser.add_argument( "-o", "--output", @@ -297,6 +328,11 @@ def main(args): elif args.verbose: LOG.setLevel(logging.DEBUG) + if args.teams: + with open(args.teams) as f: + global TEAMS + TEAMS = json.load(f) + try: db = sqlite3.connect(args.output) _init_db(db) @@ -315,7 +351,13 @@ def main(args): VALUES (?, ?, ?, ?, ?) ON CONFLICT DO NOTHING """, - (replay.id, replay.p1, replay.p2, replay.format, replay.uploadtime), + ( + replay.id, + team(replay.p1), + team(replay.p2), + replay.format, + replay.uploadtime, + ), ) parse_log(replay.id, replay.log, into=db)