quickbot/model/bot_enum.py

from aiogram.utils.i18n import I18n
from pydantic_core.core_schema import str_schema
from sqlalchemy.types import TypeDecorator, String
from typing import Any, Self, overload


class BotEnumMetaclass(type):
    def __new__(cls, name: str, bases: tuple[type], namespace: dict[str, Any]):
        all_members = {}
        if (
            bases
            and bases[0].__name__ != "BotEnum"
            and "all_members" in bases[0].__dict__
        ):
            all_members = bases[0].__dict__["all_members"]

        annotations = {}

        for key, value in namespace.items():
            if key.isupper() and not key.startswith("__") and not key.endswith("__"):
                if not isinstance(value, EnumMember):
                    value = EnumMember(value, None)

                if key in all_members.keys() and all_members[key].value != value.value:
                    raise ValueError(
                        f"Enum member {key} already exists with different value. Use same value to extend it."
                    )

                if (
                    value.value in [member.value for member in all_members.values()]
                    and key not in all_members.keys()
                ):
                    raise ValueError(f"Duplicate enum value {value[0]}")

                member = EnumMember(
                    value=value.value,
                    loc_obj=value.loc_obj,
                    parent=None,
                    name=key,
                    casting=False,
                )

                namespace[key] = member
                all_members[key] = member
                annotations[key] = type(member)

        namespace["__annotations__"] = annotations
        namespace["all_members"] = all_members

        type_ = super().__new__(cls, name, bases, namespace)

        for key, value in all_members.items():
            if not value._parent:
                value._parent = type_

        return type_


class EnumMember(object):
    @overload
    def __init__(self, value: str) -> "EnumMember": ...

    @overload
    def __init__(self, value: "EnumMember") -> "EnumMember": ...

    @overload
    def __init__(self, value: str, loc_obj: dict[str, str]) -> "EnumMember": ...

    def __init__(
        self,
        value: str = None,
        loc_obj: dict[str, str] = None,
        parent: type = None,
        name: str = None,
        casting: bool = True,
    ) -> "EnumMember":
        if not casting:
            self._parent = parent
            self._name = name
            self.value = value
            self.loc_obj = loc_obj

    @overload
    def __new__(cls: Self, *args, **kwargs) -> "EnumMember": ...

    def __new__(cls, *args, casting: bool = True, **kwargs) -> "EnumMember":
        if (cls.__name__ == "EnumMember") or not casting:
            obj = super().__new__(cls)
            kwargs["casting"] = False
            obj.__init__(*args, **kwargs)
            return obj
        if args.__len__() == 0:
            return list(cls.all_members.values())[0]
        if args.__len__() == 1 and isinstance(args[0], str):
            return {member.value: member for key, member in cls.all_members.items()}[
                args[0]
            ]
        elif args.__len__() == 1:
            return {member.value: member for key, member in cls.all_members.items()}[
                args[0].value
            ]
        else:
            return args[0]

    def __get_pydantic_core_schema__(cls, *args, **kwargs):
        return str_schema()

    def __get__(self, instance, owner) -> Self:
        return {
            member.value: member for key, member in self._parent.all_members.items()
        }[self.value]

    def __set__(self, instance, value):
        instance.__dict__[self] = value

    def __repr__(self):
        return f"<{self._parent.__name__ if self._parent else 'EnumMember'}.{self._name}: '{self.value}'>"

    def __str__(self):
        return self.value

    def __eq__(self, other: Self | str) -> bool:
        if other is None:
            return False
        if isinstance(other, str):
            return self.value == other
        return self.value == other.value

    def __hash__(self):
        return hash(self.value)

    def localized(self, lang: str = None) -> str:
        if self.loc_obj:
            if not lang:
                i18n = I18n.get_current()
                if i18n:
                    lang = i18n.current_locale
                else:
                    lang = list(self.loc_obj.keys())[0]

            if lang in self.loc_obj.keys():
                return self.loc_obj[lang]
            else:
                return self.loc_obj[list(self.loc_obj.keys())[0]]

        return self.value


class BotEnum(EnumMember, metaclass=BotEnumMetaclass):
    all_members: dict[str, EnumMember]


class EnumType(TypeDecorator):
    impl = String(64)
    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):
        if value and isinstance(value, EnumMember):
            return value.value
        return str(value)

    def process_result_value(self, value, dialect):
        if value:
            return self._enum_type(value)
        return None