Use enums for faction, leave reason in summary

This should avoid any future failures to update magic strings
everywhere. I've put in custom (de)serialization annotations for
pydantic so that the JSON version of the summary uses strings, not ints.

shroudstone/renamer.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+from enum import Enum
 import logging
 import os
 import platform
@@ -19,6 +20,7 @@
 
 from shroudstone import __version__
 from shroudstone.config import data_dir
+from shroudstone.replay.parser import LeftGameReason
 from shroudstone.replay.summary import Player, ReplaySummary, summarize_replay
 from shroudstone.replay.versions import FRIGATE
 
@@ -246,9 +248,9 @@ def get_result(replay: Replay):
         return None
     if replay.summary.build_number >= FRIGATE:
         # Since Frigate we've had explicit surrender messages, so we rely on them alone for certainty:
-        if replay.us.leave_reason == "Surrender":
+        if replay.us.leave_reason == LeftGameReason.Surrender:
             return "loss"
-        if replay.them.leave_reason == "Surrender":
+        if replay.them.leave_reason == LeftGameReason.Surrender:
             return "win"
         return None
     else:
@@ -263,6 +265,12 @@ def get_result(replay: Replay):
             return "win"
 
 
+def _cap(x: Optional[Enum]):
+    if x is None:
+        return ""
+    return x.name.capitalize()
+
+
 def new_name_for(replay: Replay, format_1v1: str, format_generic: str) -> str:
     parts = {}
     parts["map_name"] = replay.summary.map_name
@@ -283,8 +291,8 @@ def new_name_for(replay: Replay, format_1v1: str, format_generic: str) -> str:
         parts["us"] = parts["p1"] = us.nickname
         parts["them"] = parts["p2"] = them.nickname
 
-        parts["r1"] = parts["f1"] = (us.faction or "").capitalize()
-        parts["r2"] = parts["f2"] = (them.faction or "").capitalize()
+        parts["r1"] = parts["f1"] = _cap(us.faction)
+        parts["r2"] = parts["f2"] = _cap(them.faction)
 
         result = get_result(replay)
         parts["result"] = (result or "unknown").capitalize()
@@ -295,7 +303,7 @@ def new_name_for(replay: Replay, format_1v1: str, format_generic: str) -> str:
             p.nickname.capitalize() for p in replay.summary.players
         )
         parts["players_with_factions"] = ", ".join(
-            f"{p.nickname.capitalize()} {(p.faction or '').upper():.1}"
+            f"{p.nickname.capitalize()} {_cap(p.faction).upper():.1}"
             for p in replay.summary.players
         )
         newname = format_generic.format(**parts)

shroudstone/replay/summary.py

@@ -1,14 +1,17 @@
 """Summarize Stormgate replays into a JSON schema."""
+
 from __future__ import annotations
 
 import logging
+from enum import Enum
 from pathlib import Path
-from typing import BinaryIO, List, Optional, Union
+from typing import BinaryIO, List, Optional, Type, Union
 from uuid import UUID
 
-from pydantic import BaseModel
+from pydantic import BaseModel, BeforeValidator, PlainSerializer
+from typing_extensions import Annotated, TypeVar
 
-from .parser import GameState, get_build_number
+from .parser import Faction, GameState, LeftGameReason, get_build_number
 
 logger = logging.getLogger(__name__)
 
@@ -17,6 +20,25 @@
 REPLAY_TIMESTAMP_UNIT = 1 / 1024
 
 
+T = TypeVar("T", bound=Enum)
+
+
+def _from_name(t: Type[T]):
+    def get(n: Union[str, T]) -> T:
+        if isinstance(n, Enum):
+            return n
+        return getattr(t, n)
+
+    return get
+
+
+def _validator(t: Type[T]):
+    return BeforeValidator(_from_name(t))
+
+
+_serialize = PlainSerializer(lambda x: x.name, when_used="json-unless-none")
+
+
 class InvalidGameState(Exception):
     """Replay parsed successfully, but we don't understand the final gamestate."""
 
@@ -31,10 +53,12 @@ class Player(BaseModel):
     nickname: str
     nickname_discriminator: Optional[str] = None
     uuid: Optional[UUID] = None
-    faction: Optional[str] = None
+    faction: Annotated[Optional[Faction], _validator(Faction), _serialize]
     is_ai: bool = False
     disconnect_time: Optional[float] = None
-    leave_reason: str = "unknown"
+    leave_reason: Annotated[LeftGameReason, _validator(LeftGameReason), _serialize] = (
+        LeftGameReason.Unknown
+    )
 
 
 class ReplaySummary(BaseModel):
@@ -67,7 +91,7 @@ def summarize_replay(replay: Union[Path, BinaryIO]) -> ReplaySummary:
                 Player(
                     nickname=slot.ai_type.name,
                     is_ai=True,
-                    faction=slot.faction.name,
+                    faction=slot.faction,
                 )
             )
         elif slot.client_id is not None:
@@ -82,7 +106,7 @@ def summarize_replay(replay: Union[Path, BinaryIO]) -> ReplaySummary:
                     nickname_discriminator=client.discriminator,
                     uuid=client.uuid,
                     is_ai=False,
-                    faction=slot.faction.name,
+                    faction=slot.faction,
                 )
             )
             if (
@@ -92,7 +116,7 @@ def summarize_replay(replay: Union[Path, BinaryIO]) -> ReplaySummary:
                 p.disconnect_time = (
                     client.left_game_time - state.game_started_time
                 ) * REPLAY_TIMESTAMP_UNIT
-            p.leave_reason = client.left_game_reason.name
+            p.leave_reason = client.left_game_reason
     for client in state.clients.values():
         if client.slot_number != 255:
             raise InvalidGameState("Player not in a slot but slot_number != 255?")

tests/replays/c6b4eb4e-4994-4e96-b098-3e1953a02033/CL44821-2024.02.03-18.48.json

@@ -1 +1 @@
-{"build_number": 44821, "map_name": "Boneyard", "players": [{"nickname": "Pox", "nickname_discriminator": "8082", "uuid": "c6b4eb4e-4994-4e96-b098-3e1953a02033", "faction": "Infernals", "is_ai": false, "disconnect_time": 10.94140625, "leave_reason": "Unknown"}, {"nickname": "PeacefulBot", "nickname_discriminator": null, "uuid": null, "faction": "Vanguard", "is_ai": true, "disconnect_time": null, "leave_reason": "unknown"}], "spectators": [], "duration_seconds": 10.94140625, "is_1v1_ladder_game": false}
+{"build_number": 44821, "map_name": "Boneyard", "players": [{"nickname": "Pox", "nickname_discriminator": "8082", "uuid": "c6b4eb4e-4994-4e96-b098-3e1953a02033", "faction": "Infernals", "is_ai": false, "disconnect_time": 10.94140625, "leave_reason": "Unknown"}, {"nickname": "PeacefulBot", "nickname_discriminator": null, "uuid": null, "faction": "Vanguard", "is_ai": true, "disconnect_time": null, "leave_reason": "Unknown"}], "spectators": [], "duration_seconds": 10.94140625, "is_1v1_ladder_game": false}

tests/replays/c6b4eb4e-4994-4e96-b098-3e1953a02033/CL44821-2024.02.13-22.25.json

@@ -1 +1 @@
-{"build_number": 44821, "map_name": "JaggedMaw", "players": [{"nickname": "Pox", "nickname_discriminator": "8082", "uuid": "c6b4eb4e-4994-4e96-b098-3e1953a02033", "faction": "Infernals", "is_ai": false, "disconnect_time": 256.95703125, "leave_reason": "Unknown"}, {"nickname": "PeacefulBot", "nickname_discriminator": null, "uuid": null, "faction": "Vanguard", "is_ai": true, "disconnect_time": null, "leave_reason": "unknown"}], "spectators": [], "duration_seconds": 256.95703125, "is_1v1_ladder_game": false}
+{"build_number": 44821, "map_name": "JaggedMaw", "players": [{"nickname": "Pox", "nickname_discriminator": "8082", "uuid": "c6b4eb4e-4994-4e96-b098-3e1953a02033", "faction": "Infernals", "is_ai": false, "disconnect_time": 256.95703125, "leave_reason": "Unknown"}, {"nickname": "PeacefulBot", "nickname_discriminator": null, "uuid": null, "faction": "Vanguard", "is_ai": true, "disconnect_time": null, "leave_reason": "Unknown"}], "spectators": [], "duration_seconds": 256.95703125, "is_1v1_ladder_game": false}