First passthrough of locations

src/functions.py

@@ -20,7 +20,6 @@
     "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"
     ]
 
 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

@@ -1,10 +1,12 @@
 from __future__ import annotations
 
-import copy
 from typing import Any, Optional, TYPE_CHECKING
+
+import copy
 import time
 
 from src.containers import UserSet, UserDict, UserList
+from src.locations import Location, Square, Graveyard, House, Reason
 from src.messages import messages
 from src.cats import All
 from src import config
@@ -57,6 +59,7 @@ def __init__(self, pregame_state: PregameState):
         self.setup_started: bool = False
         self.setup_completed: bool = False
         self._torndown: bool = False
+        self.tearing_down: bool = False
         self.current_mode: GameMode = pregame_state.current_mode
         self.game_settings: dict[str, Any] = {}
         self.game_id: float = pregame_state.game_id
@@ -71,7 +74,9 @@ 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()
+        self.village_square = Square(self)
+        self.graveyard = Graveyard(self)
+        self._locations: set[Location] = {self.village_square, self.graveyard}
 
     def begin_setup(self):
         if self.setup_completed:
@@ -91,15 +96,26 @@ def finish_setup(self):
         assert not self._original_roles and not self._original_main_roles
         self._original_roles = copy.deepcopy(self.roles)
         self._original_main_roles = self.main_roles.copy()
+        for i, player in enumerate(self.players):
+            house = House(self, player, i)
+            house.users[player] = (Reason.home, None)
+            self._locations.add(house)
         self.setup_completed = True
 
     def teardown(self):
-        self.roles.clear()
-        self._original_roles.clear()
-        self._original_main_roles.clear()
-        self._rolestats.clear()
-        self.current_mode.teardown()
-        self._torndown = True
+        assert not self._torndown, "cannot tear down already torn-down GameState"
+        self.tearing_down = True
+        try:
+            self.roles.clear()
+            self._original_roles.clear()
+            self._original_main_roles.clear()
+            self._rolestats.clear()
+            self.current_mode.teardown()
+            self._torndown = True
+            for loc in self._locations:
+                loc.teardown()
+        finally:
+            self.tearing_down = False
 
     def _get_value(self, key: str) -> Any:
         # we don't actually need to complete setup before this can be used
@@ -115,6 +131,41 @@ def _get_value(self, key: str) -> Any:
     def in_game(self):
         return self.setup_completed and not self._torndown
 
+    @property
+    def locations(self) -> dict[User, Location]:
+        value = {}
+        for loc in self._locations:
+            for user in loc.users:
+                value[user] = loc
+        return value
+
+    def get_user_location(self, user: User):
+        for x in self._locations:
+            if user in x.users:
+                return (x,) + x.users[user]
+        raise ValueError(f"User {user} is not anywhere")
+
+    def set_user_location(self, user: User, loc: Location, reason: Reason | None = None, key: str | None = None):
+        if user not in loc.users:
+            for x in self._locations:
+                if user in x.users:
+                    old_r, old_k = x.users.pop(user)
+                    if reason is None:
+                        reason = old_r
+                        key = old_k
+                    loc.users[user] = (reason, key)
+                    break
+            else:
+                raise RuntimeError(f"Failed setting user {user} to location {loc}")
+
+    def find_house(self, user: User):
+        for x in self._locations:
+            if not isinstance(x, House):
+                continue
+            if x.owner is user:
+                return x
+        raise ValueError(f"Could not find house for {user}")
+
     def begin_phase_transition(self, phase: str):
         if self.next_phase is not None:
             raise RuntimeError("already in phase transition")

src/locations.py

@@ -0,0 +1,72 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from enum import Enum
+
+from src.users import User
+from src.containers import UserDict, UserList
+
+if TYPE_CHECKING:
+    from src.gamestate import GameState
+
+class Reason(Enum):
+    home = 0
+    day = 1
+    visiting = 2
+    killing = 3
+    dead = 4
+    special = 5
+    prison = 6
+
+class Location:
+    """Base class for locations."""
+    def __init__(self, var: GameState, name: str):
+        self._gs = var
+        self.name = name
+        self.users: UserDict[User, tuple[Reason, str | None]] = UserDict()
+
+    def __contains__(self, item):
+        return item in self.users
+
+    def __iter__(self):
+        return iter(self.users)
+
+    def __len__(self):
+        return len(self.users)
+
+    def __str__(self):
+        return self.name
+
+    def __hash__(self):
+        return hash(self.name)
+
+    def _teardown(self):
+        """Method for subclasses to define if needed."""
+
+    def teardown(self):
+        assert self._gs.tearing_down, "cannot tear down locations from outside of GameState.teardown()"
+        self._teardown()
+        self.name = "<DELETED>"
+        self.users.clear()
+
+class Square(Location):
+    def __init__(self, var: GameState):
+        super().__init__(var, "Village Square")
+
+class Graveyard(Location):
+    def __init__(self, var: GameState):
+        super().__init__(var, "Graveyard")
+
+class House(Location):
+    def __init__(self, var: GameState, player: User, pos: int):
+        super().__init__(var, f"{player.account}'s house")
+        self._owner = UserList([player])
+        self.pos = pos
+
+    @property
+    def owner(self) -> User:
+        return self._owner[0]
+
+    def _teardown(self):
+        self._owner.clear()

src/roles/doomsayer.py

@@ -8,6 +8,7 @@
 from src.cats import All
 from src.containers import UserSet, UserDict
 from src.decorators import command
+from src.locations import Reason
 from src.events import Event, event_listener
 from src.functions import get_all_players, get_target
 from src.messages import messages
@@ -109,7 +110,7 @@ def on_transition_night_end(evt: Event, var: GameState):
 @event_listener("begin_day")
 def on_begin_day(evt: Event, var: GameState):
     for sick in SICK.values():
-        status.add_absent(var, sick, "illness")
+        var.set_user_location(sick, var.find_house(sick), Reason.prison, "illness")
         status.add_silent(var, sick)
 
     # clear out LASTSEEN for people that didn't see last night

src/roles/harlot.py

@@ -9,6 +9,7 @@
 from src.containers import UserSet, UserDict
 from src.decorators import command
 from src.dispatcher import MessageDispatcher
+from src.locations import Reason
 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.gamestate import GameState
@@ -43,8 +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}"
+    var.set_user_location(wrapper.source, var.find_house(target), Reason.visiting)
 
     wrapper.pm(messages["harlot_success"].format(target))
     if target is not wrapper.source:

src/roles/helper/gunners.py

@@ -9,10 +9,11 @@
 from src.cats import Wolf, Killer
 from src.containers import UserDict
 from src.decorators import command
+from src.locations import Reason
 from src.events import Event, event_listener
 from src.functions import get_players, get_all_players, get_target, get_main_role, get_reveal_role
 from src.messages import messages
-from src.status import try_misdirection, try_exchange, add_dying, kill_players, add_absent, try_protection, is_dying
+from src.status import try_misdirection, try_exchange, add_dying, kill_players, try_protection, is_dying
 from src.trans import chk_win
 from src.dispatcher import MessageDispatcher
 from src.gamestate import GameState
@@ -88,7 +89,7 @@ def shoot(wrapper: MessageDispatcher, message: str):
                     kill_players(var)
             else:
                 wrapper.send(messages["gunner_victim_injured"].format(target))
-                add_absent(var, target, "wounded")
+                var.set_user_location(target, var.find_house(target), Reason.prison, "wounded")
                 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
@@ -139,7 +140,7 @@ def on_del_player(evt: Event, var: GameState, victim: User, all_roles: set[str],
                     elif event.data["hit"]:
                         # shot hit, but didn't kill
                         channels.Main.send(messages["gunner_shoot_overnight_hit"].format(victim))
-                        add_absent(var, shot, "wounded")
+                        var.set_user_location(shot, var.find_house(shot), Reason.prison, "wounded")
                     else:
                         # shot was fired and missed
                         channels.Main.send(messages["gunner_shoot_overnight_missed"].format(victim))

src/roles/helper/shamans.py

@@ -9,6 +9,7 @@
 from src import channels, users, status
 from src.cats import All, Wolf, Killer
 from src.containers import UserList, UserSet, UserDict, DefaultUserDict
+from src.locations import Reason
 from src.events import Event, event_listener
 from src.functions import (get_players, get_all_players, get_main_role, get_all_roles, get_reveal_role, get_target,
                            match_totem)
@@ -512,7 +513,7 @@ def on_transition_night_end(evt: Event, var: GameState):
 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")
+        var.set_user_location(player, var.find_house(player), Reason.prison, "totem")
     for player in IMPATIENCE:
         status.add_force_vote(var, player, get_all_players(var) - {player})
     for player in PACIFISM:

src/roles/priest.py

@@ -8,8 +8,9 @@
 from src.decorators import command
 from src.events import Event, event_listener
 from src.functions import get_players, get_all_players, get_target
+from src.locations import Reason
 from src.messages import messages
-from src.status import try_misdirection, try_exchange, add_absent
+from src.status import try_misdirection, try_exchange
 from src.trans import chk_win
 from src.dispatcher import MessageDispatcher
 from src.gamestate import GameState
@@ -64,7 +65,7 @@ def consecrate(wrapper: MessageDispatcher, message: str):
     evt.dispatch(var, wrapper.source, target)
 
     wrapper.pm(messages["consecrate_success"].format(target))
-    add_absent(var, wrapper.source, "consecrating")
+    var.set_user_location(wrapper.source, var.graveyard, Reason.special, "consecrating")
     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

@@ -8,6 +8,7 @@
 from src.containers import UserSet, UserDict
 from src.decorators import command
 from src.dispatcher import MessageDispatcher
+from src.locations import Reason
 from src.events import Event, event_listener
 from src.functions import get_players, get_all_players, get_reveal_role, get_target
 from src.gamestate import GameState
@@ -44,8 +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}"
+    var.set_user_location(wrapper.source, var.find_house(target), Reason.visiting)
 
     if target not in get_all_players(var, ("succubus",)):
         ENTRANCED.add(target)