diff --git a/model/bot_entity.py b/model/bot_entity.py index b4a2705..07ffc9f 100644 --- a/model/bot_entity.py +++ b/model/bot_entity.py @@ -9,10 +9,11 @@ from typing import ( get_args, get_origin, TYPE_CHECKING, + dataclass_transform, ) from pydantic import BaseModel from sqlmodel import SQLModel, BigInteger, Field, select, func, column, col - +from sqlmodel.main import FieldInfo from sqlmodel.ext.asyncio.session import AsyncSession from sqlmodel.sql.expression import SelectOfScalar from sqlmodel.main import SQLModelMetaclass, RelationshipInfo @@ -26,6 +27,10 @@ if TYPE_CHECKING: from .user import UserBase +@dataclass_transform( + kw_only_default=True, + field_specifiers=(Field, FieldInfo, EntityField, FieldDescriptor), +) class BotEntityMetaclass(SQLModelMetaclass): _future_references = {} diff --git a/model/bot_enum.py b/model/bot_enum.py index 28bdc7d..6fa3eba 100644 --- a/model/bot_enum.py +++ b/model/bot_enum.py @@ -1,6 +1,7 @@ from aiogram.utils.i18n import I18n from pydantic_core.core_schema import str_schema -from sqlalchemy.types import TypeDecorator, String +from sqlalchemy.types import TypeDecorator +from sqlmodel import AutoString from typing import Any, Self, overload @@ -119,12 +120,14 @@ class EnumMember(object): def __str__(self): return self.value - def __eq__(self, other: Self | str) -> bool: + def __eq__(self, other: Self | str | Any | None) -> bool: if other is None: return False if isinstance(other, str): return self.value == other - return self.value == other.value + if isinstance(other, EnumMember): + return self.value == other.value and self._parent is other._parent + return other.__eq__(self.value) def __hash__(self): return hash(self.value) @@ -151,24 +154,25 @@ class BotEnum(EnumMember, metaclass=BotEnumMetaclass): class EnumType(TypeDecorator): - impl = String(64) + impl = AutoString cache_ok = True - # class comparator_factory(TypeDecorator.Comparator): - # def __eq__(self, other): - # expr = type_coerce(self.expr, String) - # return expr != other.value - def __init__(self, enum_type: BotEnum): self._enum_type = enum_type super().__init__() - def process_bind_param(self, value, dialect): + def _process_param(self, value): if value and isinstance(value, EnumMember): return value.value return str(value) + def process_bind_param(self, value, dialect): + return self._process_param(value) + def process_result_value(self, value, dialect): if value: return self._enum_type(value) return None + + def process_literal_param(self, value, dialect): + return self._process_param(value) diff --git a/router.py b/router.py index 5aeacb7..4dba946 100644 --- a/router.py +++ b/router.py @@ -1,7 +1,6 @@ -from aiogram.types import CallbackQuery, Message from functools import wraps from types import UnionType -from typing import Callable, Union, get_args, get_origin, Any +from typing import Callable, Union, get_args, get_origin from .model.descriptors import ( BotCommand,