some typing shit

Author: esblinov

transfunctions/decorators/superfunction.py

@@ -2,7 +2,8 @@
 from ast import AST, NodeTransformer, Return
 from functools import wraps
 from inspect import currentframe
-from typing import Any, Dict, Generic, List, Optional, Union, overload
+from typing import Any, Dict, Generic, List, Optional, Union, Type, overload, cast
+from types import FrameType, TracebackType
 
 from displayhooks import not_display
 
@@ -60,7 +61,7 @@ def __invert__(self) -> ReturnType:
     def send(self, value: Any) -> Any:
         return self.coroutine.send(value)
 
-    def throw(self, exception_type: Any, value: Any = None, traceback: Any = None) -> None:  # pragma: no cover
+    def throw(self, exception_type: Type[BaseException], value: Any = None, traceback: Optional[TracebackType] = None) -> None:  # pragma: no cover
         pass
 
     def close(self) -> None:  # pragma: no cover
@@ -111,9 +112,9 @@ def superfunction(  # type: ignore[misc]
     def decorator(function: Callable[FunctionParams, ReturnType]) -> Callable[FunctionParams, UsageTracer[FunctionParams, ReturnType]]:
         transformer = FunctionTransformer(
             function,
-            currentframe().f_back.f_lineno,  # type: ignore[union-attr]
+            cast(FrameType, cast(FrameType, currentframe()).f_back).f_lineno,
             "superfunction",
-            currentframe().f_back,
+            cast(FrameType, cast(FrameType, currentframe()).f_back),
         )
 
         if not tilde_syntax:

transfunctions/decorators/transfunction.py

@@ -1,4 +1,6 @@
 from inspect import currentframe
+from typing import cast
+from types import FrameType
 
 from transfunctions.transformer import FunctionTransformer
 from transfunctions.typing import Callable, FunctionParams, ReturnType
@@ -7,4 +9,9 @@
 def transfunction(
     function: Callable[FunctionParams, ReturnType],
 ) -> FunctionTransformer[FunctionParams, ReturnType]:
-    return FunctionTransformer(function, currentframe().f_back.f_lineno, "transfunction", currentframe().f_back)  # type: ignore[union-attr]
+    return FunctionTransformer(
+        function,
+        cast(FrameType, cast(FrameType, currentframe()).f_back).f_lineno,
+        "transfunction",
+        cast(FrameType, cast(FrameType, currentframe()).f_back),
+    )  # type: ignore[union-attr]

transfunctions/transformer.py

@@ -7,6 +7,7 @@
     Call,
     Constant,
     FunctionDef,
+    Module,
     Load,
     Name,
     NodeTransformer,
@@ -23,7 +24,7 @@
 from inspect import getfile, getsource, iscoroutinefunction, isfunction
 from sys import version_info
 from types import FunctionType, MethodType, FrameType
-from typing import Any, Dict, Generic, List, Optional, Union, cast
+from typing import Any, Dict, Generic, List, Optional, Union, Type, cast
 
 from dill.source import getsource as dill_getsource  # type: ignore[import-untyped]
 
@@ -34,7 +35,7 @@
     WrongDecoratorSyntaxError,
     WrongMarkerSyntaxError,
 )
-from transfunctions.typing import Coroutine, Callable, Generator, FunctionParams, ReturnType
+from transfunctions.typing import Coroutine, Callable, Generator, FunctionParams, ReturnType, SomeClassInstance
 from transfunctions.universal_namespace import UniversalNamespaceAroundFunction
 
 
@@ -55,13 +56,17 @@ def __init__(
         self.decorator_lineno = decorator_lineno
         self.decorator_name = decorator_name
         self.frame = frame
-        self.base_object = None
+        self.base_object: Optional[SomeClassInstance] = None  # type: ignore[valid-type]
         self.cache: Dict[str, Callable] = {}
 
     def __call__(self, *args: Any, **kwargs: Any) -> None:
         raise CallTransfunctionDirectlyError("You can't call a transfunction object directly, create a function, a generator function or a coroutine function from it.")
 
-    def __get__(self, base_object, type=None):
+    def __get__(
+        self,
+        base_object: SomeClassInstance,
+        owner: Type[SomeClassInstance],
+    ) -> 'FunctionTransformer[FunctionParams, ReturnType]':
         self.base_object = base_object
         return self
 
@@ -253,7 +258,7 @@ def visit_FunctionDef(self, node: FunctionDef) -> Optional[Union[AST, List[AST]]
 
         return result
 
-    def wrap_ast_by_closures(self, tree):
+    def wrap_ast_by_closures(self, tree: Module) -> Module:
         old_functiondef = tree.body[0]
 
         tree.body[0] = FunctionDef(
@@ -276,7 +281,7 @@ def wrap_ast_by_closures(self, tree):
         return tree
 
 
-    def rewrite_globals_and_closure(self, function):
+    def rewrite_globals_and_closure(self, function: FunctionType) -> FunctionType:
         # https://stackoverflow.com/a/13503277/14522393
         all_new_closure_names = set(self.function.__code__.co_freevars)
 
@@ -297,6 +302,6 @@ def rewrite_globals_and_closure(self, function):
             closure=filtered_closure,
         )
 
-        new_function = update_wrapper(new_function, function)
+        new_function = cast(FunctionType, update_wrapper(new_function, function))
         new_function.__kwdefaults__ = function.__kwdefaults__
         return new_function

transfunctions/typing.py

@@ -20,10 +20,11 @@
 
 ReturnType = TypeVar('ReturnType')
 FunctionParams = ParamSpec('FunctionParams')
+SomeClassInstance = TypeVar('SomeClassInstance')
 
 if sys.version_info >= (3, 9):
     IterableWithResults = Iterable[ReturnType]
 else:
     IterableWithResults = Iterable
 
-__all__ = ('ParamSpec', 'TypeAlias', 'Callable', 'Coroutine', 'Generator', 'ReturnType', 'FunctionParams', 'IterableWithResults')
+__all__ = ('ParamSpec', 'TypeAlias', 'Callable', 'Coroutine', 'Generator', 'ReturnType', 'FunctionParams', 'IterableWithResults', 'SomeClassInstance')

transfunctions/universal_namespace.py

@@ -1,4 +1,4 @@
-from typing import Dict, Any
+from typing import Dict, Optional, Any
 from types import FrameType
 import builtins
 
@@ -7,7 +7,7 @@ class Nothing:
     pass
 
 class UniversalNamespaceAroundFunction(dict):
-    def __init__(self, function, frame: FrameType) -> None:
+    def __init__(self, function, frame: Optional[FrameType]) -> None:
         self.function = function
         self.frame = frame
         self.results: Dict[str, Any] = {}