resolve teams

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)