43 lines
1.2 KiB
Python
43 lines
1.2 KiB
Python
from typing import Union
|
|
|
|
from libbot.pyrogram.classes import PyroClient
|
|
from pyrogram.types import User
|
|
|
|
from classes.pyrouser import PyroUser
|
|
from modules.database import col_users
|
|
|
|
|
|
class PyroClient(PyroClient):
|
|
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
|
|
"""
|
|
if (
|
|
await col_users.find_one(
|
|
{"id": user.id if isinstance(user, User) else user}
|
|
)
|
|
is None
|
|
):
|
|
await col_users.insert_one(
|
|
{
|
|
"id": user.id if isinstance(user, User) else user,
|
|
"locale": user.language_code if isinstance(user, User) else None,
|
|
}
|
|
)
|
|
|
|
db_record = await col_users.find_one(
|
|
{"id": user.id if isinstance(user, User) else user}
|
|
)
|
|
|
|
if db_record is None:
|
|
raise TypeError(
|
|
f"User with ID {user.id if isinstance(user, User) else user} was not found in the database"
|
|
)
|
|
|
|
return PyroUser(**db_record)
|