Refactor locations (version 2)

src/__init__.py

@@ -22,7 +22,7 @@
 from src import users
 from src import channels, containers
 from src import dispatcher, gamestate
-from src import decorators
+from src import decorators, locations
 from src import game_stats, handler, hooks, status, warnings, relay
 from src import reaper
 from src import gamejoin, pregame

src/functions.py

@@ -19,8 +19,7 @@
     "get_players", "get_all_players", "get_participants",
     "get_target", "change_role",
     "get_main_role", "get_all_roles", "get_reveal_role",
-    "match_role", "match_mode", "match_totem",
-    "get_players_in_location", "get_location"
+    "match_role", "match_mode", "match_totem"
     ]
 
 def get_players(var: Optional[GameState | PregameState], roles=None, *, mainroles=None) -> list[User]:
@@ -307,22 +306,3 @@ def match_totem(totem: str, scope: Optional[Iterable[str]] = None) -> Match[Loca
             filtered_matches.add(LocalTotem(totem_map[match], match))
 
     return Match(filtered_matches)
-
-def get_players_in_location(var: GameState, location: str) -> set[User]:
-    """ Get all players in a particular location.
-
-    :param var: Game state
-    :param location: Location to check
-    :return: All users present in the given location, or an empty set if the location is vacant
-    """
-    pl = get_players(var)
-    return {p for p, loc in var.locations.items() if loc == location and p in pl}
-
-def get_location(var: GameState, player: User) -> str:
-    """ Get the location this player is present in.
-
-    :param var: Game state
-    :param player: Player to check
-    :return: Location player is present in
-    """
-    return var.locations[player]

src/gamestate.py

@@ -76,7 +76,6 @@ def __init__(self, pregame_state: PregameState):
         self.next_phase: Optional[str] = None
         self.night_count: int = 0
         self.day_count: int = 0
-        self.locations: UserDict[User, str] = UserDict()
 
     def __init_subclass__(cls, **kwargs):
         super().__init_subclass__(**kwargs)

src/locations.py

@@ -0,0 +1,102 @@
+from __future__ import annotations
+
+from src import gamestate
+from src.containers import UserDict
+from src.users import User
+from src.events import Event, event_listener
+
+__all__ = ["Location", "VillageSquare", "Graveyard", "Forest",
+           "get_players_in_location", "get_location", "get_home",
+           "move_player", "move_player_home", "set_home"]
+
+# singleton cache of known locations; persisted between games
+LOCATION_CACHE: dict[str, Location] = {}
+
+# GameState extension to store location data for internal use
+# Other modules **must** use the location API exposed in __all__
+# instead of directly accessing these members
+class GameState(gamestate.GameState):
+    def __init__(self):
+        self.home_locations: UserDict[User, Location] = UserDict()
+        self.current_locations: UserDict[User, Location] = UserDict()
+
+class Location:
+    __slots__ = ("_name",)
+
+    def __init__(self, name: str):
+        self._name = name
+
+    def __new__(cls, name: str):
+        if name not in LOCATION_CACHE:
+            obj = super().__new__(cls)
+            obj.__init__(name)
+            LOCATION_CACHE[name] = obj
+
+        return LOCATION_CACHE[name]
+
+    @property
+    def name(self):
+        return self._name
+
+# default locations, always defined
+VillageSquare = Location("square")
+Graveyard = Location("graveyard")
+Forest = Location("forest")
+
+def get_players_in_location(var: GameState, location: Location) -> set[User]:
+    """ Get all players in a particular location.
+
+    :param var: Game state
+    :param location: Location to check
+    :return: All users present in the given location, or an empty set if the location is vacant
+    """
+    return {p for p, loc in var.current_locations.items() if loc is location}
+
+def get_location(var: GameState, player: User) -> Location:
+    """ Get the location this player is present in.
+
+    :param var: Game state
+    :param player: Player to check
+    :return: Location player is present in
+    """
+    return var.current_locations[player]
+
+def get_home(var: GameState, player: User) -> Location:
+    """ Get the player's home location.
+
+    :param var: Game state
+    :param player: Player to check
+    :return: Player's home location
+    """
+    return var.home_locations[player]
+
+def move_player(var: GameState, player: User, to: Location):
+    """ Move player to a new location.
+
+    :param var: Game state
+    :param player: Player to move
+    :param to: Player's new location
+    """
+    var.current_locations[player] = to
+
+def move_player_home(var: GameState, player: User):
+    """ Move player to their home location.
+
+    :param var: Game state
+    :param player: Player to move
+    """
+    var.current_locations[player] = var.home_locations[player]
+
+def set_home(var: GameState, player: User, home: Location):
+    """ Set player's home location.
+
+    :param var: Game state
+    :param player: Player to set
+    :param home: New home
+    """
+    var.home_locations[player] = home
+
+@event_listener("del_player")
+def on_del_player(evt: Event, var: GameState, player: User, allroles: set[str], death_triggers: bool):
+    del var.home_locations[:player:]
+    del var.current_locations[:player:]

src/pregame.py

@@ -23,6 +23,7 @@
 from src.users import User
 from src.dispatcher import MessageDispatcher
 from src.channels import Channel
+from src.locations import Location, set_home
 
 WAIT_TOKENS = 0
 WAIT_LAST = 0
@@ -31,7 +32,7 @@
 LAST_WAIT: UserDict[User, datetime] = UserDict()
 START_VOTES: UserSet = UserSet()
 CAN_START_TIME: datetime = datetime.now()
-FORCE_ROLES: DefaultUserDict[UserSet] = DefaultUserDict(UserSet)
+FORCE_ROLES: DefaultUserDict[str, UserSet] = DefaultUserDict(UserSet)
 
 @command("wait", playing=True, phases=("join",))
 def wait(wrapper: MessageDispatcher, message: str):
@@ -430,6 +431,12 @@ def _isvalid(mode, allow_vote_only):
         else:
             raise KeyError("Invalid action for role_attribution_end")
 
+    # set default location for each player to a unique house
+    for i, p in enumerate(get_players(ingame_state)):
+        home_event = Event("player_home", {"home": Location("house_{0}".format(i))})
+        home_event.dispatch(ingame_state, p)
+        set_home(ingame_state, p, home_event.data["home"])
+
     with locks.join_timer: # cancel timers
         for name in ("join", "join_pinger", "start_votes"):
             if name in TIMERS:
@@ -469,11 +476,11 @@ def _isvalid(mode, allow_vote_only):
     wrapper.send(messages[key].format(villagers, gamemode, options))
     wrapper.target.mode("+m")
 
-    if not ingame_state.start_with_day:
+    if start_event.data["custom_game_callback"]:
+        start_event.data["custom_game_callback"](ingame_state)
+    elif not ingame_state.start_with_day:
         from src.trans import transition_night
         transition_night(ingame_state)
-    elif start_event.data["custom_game_callback"]:
-        start_event.data["custom_game_callback"](ingame_state)
     else:
         # send role messages
         evt = Event("send_role", {})

src/roles/alphawolf.py

@@ -14,6 +14,7 @@
 from src.users import User
 from src.dispatcher import MessageDispatcher
 from src.gamestate import GameState
+from src.locations import get_home
 
 register_wolf("alpha wolf")
 
@@ -70,9 +71,9 @@ def on_night_kills(evt: Event, var: GameState):
         # simplify a lot of the code by offloading it to relevant pieces
         add_lycanthropy(var, target, "bitten")
         add_lycanthropy_scope(var, All)
-        house = var.players.index(target)
-        evt.data["victims"].add(f"house_{house}")
-        evt.data["killers"][f"house_{house}"].append("@wolves")
+        house = get_home(var, target)
+        evt.data["victims"].add(house)
+        evt.data["killers"][house].append("@wolves")
 
     # reset ENABLED here instead of begin_day so that night deaths can enable alpha wolf the next night
     ENABLED = False

src/roles/doomsayer.py

@@ -16,6 +16,7 @@
 from src.users import User
 from src.dispatcher import MessageDispatcher
 from src.gamestate import GameState
+from src.locations import move_player_home
 
 register_wolf("doomsayer")
 
@@ -111,6 +112,7 @@ def on_begin_day(evt: Event, var: GameState):
     for sick in SICK.values():
         status.add_absent(var, sick, "illness")
         status.add_silent(var, sick)
+        move_player_home(var, sick)
 
     # clear out LASTSEEN for people that didn't see last night
     for doom in list(LASTSEEN.keys()):

src/roles/harlot.py

@@ -10,11 +10,12 @@
 from src.decorators import command
 from src.dispatcher import MessageDispatcher
 from src.events import Event, event_listener
-from src.functions import get_players, get_all_players, get_main_role, get_reveal_role, get_target
+from src.functions import get_players, get_all_players, get_target
 from src.gamestate import GameState
 from src.messages import messages
 from src.status import try_misdirection, try_exchange
 from src.users import User
+from src.locations import move_player, get_home
 
 VISITED: UserDict[users.User, users.User] = UserDict()
 PASSED = UserSet()
@@ -43,9 +44,7 @@ def hvisit(wrapper: MessageDispatcher, message: str):
 
     VISITED[wrapper.source] = target
     PASSED.discard(wrapper.source)
-    house = var.players.index(target)
-    var.locations[wrapper.source] = f"house_{house}"
-
+    move_player(var, wrapper.source, get_home(var, target))
     wrapper.pm(messages["harlot_success"].format(target))
     if target is not wrapper.source:
         target.send(messages["harlot_success"].format(wrapper.source))

src/roles/helper/gunners.py

@@ -17,6 +17,7 @@
 from src.dispatcher import MessageDispatcher
 from src.gamestate import GameState
 from src.users import User
+from src.locations import move_player_home
 
 _rolestate: dict[str, dict[str, Any]] = {}
 
@@ -89,6 +90,7 @@ def shoot(wrapper: MessageDispatcher, message: str):
             else:
                 wrapper.send(messages["gunner_victim_injured"].format(target))
                 add_absent(var, target, "wounded")
+                move_player_home(var, target)
                 from src.votes import chk_decision
                 if not chk_win(var):
                     # game didn't immediately end due to injury, see if we should force through a vote
@@ -140,6 +142,8 @@ def on_del_player(evt: Event, var: GameState, victim: User, all_roles: set[str],
                         # shot hit, but didn't kill
                         channels.Main.send(messages["gunner_shoot_overnight_hit"].format(victim))
                         add_absent(var, shot, "wounded")
+                        # player will be moved back to home after daytime locations are fixed;
+                        # doing it here will simply get overwritten
                     else:
                         # shot was fired and missed
                         channels.Main.send(messages["gunner_shoot_overnight_missed"].format(victim))

src/roles/helper/shamans.py

@@ -513,6 +513,7 @@ def on_begin_day(evt: Event, var: GameState):
     # Apply totem effects that need to begin on day proper
     for player in NARCOLEPSY:
         status.add_absent(var, player, "totem")
+        move_player_home(var, player)
     for player in IMPATIENCE:
         status.add_force_vote(var, player, get_all_players(var) - {player})
     for player in PACIFISM:

src/roles/helper/wolves.py

@@ -16,6 +16,7 @@
 from src.dispatcher import MessageDispatcher
 from src.gamestate import GameState
 from src.users import User
+from src.locations import get_home
 
 KILLS: UserDict[users.User, UserList] = UserDict()
 
@@ -134,9 +135,9 @@ def on_night_kills(evt: Event, var: GameState):
             # wolfchat such as sorcerer or traitor, unlike main role wolves)
             for victim in victims:
                 if victim not in wolves:
-                    house = var.players.index(victim)
-                    evt.data["victims"].add(f"house_{house}")
-                    evt.data["killers"][f"house_{house}"].append(wolf)
+                    house = get_home(var, victim)
+                    evt.data["victims"].add(house)
+                    evt.data["killers"][house].append(wolf)
     # for wolves in wolfchat, determine who had the most kill votes and kill them,
     # choosing randomly in case of ties
     for i in range(total_kills):
@@ -150,10 +151,10 @@ def on_night_kills(evt: Event, var: GameState):
                 dups.append(v)
         if maxc and dups:
             target = random.choice(dups)
-            house = var.players.index(target)
-            evt.data["victims"].add(f"house_{house}")
+            house = get_home(var, target)
+            evt.data["victims"].add(house)
             # special key to let us know to randomly select a wolf in case of retribution totem
-            evt.data["killers"][f"house_{house}"].append("@wolves")
+            evt.data["killers"][house].append("@wolves")
             del found[target]
 
 @event_listener("retribution_kill")

src/roles/priest.py

@@ -13,6 +13,7 @@
 from src.trans import chk_win
 from src.dispatcher import MessageDispatcher
 from src.gamestate import GameState
+from src.locations import move_player, Graveyard
 
 PRIESTS = UserSet()
 
@@ -65,6 +66,7 @@ def consecrate(wrapper: MessageDispatcher, message: str):
 
     wrapper.pm(messages["consecrate_success"].format(target))
     add_absent(var, wrapper.source, "consecrating")
+    move_player(var, wrapper.source, Graveyard)
     from src.votes import chk_decision
     if not chk_win(var):
         # game didn't immediately end due to marking as absent, see if we should force through a lynch

src/roles/succubus.py

@@ -9,11 +9,12 @@
 from src.decorators import command
 from src.dispatcher import MessageDispatcher
 from src.events import Event, event_listener
-from src.functions import get_players, get_all_players, get_reveal_role, get_target
+from src.functions import get_players, get_all_players, get_target
 from src.gamestate import GameState
 from src.messages import messages
 from src.status import try_misdirection, try_exchange, is_dead
 from src.users import User
+from src.locations import move_player, get_home
 
 ENTRANCED = UserSet()
 VISITED: UserDict[users.User, users.User] = UserDict()
@@ -44,9 +45,7 @@ def hvisit(wrapper: MessageDispatcher, message: str):
 
     VISITED[wrapper.source] = target
     PASSED.discard(wrapper.source)
-    house = var.players.index(target)
-    var.locations[wrapper.source] = f"house_{house}"
-
+    move_player(var, wrapper.source, get_home(var, target))
     if target not in get_all_players(var, ("succubus",)):
         ENTRANCED.add(target)
         wrapper.send(messages["succubus_target_success"].format(target))

src/status/dying.py

@@ -6,7 +6,6 @@
 
 from src.containers import UserDict, UserSet
 from src.functions import get_main_role, get_all_roles, get_reveal_role
-from src.messages import messages
 from src.gamestate import GameState, PregameState
 from src.events import Event, event_listener
 from src.users import User
@@ -104,9 +103,6 @@ def kill_players(var: Optional[GameState | PregameState], *, end_game: bool = Tr
                     var.roles[role].remove(player)
                 dead.add(player)
                 DEAD.add(player)
-
-                # move their body to the graveyard
-                var.locations[player] = "graveyard"
             else:
                 # left during join phase
                 var.players.remove(player)