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 and len(self.loc_obj) > 0: if lang and 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 None def process_result_value(self, value, dialect): if value: return self._enum_type(value) return None