wow-token-app-bot/token_bot/controller/users.py

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)