from typing import Union

from libbot.pyrogram.classes import PyroClient as LibPyroClient
from pyrogram.types import User

from classes.pyrouser import PyroUser


class PyroClient(LibPyroClient):
    async def find_user(self, user: Union[int, User]) -> PyroUser:
        """Find User by it's ID or User object.

        ### Args:
            * user (`Union[int, User]`): ID or User object to extract ID from

        ### Returns:
            * `PyroUser`: PyroUser object
        """

        return (
            await PyroUser.find(user)
            if isinstance(user, int)
            else await PyroUser.find(user.id, locale=user.language_code)
        )