feat(protobuf): enhanced type annotations (mypy)

Author: pk5ls20

lagrange/utils/binary/protobuf/models.py

@@ -1,28 +1,25 @@
 import inspect
-import importlib
 from types import GenericAlias
-from typing import cast, Dict, List, Tuple, Type, TypeVar, Union, Generic, Any, Callable, Mapping, overload, ForwardRef
+from typing import cast, Dict, List, Tuple, Type, TypeVar, Union, Generic, Any, Callable, Mapping, overload
 from typing_extensions import Optional, Self, TypeAlias, dataclass_transform
 
 from .coder import Proto, proto_decode, proto_encode
 
-_ProtoTypes = Union[str, list, dict, bytes, int, float, bool, "ProtoStruct"]
+_ProtoBasicTypes = Union[str, list, dict, bytes, int, float, bool]
+_ProtoTypes = Union[_ProtoBasicTypes, "ProtoStruct"]
 
 T = TypeVar("T", str, list, dict, bytes, int, float, bool, "ProtoStruct")
 V = TypeVar("V")
 NT: TypeAlias = Dict[int, Union[_ProtoTypes, "NT"]]
-AMT: TypeAlias = Dict[str, Tuple[Type[_ProtoTypes], "ProtoField"]]
-DAMT: TypeAlias = Dict[str, "DelayAnnoType"]
-DelayAnnoType = Union[str, type(List[str])]
 NoneType = type(None)
 
 
 class ProtoField(Generic[T]):
     def __init__(self, tag: int, default: T):
         if tag <= 0:
             raise ValueError("Tag must be a positive integer")
-        self._tag = tag
-        self._default = default
+        self._tag: int = tag
+        self._default: T = default
 
     @property
     def tag(self) -> int:
@@ -86,16 +83,14 @@ def proto_field(
 @dataclass_transform(kw_only_default=True, field_specifiers=(proto_field,))
 class ProtoStruct:
     _anno_map: Dict[str, Tuple[Type[_ProtoTypes], ProtoField[Any]]]
-    _delay_anno_map: Dict[str, DelayAnnoType]
     _proto_debug: bool
 
     def __init__(self, *args, **kwargs):
         undefined_params: List[str] = []
-        args = list(args)
-        self._resolve_annotations(self)
+        arg_list = list(args)
         for name, (typ, field) in self._anno_map.items():
             if args:
-                self._set_attr(name, typ, args.pop(0))
+                self._set_attr(name, typ, arg_list.pop(0))
             elif name in kwargs:
                 self._set_attr(name, typ, kwargs.pop(name))
             else:
@@ -104,13 +99,11 @@ def __init__(self, *args, **kwargs):
                 else:
                     undefined_params.append(name)
         if undefined_params:
-            raise AttributeError(
-                "Undefined parameters in '{}': {}".format(self, undefined_params)
-            )
+            raise AttributeError(f"Undefined parameters in '{self}': {undefined_params}")
 
     def __init_subclass__(cls, **kwargs):
+        cls._anno_map = cls._get_annotations()
         cls._proto_debug = kwargs.pop("debug") if "debug" in kwargs else False
-        cls._anno_map, cls._delay_anno_map = cls._get_annotations()
         super().__init_subclass__(**kwargs)
 
     def __repr__(self) -> str:
@@ -125,17 +118,14 @@ def _set_attr(self, name: str, data_typ: Type[V], value: V) -> None:
         if isinstance(data_typ, GenericAlias):  # force ignore
             pass
         elif not isinstance(value, data_typ) and value is not None:
-            raise TypeError(
-                "'{}' is not a instance of type '{}'".format(value, data_typ)
-            )
+            raise TypeError(f"{value} is not a instance of type {data_typ}")
         setattr(self, name, value)
 
     @classmethod
     def _get_annotations(
         cls,
-    ) -> Tuple[AMT, DAMT]:  # Name: (ReturnType, ProtoField)
-        annotations: AMT = {}
-        delay_annotations: DAMT = {}
+    ) -> Dict[str, Tuple[Type[_ProtoTypes], "ProtoField"]]:  # Name: (ReturnType, ProtoField)
+        annotations: Dict[str, Tuple[Type[_ProtoTypes], "ProtoField"]] = {}
         for obj in reversed(inspect.getmro(cls)):
             if obj in (ProtoStruct, object):  # base object, ignore
                 continue
@@ -149,34 +139,15 @@ def _get_annotations(
                 if not isinstance(field, ProtoField):
                     raise TypeError("attribute '{name}' is not a ProtoField object")
 
-                _typ = typ
-                annotations[name] = (_typ, field)
-                if isinstance(typ, str):
-                    delay_annotations[name] = typ
                 if hasattr(typ, "__origin__"):
-                    typ = cast(GenericAlias, typ)
-                    _inner = typ.__args__[0]
-                    _typ = typ.__origin__[typ.__args__[0]]
-                    annotations[name] = (_typ, field)
-
-                    if isinstance(_inner, type):
-                        continue
-                    if isinstance(_inner, GenericAlias) and isinstance(_inner.__args__[0], type):
-                        continue
-                    if isinstance(_inner, str):
-                        delay_annotations[name] = _typ.__origin__[_inner]
-                    if isinstance(_inner, ForwardRef):
-                        delay_annotations[name] = _inner.__forward_arg__
-                    if isinstance(_inner, GenericAlias):
-                        delay_annotations[name] = _typ
-
-        return annotations, delay_annotations
+                    typ = typ.__origin__[typ.__args__[0]]
+                annotations[name] = (typ, field)
+
+        return annotations
 
     @classmethod
     def _get_field_mapping(cls) -> Dict[int, Tuple[str, Type[_ProtoTypes]]]:  # Tag, (Name, Type)
         field_mapping: Dict[int, Tuple[str, Type[_ProtoTypes]]] = {}
-        if cls._delay_anno_map:
-            cls._resolve_annotations(cls)
         for name, (typ, field) in cls._anno_map.items():
             field_mapping[field.tag] = (name, typ)
         return field_mapping
@@ -187,17 +158,7 @@ def _get_stored_mapping(self) -> Dict[str, NT]:
             stored_mapping[name] = getattr(self, name)
         return stored_mapping
 
-    @staticmethod
-    def _resolve_annotations(arg: Union[Type["ProtoStruct"], "ProtoStruct"]) -> None:
-        for k, v in arg._delay_anno_map.copy().items():
-            module = importlib.import_module(arg.__module__)
-            if hasattr(v, "__origin__"):  # resolve GenericAlias, such as list[str]
-                arg._anno_map[k] = (v.__origin__[module.__getattribute__(v.__args__[0])], arg._anno_map[k][1])
-            else:
-                arg._anno_map[k] = (module.__getattribute__(v), arg._anno_map[k][1])
-            arg._delay_anno_map.pop(k)
-
-    def _encode(self, v: _ProtoTypes) -> NT:
+    def _encode(self, v: _ProtoTypes) -> _ProtoBasicTypes:
         if isinstance(v, ProtoStruct):
             v = v.encode()
         return v