Use GameType Enum properly.

Add header to requests.
Close session when cog is unloaded.

hockey/api.py

@@ -10,7 +10,7 @@
 from redbot.core.i18n import Translator
 
 from .constants import TEAMS
-from .game import Game, GameState
+from .game import Game, GameState, GameType
 from .goal import Goal
 
 TEAM_IDS = {v["id"]: k for k, v in TEAMS.items()}
@@ -89,30 +89,6 @@ class GameData(TypedDict):
     link: Optional[str]
 
 
-class GameType(Enum):
-    unknown = "Unknown"
-    pre_season = "PR"
-    regular_season = "R"
-    playoffs = "P"
-    allstars = "A"
-    allstars_women = "WA"
-    olympics = "O"
-    world_cup_exhibition = "WCOH_EXH"
-    world_cup_prelim = "WCOH_PRELIM"
-    world_cup_final = "WCOH_FINAL"
-
-    def __str__(self):
-        return str(self.value)
-
-    @classmethod
-    def from_int(cls, value: int) -> GameType:
-        return {
-            1: GameType.pre_season,
-            2: GameType.regular_season,
-            3: GameType.playoffs,
-        }.get(value, GameType.unknown)
-
-
 class GameEventTypeCode(Enum):
     UNKNOWN = 0
     FACEOFF = 502
@@ -353,7 +329,9 @@ def from_nhle(cls, data: dict) -> Schedule:
 
 class HockeyAPI:
     def __init__(self):
-        self.session = aiohttp.ClientSession()
+        self.session = aiohttp.ClientSession(
+            headers={"User-Agent": "Red-DiscordBot Trusty-cogs Hockey"}
+        )
         self.base_url = None
 
     async def close(self):

hockey/game.py

@@ -96,6 +96,30 @@ def from_nhle(cls, game_state: str, period: int) -> GameState:
         }.get(game_state, GameState.unknown)
 
 
+class GameType(Enum):
+    unknown = "Unknown"
+    pre_season = "PR"
+    regular_season = "R"
+    playoffs = "P"
+    allstars = "A"
+    allstars_women = "WA"
+    olympics = "O"
+    world_cup_exhibition = "WCOH_EXH"
+    world_cup_prelim = "WCOH_PRELIM"
+    world_cup_final = "WCOH_FINAL"
+
+    def __str__(self):
+        return str(self.value)
+
+    @classmethod
+    def from_int(cls, value: int) -> GameType:
+        return {
+            1: GameType.pre_season,
+            2: GameType.regular_season,
+            3: GameType.playoffs,
+        }.get(value, GameType.unknown)
+
+
 @dataclass
 class GameStatus:
     abstractGameState: str
@@ -294,7 +318,7 @@ def __init__(self, **kwargs):
         self.third_star = kwargs.get("third_star")
         self.away_roster = kwargs.get("away_roster")
         self.home_roster = kwargs.get("home_roster")
-        self.game_type: str = kwargs.get("game_type", "")
+        self.game_type: GameType = kwargs.get("game_type", GameType.unknown)
         self.link = kwargs.get("link")
         self.season = kwargs.get("season")
         self._recap_url: Optional[str] = kwargs.get("recap_url", None)
@@ -326,7 +350,11 @@ def timestamp(self) -> int:
         return int(self.game_start.timestamp())
 
     def game_type_str(self):
-        game_types = {"PR": _("Pre Season"), "R": _("Regular Season"), "P": _("Post Season")}
+        game_types = {
+            GameType.pre_season: _("Pre Season"),
+            GameType.regular_season: _("Regular Season"),
+            GameType.playoffs: _("Post Season"),
+        }
         return game_types.get(self.game_type, _("Unknown"))
 
     def to_json(self) -> dict:
@@ -353,7 +381,7 @@ def to_json(self) -> dict:
             "first_star": self.first_star,
             "second_star": self.second_star,
             "third_star": self.third_star,
-            "game_type": self.game_type,
+            "game_type": self.game_type.value,
             "link": self.link,
         }
 
@@ -603,7 +631,7 @@ async def get_stats_msg(self) -> Tuple[str, str, Optional[str]]:
         home_str = _("GP:**0** W:**0** L:**0\n**OT:**0** PTS:**0** S:**0**\n")
         away_str = _("GP:**0** W:**0** L:**0\n**OT:**0** PTS:**0** S:**0**\n")
         desc = None
-        if self.game_type != "P":
+        if self.game_type is GameType.playoffs:
             msg = _(
                 "GP:**{gp}** W:**{wins}** L:**{losses}\n**OT:**{ot}** PTS:**{pts}** S:**{streak}**\n"
             )
@@ -874,7 +902,7 @@ async def actually_post_state(
                 allowed_mentions = {"allowed_mentions": discord.AllowedMentions(roles=True)}
             else:
                 allowed_mentions = {"allowed_mentions": discord.AllowedMentions(roles=False)}
-            if self.game_type == "R" and "OT" in self.period_ord:
+            if self.game_type is GameType.regular_season and "OT" in self.period_ord:
                 if not guild_settings["ot_notifications"]:
                     allowed_mentions = {"allowed_mentions": discord.AllowedMentions(roles=False)}
             if "SO" in self.period_ord:

hockey/hockey.py

@@ -179,6 +179,7 @@ async def cog_unload(self):
         if self.loop is not None:
             self.loop.cancel()
         await self.session.close()
+        await self.api.close()
         self.pickems_loop.cancel()
         await self.after_pickems_loop()
 

hockey/hockeypickems.py

@@ -13,7 +13,7 @@
 from hockey.helper import utc_to_local
 
 from .abc import HockeyMixin
-from .game import Game
+from .game import Game, GameType
 from .pickems import Pickems
 
 _ = Translator("Hockey", __file__)
@@ -91,7 +91,7 @@ async def save_pickems_data(self) -> None:
                         log.trace("Saving pickem %r", pickem)
                         data[name] = pickem.to_json()
                     self.all_pickems[guild_id][name]._should_save = False
-                    if pickem.game_type in ["P", "PR"]:
+                    if pickem.game_type in [GameType.pre_season, GameType.playoffs]:
                         if (datetime.now(timezone.utc) - pickem.game_start) >= timedelta(days=7):
                             del data[name]
                             if guild_id not in to_del:
@@ -748,11 +748,11 @@ async def tally_guild_leaderboard(self, guild: discord.Guild) -> None:
                                 await bank.deposit_credits(member, int(base_credits))
                             except Exception:
                                 log.debug("Could not deposit pickems credits for %r", member)
-                        if pickems.game_type == "P":
+                        if pickems.game_type is GameType.playoffs:
                             leaderboard[str(user)]["playoffs"] += 1
                             leaderboard[str(user)]["playoffs_weekly"] += 1
                             leaderboard[str(user)]["playoffs_total"] += 1
-                        elif pickems.game_type == "PR":
+                        elif pickems.game_type is GameType.pre_season:
                             leaderboard[str(user)]["pre-season"] += 1
                             leaderboard[str(user)]["pre-season_weekly"] += 1
                             leaderboard[str(user)]["pre-season_total"] += 1
@@ -763,9 +763,9 @@ async def tally_guild_leaderboard(self, guild: discord.Guild) -> None:
                             # playoffs is finished
                             leaderboard[str(user)]["weekly"] += 1
                     else:
-                        if pickems.game_type == "P":
+                        if pickems.game_type is GameType.playoffs:
                             leaderboard[str(user)]["playoffs_total"] += 1
-                        elif pickems.game_type == "PR":
+                        elif pickems.game_type is GameType.pre_season:
                             leaderboard[str(user)]["pre-season_total"] += 1
                         else:
                             leaderboard[str(user)]["total"] += 1

hockey/pickems.py

@@ -9,7 +9,7 @@
 
 from .constants import TEAMS
 from .errors import NotAValidTeamError, UserHasVotedError, VotingHasEndedError
-from .game import Game
+from .game import Game, GameState, GameType
 
 _ = Translator("Hockey", __file__)
 log = getLogger("red.trusty-cogs.Hockey")
@@ -103,7 +103,7 @@ class Pickems(discord.ui.View):
     def __init__(
         self,
         game_id: int,
-        game_state: str,
+        game_state: GameState,
         messages: List[str],
         guild: int,
         game_start: datetime,
@@ -113,7 +113,7 @@ def __init__(
         name: str,
         winner: Optional[str],
         link: Optional[str],
-        game_type: str,
+        game_type: GameType,
         should_edit: bool,
     ):
         self.game_id = game_id
@@ -222,7 +222,7 @@ def add_vote(self, user_id: int, team: discord.Emoji) -> None:
     def to_json(self) -> Dict[str, Any]:
         return {
             "game_id": self.game_id,
-            "game_state": self.game_state,
+            "game_state": self.game_state.value,
             "messages": self.messages,
             "guild": self.guild,
             "game_start": self.game_start.strftime("%Y-%m-%dT%H:%M:%SZ"),
@@ -232,7 +232,7 @@ def to_json(self) -> Dict[str, Any]:
             "name": self.name,
             "winner": self.winner,
             "link": self.link,
-            "game_type": self.game_type,
+            "game_type": self.game_type.value,
             "should_edit": self.should_edit,
         }
 
@@ -243,7 +243,7 @@ def from_json(cls, data: dict) -> Pickems:
         game_start = game_start.replace(tzinfo=timezone.utc)
         return cls(
             game_id=data["game_id"],
-            game_state=data["game_state"],
+            game_state=GameState(data["game_state"]),
             messages=data.get("messages", []),
             guild=data["guild"],
             game_start=game_start,
@@ -253,7 +253,7 @@ def from_json(cls, data: dict) -> Pickems:
             name=data.get("name", ""),
             winner=data.get("winner", None),
             link=data.get("link", None),
-            game_type=data.get("game_type", "R"),
+            game_type=GameType(data.get("game_type", "R")),
             should_edit=data.get("should_edit", True),
         )