feat: add type annotations to modelclass (#425)

fixes: #424
polygon-io · Aug 4, 2023 · 1f9dac6 · 1f9dac6
1 parent 73b21f8
commit 1f9dac6
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/polygon/modelclass.py b/polygon/modelclass.py
@@ -1,8 +1,12 @@
 import inspect
+import typing
 from dataclasses import dataclass
 
 
-def modelclass(cls):
+_T = typing.TypeVar("_T")
+
+
+def modelclass(cls: typing.Type[_T]) -> typing.Type[_T]:
     cls = dataclass(cls)
     attributes = [
         a
@@ -18,6 +22,6 @@ def init(self, *args, **kwargs):
             if k in attributes:
                 self.__dict__[k] = v
 
-    cls.__init__ = init
+    cls.__init__ = init  # type: ignore[assignment]
 
     return cls