import os from typing import List import aiodynamo.client import aiohttp from token_bot.persistant_database import database as pdb, Alert from token_bot.persistant_database.user_schema import User class UsersController: def __init__(self, session: aiohttp.ClientSession): self._pdb: pdb.Database = pdb.Database(session) self.table: aiodynamo.client.Table = self._pdb.client.table( os.getenv("USERS_TABLE") ) @staticmethod def _user_to_obj(user: int | User) -> User: if isinstance(user, int): return User(user) return user @staticmethod def _alert_to_obj(alert: str | Alert) -> Alert: if isinstance(alert, str): return Alert.from_str(alert) return alert async def add(self, user: int | User): user = self._user_to_obj(user) await user.put(self.table) async def delete(self, user: int | User): user = self._user_to_obj(user) await user.delete(self.table) async def exists(self, user: int | User) -> bool: user = self._user_to_obj(user) return await user.get(self.table) async def get(self, user: int | User) -> User: user = self._user_to_obj(user) await user.get(self.table) return user async def list_alerts(self, user: int | User) -> List[Alert]: user = self._user_to_obj(user) await user.get(self.table) return user.subscribed_alerts async def is_subscribed(self, user: int | User, alert: str | Alert) -> bool: user = self._user_to_obj(user) alert = self._alert_to_obj(alert) await user.get(self.table) return alert in user.subscribed_alerts async def add_alert(self, user: int | User, alert: str | Alert) -> None: user = self._user_to_obj(user) alert = self._alert_to_obj(alert) await user.get(self.table) user.subscribed_alerts.append(alert) await user.put(self.table) async def remove_alert(self, user: int | User, alert: str | Alert) -> None: user = self._user_to_obj(user) alert = self._alert_to_obj(alert) await user.get(self.table) user.subscribed_alerts.remove(alert) await user.put(self.table)