Store objects in a subclass

orochi/bot.py

@@ -1,161 +1,16 @@
-from dataclasses import dataclass, field
-from datetime import datetime
-from enum import Enum
-import pickle
-
 import disnake
 from disnake import CategoryChannel, PermissionOverwrite, TextChannel
 from disnake.ext import commands
 import logging
 
 from orochi.config import Config
+from orochi.models import Game
 
 bot = commands.Bot(command_prefix='!')
-GAME: "Game"
-
-
-class Room(Enum):
-    A = 'A'
-    B = 'B'
-    C = 'C'
-
-
-class Vote(Enum):
-    ALLY = 'A'
-    BETRAY = 'B'
-
-
-@dataclass(frozen=True)
-class Player:
-    name: str
-    private_channel_id: int = field(hash=False)
-
-    @property
-    def round_votes(self):
-        for r in GAME.rounds:
-            for room in r.rooms:
-                for vote in room.votes:
-                    if self in vote.players:
-                        yield vote
-
-    @property
-    def score(self):
-        s = 3
-
-        for vote in self.round_votes:
-            room = vote.room
-            other_vote = room.vote1 if room.vote1 is not vote else room.vote2
-            match vote.vote, other_vote.vote:
-                case Vote.ALLY, Vote.ALLY:
-                    s += 2
-                case Vote.ALLY, Vote.BETRAY:
-                    s -= 2
-                case Vote.BETRAY, Vote.ALLY:
-                    s += 3
-                case Vote.BETRAY, Vote.BETRAY:
-                    pass
-
-        return s
-
-
-@dataclass
-class RoundVote:
-    player1: Player
-    player2: Player | None = None
-    vote: Vote | None = None
-    timestamp: datetime | None = None
-
-    @property
-    def players(self):
-        return self.player1, self.player2
-
-    @property
-    def room(self):
-        for r in GAME.rounds:
-            for room in r.rooms:
-                if self in room.votes:
-                    return room
-
-
-@dataclass
-class RoundRoom:
-    room: Room
-    vote1: RoundVote
-    vote2: RoundVote
-
-    @property
-    def votes(self):
-        return self.vote1, self.vote2
-
-    @property
-    def round(self):
-        for r in GAME.rounds:
-            if self in r.rooms:
-                return r
-
-
-@dataclass
-class Round:
-    round: int
-    room_a: RoundRoom
-    room_b: RoundRoom
-    room_c: RoundRoom
-
-    @property
-    def rooms(self):
-        return self.room_a, self.room_b, self.room_c
-
-
-@dataclass
-class Game:
-    rounds: list[Round] = field(default_factory=list)
-    players: dict[str, Player] = field(default_factory=dict)
-
-    def register_player(self, name: str, vote_channel_id: int) -> Player:
-        player = Player(name, vote_channel_id)
-        self.players[name] = player
-        return player
-
-    def default_first_round(self) -> Round:
-        return Round(
-            round=1,
-            room_a=RoundRoom(room=Room.A,
-                             vote1=RoundVote(player1=GAME.players['Tora']),
-                             vote2=RoundVote(player1=GAME.players['Kamui'],
-                                             player2=GAME.players['Philia'])),
-            room_b=RoundRoom(room=Room.A,
-                             vote1=RoundVote(player1=GAME.players['Dan']),
-                             vote2=RoundVote(player1=GAME.players['Ennea'],
-                                             player2=GAME.players['Delphine'])),
-            room_c=RoundRoom(room=Room.A,
-                             vote1=RoundVote(player1=GAME.players['Hanabi']),
-                             vote2=RoundVote(player1=GAME.players['Nona'],
-                                             player2=GAME.players['Oji'])),
-        )
-
-    def save(self, filename: str) -> None:
-        """
-        Uses pickle to save the current state of the game.
-        """
-        with open(filename, 'wb') as f:
-            pickle.dump(self, f)
-
-    @classmethod
-    def load(cls, filename: str) -> "Game | None":
-        """
-        Reload the game from a saved file.
-        """
-        try:
-            with open(filename, 'rb') as f:
-                return pickle.load(f)
-        except FileNotFoundError:
-            return None
 
 
 @bot.event
 async def on_ready():
-    global GAME
-
     config: Config = bot.config
     logger = bot.logger
 
@@ -231,23 +86,23 @@ async def on_ready():
             role, overwrite=PermissionOverwrite(read_message_history=True, read_messages=True)
         )
 
-    GAME = Game.load('game.save')
+    game = Game.load('game.save')
-    if not GAME:
+    if not game:
-        GAME = Game()
+        game = Game()
         for player in config.PLAYERS:
-            GAME.register_player(player, config.vote_channels[player.lower()])
+            game.register_player(player, config.vote_channels[player.lower()])
-        GAME.save('game.save')
+        game.save('game.save')
 
     # Update private channel id if necessary
-    for player in list(GAME.players.values()):
+    for player in list(game.players.values()):
         if player.private_channel_id != config.vote_channels[player.name.lower()]:
-            GAME.register_player(player.name, config.vote_channels[player.name.lower()])
+            game.register_player(player.name, config.vote_channels[player.name.lower()])
-            GAME.save('game.save')
+            game.save('game.save')
 
     # Setup first round if not exists
-    if not GAME.rounds:
+    if not game.rounds:
-        GAME.rounds.append(GAME.default_first_round())
+        game.rounds.append(game.default_first_round())
-        GAME.save('game.save')
+        game.save('game.save')
 
     if not config.telepathy_channel:
         channel: TextChannel = await secret_category.create_text_channel("bigbrain")

orochi/models.py

@@ -0,0 +1,150 @@
+import pickle
+from dataclasses import dataclass, field
+from datetime import datetime
+from enum import Enum
+from typing import ClassVar
+
+
+class Room(Enum):
+    A = 'A'
+    B = 'B'
+    C = 'C'
+
+
+class Vote(Enum):
+    ALLY = 'A'
+    BETRAY = 'B'
+
+
+@dataclass(frozen=True)
+class Player:
+    name: str
+    private_channel_id: int = field(hash=False)
+
+    @property
+    def round_votes(self):
+        for r in Game.INSTANCE.rounds:
+            for room in r.rooms:
+                for vote in room.votes:
+                    if self in vote.players:
+                        yield vote
+
+    @property
+    def score(self):
+        s = 3
+
+        for vote in self.round_votes:
+            room = vote.room
+            other_vote = room.vote1 if room.vote1 is not vote else room.vote2
+            match vote.vote, other_vote.vote:
+                case Vote.ALLY, Vote.ALLY:
+                    s += 2
+                case Vote.ALLY, Vote.BETRAY:
+                    s -= 2
+                case Vote.BETRAY, Vote.ALLY:
+                    s += 3
+                case Vote.BETRAY, Vote.BETRAY:
+                    pass
+
+        return s
+
+
+@dataclass
+class RoundVote:
+    player1: Player
+    player2: Player | None = None
+    vote: Vote | None = None
+    timestamp: datetime | None = None
+
+    @property
+    def players(self):
+        return self.player1, self.player2
+
+    @property
+    def room(self):
+        for r in Game.rounds:
+            for room in r.rooms:
+                if self in room.votes:
+                    return room
+
+
+@dataclass
+class RoundRoom:
+    room: Room
+    vote1: RoundVote
+    vote2: RoundVote
+
+    @property
+    def votes(self):
+        return self.vote1, self.vote2
+
+    @property
+    def round(self):
+        for r in Game.INSTANCE.rounds:
+            if self in r.rooms:
+                return r
+
+
+@dataclass
+class Round:
+    round: int
+    room_a: RoundRoom
+    room_b: RoundRoom
+    room_c: RoundRoom
+
+    @property
+    def rooms(self):
+        return self.room_a, self.room_b, self.room_c
+
+
+@dataclass
+class Game:
+    INSTANCE: ClassVar["Game"] = None
+
+    rounds: list[Round] = field(default_factory=list)
+    players: dict[str, Player] = field(default_factory=dict)
+    
+    def __post_init__(self):
+        Game.INSTANCE = self
+
+    def register_player(self, name: str, vote_channel_id: int) -> Player:
+        player = Player(name, vote_channel_id)
+        self.players[name] = player
+        return player
+
+    def default_first_round(self) -> Round:
+        return Round(
+            round=1,
+            room_a=RoundRoom(room=Room.A,
+                             vote1=RoundVote(player1=self.players['Tora']),
+                             vote2=RoundVote(player1=self.players['Kamui'],
+                                             player2=self.players['Philia'])),
+            room_b=RoundRoom(room=Room.A,
+                             vote1=RoundVote(player1=self.players['Dan']),
+                             vote2=RoundVote(player1=self.players['Ennea'],
+                                             player2=self.players['Delphine'])),
+            room_c=RoundRoom(room=Room.A,
+                             vote1=RoundVote(player1=self.players['Hanabi']),
+                             vote2=RoundVote(player1=self.players['Nona'],
+                                             player2=self.players['Oji'])),
+        )
+
+    def save(self, filename: str) -> None:
+        """
+        Uses pickle to save the current state of the game.
+        """
+        with open(filename, 'wb') as f:
+            pickle.dump(self, f)
+
+    @classmethod
+    def load(cls, filename: str) -> "Game | None":
+        """
+        Reload the game from a saved file.
+        """
+        try:
+            with open(filename, 'rb') as f:
+                game = pickle.load(f)
+                Game.INSTANCE = game
+                return game
+        except FileNotFoundError:
+            return None