from typing import List, Dict from aiodynamo.client import Table from aiodynamo.errors import ItemNotFound import token_bot.persistant_database as pdb from token_bot.token_database.region import Region class User: def __init__( self, user_id: int, region: Region = None, subscribed_alerts: List["pdb.Alert"] = None, ) -> None: self.user_id: int = user_id self._loaded: bool = False self.region: Region = region self.subscribed_alerts: List[pdb.Alert] = subscribed_alerts def __eq__(self, other): return self.user_id == other.user_id def __hash__(self): return hash(self.user_id) @classmethod def from_item( cls, primary_key: int, region: Region, subscribed_alerts: List[str] ) -> "User": alerts = [pdb.Alert.from_str(alert_str) for alert_str in subscribed_alerts] return cls(primary_key, region, alerts) @property def primary_key(self) -> str: return str(self.user_id) @property def primary_key_name(self) -> str: return "user_id" @property def key(self) -> Dict[str, str]: return {self.primary_key_name: self.primary_key} def _subscribed_alerts_as_trinity_list(self) -> List[str]: return ( [str(alert) for alert in self.subscribed_alerts] if self.subscribed_alerts else [] ) 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 put(self, table: Table) -> None: await table.put_item( item={ self.primary_key_name: self.primary_key, "region": self.region, "subscribed_alerts": self._subscribed_alerts_as_trinity_list(), } ) async def delete(self, table: Table) -> None: if not self._loaded: await self._lazy_load(table, consistent=True) await table.delete_item( key={self.primary_key_name: self.primary_key}, ) 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.subscribed_alerts = [] for string_trinity in response["subscribed_alerts"]: self.subscribed_alerts.append(pdb.Alert.from_str(string_trinity)) self.region = Region(response["region"]) return True async def add_alert( self, table: Table, alert: "pdb.Alert", consistent: bool = False ) -> None: await self._lazy_load(table, consistent=consistent) if alert not in self.subscribed_alerts: self.subscribed_alerts.append(alert) await self.put(table) async def remove_alert( self, table: Table, alert: "pdb.Alert", consistent: bool = True ) -> None: await self._lazy_load(table, consistent=consistent) if alert in self.subscribed_alerts: self.subscribed_alerts.remove(alert) await self.put(table)