from typing import List from aiodynamo.client import Table from aiodynamo.errors import ItemNotFound from aiodynamo.expressions import F from token_bot.persistant_database import AlertType from token_bot.token_database.flavor import Flavor from token_bot.token_database.region import Region import token_bot.persistant_database as pdb class Alert: def __init__( self, alert: pdb.AlertType, flavor: Flavor, region: Region, price: int = 0 ) -> None: # AlertType is the Primary Key self._alert_type: pdb.AlertType = alert # Flavor (Retail, Classic) is the Sort Key self._flavor: Flavor = flavor self._region: Region = region self._price: int = price self._loaded: bool = False self._users: List[pdb.User] = [] @classmethod def from_item(cls, primary_key: int, sort_key: str, users: List[int]) -> "Alert": alert_type = pdb.AlertType(primary_key) flavor_repr, region_repr, price_repr = sort_key.split("-") flavor = Flavor(int(flavor_repr)) region = Region(region_repr) price = int(price_repr) return cls(alert_type, flavor, region, price) @classmethod def from_str(cls, string_trinity: str) -> "Alert": alert_repr, flavor_repr, region_repr, price_repr = string_trinity.split("-") if len(string_trinity.split("-")) != 4: raise ValueError alert = pdb.AlertType(int(alert_repr)) flavor = Flavor(int(flavor_repr)) region = Region(region_repr) price = int(price_repr) return cls(alert, flavor, region, price) @property def primary_key(self) -> int: return self.alert_type.value @property def primary_key_name(self) -> str: return "alert" @property def sort_key(self) -> str: return f"{self.flavor.value}-{self.region.value}-{self.price}" @property def sort_key_name(self) -> str: return "flavor-region-price" @property def key(self) -> dict[str, str | int]: return { self.primary_key_name: self.primary_key, self.sort_key_name: self.sort_key, } @property def alert_type(self) -> pdb.AlertType: return self._alert_type @property def flavor(self) -> Flavor: return self._flavor @property def region(self) -> Region: return self._region @property def price(self) -> int: return self._price @property def users(self) -> List[pdb.User]: return self._users def __str__(self): return f"{self.alert_type.value}-{self.flavor.value}-{self.region.value}-{self.price}" def __eq__(self, other): return ( self.alert_type == other.alert_type and self.flavor == other.flavor and self.price == other.price ) def to_human_string(self): if self.alert_type == AlertType.SPECIFIC_PRICE: raise NotImplementedError else: alert_type_str = " ".join(self.alert_type.name.split("_")) return f"{alert_type_str.title()}" async def _lazy_load(self, table: Table, consistent: bool = False) -> None: if consistent or not self._loaded: await self.get(table, consistent=consistent) async def _append_user(self, table: Table, user: pdb.User) -> None: self.users.append(user) await self.put(table) async def _remove_user(self, table: Table, user: pdb.User) -> None: self.users.remove(user) await self.put(table) async def put(self, table: Table) -> None: user_ids = [str(user.user_id) for user in self.users] await table.put_item( item={ self.primary_key_name: self.primary_key, self.sort_key_name: self.sort_key, "users": user_ids, } ) async def get(self, table: Table, consistent: bool = False) -> bool: try: response = await table.get_item(key=self.key, consistent_read=consistent) except ItemNotFound: return False self._users = [pdb.User(int(user_id)) for user_id in response["users"]] self._loaded = True return True async def get_users(self, table: Table, consistent: bool = False) -> List[pdb.User]: await self._lazy_load(table, consistent=consistent) return self.users async def add_user( self, table: Table, user: pdb.User, consistent: bool = False ) -> None: await self._lazy_load(table, consistent=consistent) if user not in self.users: await self._append_user(table=table, user=user) async def remove_user( self, table: Table, user: pdb.User, consistent: bool = True ) -> None: await self._lazy_load(table, consistent=consistent) if user in self.users: await self._remove_user(table=table, user=user)