From 99e134b1607acb321f6a82bafe0aef2632153681 Mon Sep 17 00:00:00 2001 From: TrustyJAID Date: Sun, 12 Nov 2023 14:04:19 -0700 Subject: [PATCH] Use GameType Enum properly. Add header to requests. Close session when cog is unloaded. --- hockey/api.py | 30 ++++-------------------------- hockey/game.py | 38 +++++++++++++++++++++++++++++++++----- hockey/hockey.py | 1 + hockey/hockeypickems.py | 12 ++++++------ hockey/pickems.py | 14 +++++++------- 5 files changed, 51 insertions(+), 44 deletions(-) diff --git a/hockey/api.py b/hockey/api.py index b95ffbc870..0964181661 100644 --- a/hockey/api.py +++ b/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): diff --git a/hockey/game.py b/hockey/game.py index 6ddcc318be..6d463f5bf1 100644 --- a/hockey/game.py +++ b/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: diff --git a/hockey/hockey.py b/hockey/hockey.py index 96f5be5aaf..acf79c801c 100644 --- a/hockey/hockey.py +++ b/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() diff --git a/hockey/hockeypickems.py b/hockey/hockeypickems.py index e277a9bd13..01ca32464f 100644 --- a/hockey/hockeypickems.py +++ b/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 diff --git a/hockey/pickems.py b/hockey/pickems.py index 223f60d023..a2c2f85dad 100644 --- a/hockey/pickems.py +++ b/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), )