From 04f51a62287523fc4d1a961fd2d6f1a98601c272 Mon Sep 17 00:00:00 2001 From: Anthony Carapetis Date: Sat, 17 Aug 2024 15:39:32 +1000 Subject: [PATCH] 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 | 18 ++++++--- shroudstone/replay/summary.py | 40 +++++++++++++++---- .../CL44821-2024.02.03-18.48.json | 2 +- .../CL44821-2024.02.13-22.25.json | 2 +- 4 files changed, 47 insertions(+), 15 deletions(-) diff --git a/shroudstone/renamer.py b/shroudstone/renamer.py index cd5ca16..2e67348 100755 --- a/shroudstone/renamer.py +++ b/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) diff --git a/shroudstone/replay/summary.py b/shroudstone/replay/summary.py index 910af8f..615b7a4 100644 --- a/shroudstone/replay/summary.py +++ b/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?") diff --git a/tests/replays/c6b4eb4e-4994-4e96-b098-3e1953a02033/CL44821-2024.02.03-18.48.json b/tests/replays/c6b4eb4e-4994-4e96-b098-3e1953a02033/CL44821-2024.02.03-18.48.json index bb4b175..50125fe 100644 --- a/tests/replays/c6b4eb4e-4994-4e96-b098-3e1953a02033/CL44821-2024.02.03-18.48.json +++ b/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} \ No newline at end of file +{"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} \ No newline at end of file diff --git a/tests/replays/c6b4eb4e-4994-4e96-b098-3e1953a02033/CL44821-2024.02.13-22.25.json b/tests/replays/c6b4eb4e-4994-4e96-b098-3e1953a02033/CL44821-2024.02.13-22.25.json index fd673cf..5f81b2d 100644 --- a/tests/replays/c6b4eb4e-4994-4e96-b098-3e1953a02033/CL44821-2024.02.13-22.25.json +++ b/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} \ No newline at end of file +{"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} \ No newline at end of file