init

model/bot_enum.py

@@ -0,0 +1,147 @@
+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 owner.all_members.items()}[self.value]
+        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(256)
+
+    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