Getting local variable annotations of function and generics (#895)

Author: vlad-perevezentsev

numba_typing/local_variable_type_annotations.py

@@ -0,0 +1,43 @@
+import ast
+import inspect
+import textwrap
+from pathlib import Path
+
+
+def get_variable_annotations(func):
+    """Get local variable annotations from the function."""
+    module_path = inspect.getfile(func)
+    func_global_variables = func.__globals__
+    func_global_variables.update(inspect.getclosurevars(func).nonlocals)
+    module_name = (Path(f'{module_path}').stem)
+    func_code = inspect.getsource(func)
+    func_code_with_dedent = textwrap.dedent(func_code)
+    func_tree = ast.parse(func_code_with_dedent)
+    analyzer = Analyzer(module_name, func_global_variables)
+    analyzer.visit(func_tree)
+    return analyzer.locals_parameter
+
+
+class Analyzer(ast.NodeVisitor):
+    def __init__(self, module_name, func_global_variables):
+        self.locals_parameter = {}
+        self.module_name = module_name
+        self.global_parameter = func_global_variables
+
+    def visit_AnnAssign(self, node):
+        target, annotation = node.target, node.annotation
+        if isinstance(annotation, ast.Subscript):  # containers and generics
+            try:
+                container_name = annotation.value.id
+                module_container_name = ''
+            except AttributeError:  # typing.
+                module_container_name = annotation.value.value.id + '.'
+                container_name = annotation.value.attr
+            if isinstance(annotation.slice.value, ast.Tuple):
+                types_as_str = ','.join(elt.id for elt in annotation.slice.value.elts)
+                exec_variables = f'{target.id} = [{module_container_name}{container_name}[{types_as_str}]]'
+            else:
+                exec_variables = f'{target.id} = [{module_container_name}{container_name}[{annotation.slice.value.id}]]'
+            exec(exec_variables, self.global_parameter, self.locals_parameter)
+        else:  # not containers
+            exec(f'{target.id} = [{annotation.id}]', self.global_parameter, self.locals_parameter)

numba_typing/test_variable_annotations.py

@@ -0,0 +1,90 @@
+from local_variable_type_annotations import get_variable_annotations
+import unittest
+from typing import Any, Union, List, Tuple, Dict, Iterable, Iterator, Generic, TypeVar
+import typing
+
+T = TypeVar('T')
+S = TypeVar('S', int, str)
+G = Generic[T, S]
+
+
+class A():
+    ...
+
+
+class B():
+    ...
+
+
+class C():
+    ...
+
+
+class TestAst(unittest.TestCase):
+    maxDiff = None
+
+    def test_get_variable_annotations_standard_types(self):
+        def test_func():
+            t = 1
+            t_int: int
+            t_str: str
+            t_float: float
+            t_bool: bool
+            t_bytes: bytes
+            t_list: list
+            t_dict: dict
+            t_tuple: tuple
+            t_set: set
+        result = get_variable_annotations(test_func)
+        expected_result = {'t_int': [int], 't_str': [str], 't_float': [float], 't_bool': [bool], 't_bytes': [bytes],
+                           't_list': [list], 't_dict': [dict], 't_tuple': [tuple], 't_set': [set]}
+        self.assertEqual(result, expected_result)
+
+    def test_get_variable_annotations_generic_types(self):
+        def test_func():
+            t_any: Any
+            t_union: Union[str, bytes]
+            t_optional: typing.Optional[float]
+            t_list: List[int]
+            t_tuple: Tuple[str]
+            t_dict: Dict[str, str]
+            t_iterable: Iterable[float]
+            t_iterator: Iterator[int]
+            t_generic: G
+            t_typevar_t: T
+            t_typevar_s: S
+            t_list_typevar: List[T]
+            t_tuple_typevar: Tuple[S]
+            t_dict_typevar: Dict[T, S]
+        result = get_variable_annotations(test_func)
+        expected_result = {'t_any': [Any], 't_union': [Union[str, bytes]], 't_optional': [typing.Optional[float]],
+                           't_list': [List[int]], 't_tuple': [Tuple[str]], 't_dict': [Dict[str, str]],
+                           't_iterable': [Iterable[float]], 't_iterator': [Iterator[int]], 't_generic': [G],
+                           't_typevar_t': [T], 't_typevar_s': [S], 't_list_typevar': [List[T]],
+                           't_tuple_typevar': [Tuple[S]], 't_dict_typevar': [Dict[T, S]]}
+        self.assertEqual(result, expected_result)
+
+    def test_get_variable_annotations_user_types(self):
+        def test_func():
+            t_class_a: A
+            t_class_b: B
+            t_class_c: C
+        result = get_variable_annotations(test_func)
+        expected_result = {'t_class_a': [A], 't_class_b': [B], 't_class_c': [C]}
+        self.assertEqual(result, expected_result)
+
+    def test_get_variable_annotations_non_locals(self):
+        def foo():
+            Q = TypeVar('Q')
+
+            def bar():
+                t_typevar: Q
+            return bar
+        test_func = foo()
+        result = get_variable_annotations(test_func)
+        expected_result = {'t_typevar': [Q]}
+        self.assertEqual(result, expected_result)
+
+
+if __name__ == "__main__":
+    unittest.main()