From 3f957db2ebf9b04c68fc769138324edd0862a7d6 Mon Sep 17 00:00:00 2001 From: Sergey Elpashev Date: Sat, 10 Feb 2024 12:59:31 +0300 Subject: [PATCH] Use pydantic instead of dataclasses --- neuroapi/_methods/api_method.py | 11 +++---- neuroapi/types/_admin.py | 21 ++---------- neuroapi/types/_api_model.py | 12 +++++++ neuroapi/types/_bot.py | 15 ++++++--- neuroapi/types/_bot_settings.py | 41 ++++++------------------ neuroapi/types/_helpers.py | 47 --------------------------- neuroapi/types/_image.py | 24 ++------------ neuroapi/types/_post.py | 57 +++++++++------------------------ neuroapi/types/_user.py | 29 +++++++---------- 9 files changed, 70 insertions(+), 187 deletions(-) create mode 100644 neuroapi/types/_api_model.py delete mode 100644 neuroapi/types/_helpers.py diff --git a/neuroapi/_methods/api_method.py b/neuroapi/_methods/api_method.py index af239e4..e3ae9e3 100644 --- a/neuroapi/_methods/api_method.py +++ b/neuroapi/_methods/api_method.py @@ -1,8 +1,7 @@ -from ..config import Config +from pydantic import BaseModel, Field + +from ..config import GlobalConfig as Config -class ApiMethod: - api_url: str - - def __init__(self) -> None: - self.api_url = Config().api_url +class ApiMethod(BaseModel): + api_url: str = Field(Config().api_url) diff --git a/neuroapi/types/_admin.py b/neuroapi/types/_admin.py index 19692c7..955576a 100644 --- a/neuroapi/types/_admin.py +++ b/neuroapi/types/_admin.py @@ -1,23 +1,6 @@ -from dataclasses import dataclass -from typing import Any - -from ._helpers import * +from ._api_model import ApiModel -@dataclass -class Admin: +class Admin(ApiModel): user_id: int user_name: str - - @staticmethod - def from_dict(obj: Any) -> 'Admin': - assert isinstance(obj, dict) - user_id = int(from_str(obj.get("user_id"))) - user_name = from_str(obj.get("user_name")) - return Admin(user_id, user_name) - - def to_dict(self) -> dict: - result: dict = {} - result["user_id"] = from_str(str(self.user_id)) - result["user_name"] = from_str(self.user_name) - return result diff --git a/neuroapi/types/_api_model.py b/neuroapi/types/_api_model.py new file mode 100644 index 0000000..083aecb --- /dev/null +++ b/neuroapi/types/_api_model.py @@ -0,0 +1,12 @@ +from typing import Any, Dict + +from pydantic import BaseModel + + +class ApiModel(BaseModel): + @classmethod + def from_dict(cls: 'ApiModel', obj: Dict[str, Any]) -> 'ApiModel': + return cls(**obj) + + def to_dict(self, **kwargs: Any) -> Dict[str, Any]: + return self.model_dump(**kwargs) \ No newline at end of file diff --git a/neuroapi/types/_bot.py b/neuroapi/types/_bot.py index 58306bd..ffce593 100644 --- a/neuroapi/types/_bot.py +++ b/neuroapi/types/_bot.py @@ -1,8 +1,12 @@ from aiogram import Bot, Dispatcher +from pydantic import BaseModel from handlers.handler import Handler +class Token(BaseModel): + token: str + class NeuroApiBot: bot: Bot dp: Dispatcher @@ -10,14 +14,15 @@ class NeuroApiBot: _instances = {} def __init__(self, token: str) -> None: - self.bot = Bot(token) + token_data = Token(token=token) + self.bot = Bot(token_data.token) self.dp = Dispatcher() def __new__(cls, token: str) -> 'NeuroApiBot': - assert isinstance(token, str) - if token not in cls._instances: - cls._instances[token] = super(NeuroApiBot, cls).__new__(cls) - return cls._instances[token] + token_data = Token(token=token) + if token_data.token not in cls._instances: + cls._instances[token_data.token] = super(NeuroApiBot, cls).__new__(cls) + return cls._instances[token_data.token] def include_router(self, *routerClasses: Handler) -> None: for routerClass in routerClasses: diff --git a/neuroapi/types/_bot_settings.py b/neuroapi/types/_bot_settings.py index fad43b9..7faa1b5 100644 --- a/neuroapi/types/_bot_settings.py +++ b/neuroapi/types/_bot_settings.py @@ -1,38 +1,17 @@ -from dataclasses import dataclass -from typing import Optional +from typing import List from uuid import UUID -from ._helpers import * +from pydantic import Field + +from ._api_model import ApiModel from ._singleton import Singleton -@dataclass -class BotSettings(Singleton): +class BotSettings(ApiModel, Singleton): uuid: UUID - message_times: List[str] + message_times: List[str] = Field([], alias='messageTimes') channel: str - is_active: bool - - @staticmethod - def from_dict(obj: Any) -> 'BotSettings': - assert isinstance(obj, dict) - uuid = UUID(obj.get("uuid")) - message_times = from_list(from_str, obj.get("messageTimes")) - channel = from_str(obj.get("channel")) - is_active = from_bool(obj.get("isActive")) - return BotSettings(uuid, message_times, channel, is_active) - - def to_dict(self) -> dict: - result: dict = {} - result["uuid"] = str(self.uuid) - result["messageTimes"] = from_list(from_str, self.message_times) - result["channel"] = from_str(self.channel) - result["isActive"] = from_bool(self.is_active) - return result - - @staticmethod - def get_active() -> Optional['BotSettings']: - try: - return BotSettings._instances[BotSettings] - except: - return None \ No newline at end of file + is_active: bool = Field(False, alias='isActive') + + def get_text(self): + return f"Канал: {self.channel}\nВремя: {', '.join(self.message_times)}" \ No newline at end of file diff --git a/neuroapi/types/_helpers.py b/neuroapi/types/_helpers.py deleted file mode 100644 index 9324d7d..0000000 --- a/neuroapi/types/_helpers.py +++ /dev/null @@ -1,47 +0,0 @@ -from datetime import datetime -from typing import Any, Callable, List, TypeVar, Type, cast -import dateutil.parser -T = TypeVar("T") - - -def from_bool(x: Any) -> bool: - assert isinstance(x, bool) - return x - - -def from_str(x: Any) -> str: - assert isinstance(x, str) - return x - - -def from_union(fs, x): - for f in fs: - try: - return f(x) - except: - pass - assert False - - -def from_none(x: Any) -> Any: - assert x is None - return x - - -def from_list(f: Callable[[Any], T], x: Any) -> List[T]: - assert isinstance(x, list) - return [f(y) for y in x] - - -def from_datetime(x: Any) -> datetime: - return dateutil.parser.parse(x) - - -def to_class(c: Type[T], x: Any) -> dict: - assert isinstance(x, c) - return cast(Any, x).to_dict() - - -def from_int(x: Any) -> int: - assert isinstance(x, int) and not isinstance(x, bool) - return x diff --git a/neuroapi/types/_image.py b/neuroapi/types/_image.py index bce6fb7..d6f9812 100644 --- a/neuroapi/types/_image.py +++ b/neuroapi/types/_image.py @@ -1,28 +1,10 @@ -from dataclasses import dataclass from uuid import UUID -from ._helpers import * + +from ._api_model import ApiModel -@dataclass -class Image: +class Image(ApiModel): message_id: int file_id: str has_spoiler: bool post_uuid: UUID - - @staticmethod - def from_dict(obj: Any) -> 'Image': - assert isinstance(obj, dict) - message_id = from_int(obj.get("message_id")) - file_id = from_str(obj.get("file_id")) - has_spoiler = from_bool(obj.get("has_spoiler")) - post_uuid = UUID(obj.get("post_uuid")) - return Image(message_id, file_id, has_spoiler, post_uuid) - - def to_dict(self) -> dict: - result: dict = {} - result["message_id"] = from_int(self.message_id) - result["file_id"] = from_str(self.file_id) - result["has_spoiler"] = from_bool(self.has_spoiler) - result["post_uuid"] = str(self.post_uuid) - return result diff --git a/neuroapi/types/_post.py b/neuroapi/types/_post.py index 60e821d..5b85755 100644 --- a/neuroapi/types/_post.py +++ b/neuroapi/types/_post.py @@ -1,59 +1,34 @@ import json -from dataclasses import dataclass from datetime import datetime -from typing import Any, List, Optional +from typing import Any, Dict, List, Optional from uuid import UUID from aiogram.types import MessageEntity +from pydantic import Field -from ._helpers import * +from ._api_model import ApiModel from ._image import Image -def to_message_dict_class(x: Any) -> dict: - assert isinstance(x, MessageEntity) - return cast(MessageEntity, x).model_dump() - - -@dataclass -class Post: +class Post(ApiModel): uuid: UUID posted: bool text: str media_group_id: int | str timestamp: datetime from_user_id: int - images: Optional[List[Image]] = None - message_entities: Optional[List[MessageEntity]] = None + images: Optional[List[Image]] = Field(None) + message_entities: Optional[List[MessageEntity]] = Field(None) + - @staticmethod - def from_dict(obj: Any) -> 'Post': - assert isinstance(obj, dict) - uuid = UUID(obj.get("uuid")) - posted = from_bool(obj.get("posted")) - text = from_str(obj.get("text")) - media_group_id = from_str(obj.get("media_group_id")) if obj.get( - "media_group_id") is not None else 'None' - timestamp = from_datetime(obj.get("timestamp")) - from_user_id = int(from_str(obj.get("from_user_id"))) - images = from_union([lambda x: from_list( - Image.from_dict, x), from_none], obj.get("images")) - mes_ent = json.loads(obj.get('message_entities','[]')) - message_entities = from_list(MessageEntity.model_validate, mes_ent) - return Post(uuid, posted, text, media_group_id, timestamp, from_user_id, images, message_entities) + @classmethod + def from_dict(cls: 'Post', obj: Dict[str, Any]) -> 'Post': + mes_ent = json.loads(obj.get('message_entities', '[]')) + obj['message_entities'] = mes_ent + return cls(**obj) def to_dict(self) -> dict: - result: dict = {} - result["uuid"] = str(self.uuid) - result["posted"] = from_bool(self.posted) - result["text"] = from_str(self.text) - result["media_group_id"] = from_str(str(self.media_group_id)) - result["timestamp"] = self.timestamp.isoformat() - result["from_user_id"] = from_str(str(self.from_user_id)) - if self.images is not None: - result["images"] = from_union([lambda x: from_list( - lambda x: to_class(Image, x), x), from_none], self.images) - if self.message_entities is not None: - result['message_entities'] = from_union([lambda x: from_list( - lambda x: to_message_dict_class(x), x), from_none], self.message_entities) - return result + obj = super().to_dict() + obj['message_entities'] = json.dumps(obj['message_entities']) + obj['media_group_id'] = str(obj['media_group_id']) + return obj \ No newline at end of file diff --git a/neuroapi/types/_user.py b/neuroapi/types/_user.py index 71300a0..908f1fe 100644 --- a/neuroapi/types/_user.py +++ b/neuroapi/types/_user.py @@ -1,22 +1,17 @@ -from dataclasses import dataclass -from typing import Any -from ._helpers import * +import json +from typing import Optional + +from pydantic import Field + +from ._api_model import ApiModel -@dataclass -class User: +class User(ApiModel): id: int - username: str - - @staticmethod - def from_dict(obj: Any) -> 'User': - assert isinstance(obj, dict) - id = int(from_str(obj.get("id"))) - username = from_str(obj.get("username")) - return User(id, username) + username: str = Field(..., alias='user_name') + banned: Optional[bool] = Field(None) def to_dict(self) -> dict: - result: dict = {} - result["id"] = from_str(str(self.id)) - result["username"] = from_str(self.username) - return result + obj = super().to_dict(exclude_unset=True) + obj['id'] = str(obj['id']) + return obj \ No newline at end of file