add unit tests and handler for union

Author: KsanaKozlova

numba_typing/tests/__init__.py

@@ -0,0 +1,46 @@
+import unittest
+import type_annotations
+from typing import Union, Dict, List
+
+
+class TestTypeAnnotations(unittest.TestCase):
+
+    def test_get_func_annotations_exception(self):
+
+        def foo(a: int, b, c: str = "string"):
+            pass
+        self.assertRaises(SyntaxError,  type_annotations.get_func_annotations, foo)
+
+    def test_get_cls_annotations(self):
+        class MyClass(object):
+            x: int = 3
+            y: str = "string"
+
+            def __init__(self, x: str, y: int):
+                self.x = x
+                self.y = y
+
+        self.assertEqual(type_annotations.get_cls_annotations(MyClass), ({'x': int, 'y': str}, {}))
+
+    def test_get_func_annotations(self):
+        def func_one(a: int, b: Union[int, float], c: str):
+            pass
+        with self.subTest("annotations"):
+            self.assertEqual(type_annotations.get_func_annotations(func_one),
+                             ({'a': int, 'b': [int, float], 'c': str}, {}))
+
+        def func_two(a: int = 2, b: str = "string", c: List[int] = [1, 2, 3]):
+            pass
+        with self.subTest("annotations and all default values"):
+            self.assertEqual(type_annotations.get_func_annotations(func_two),
+                             ({'a': int, 'b': str, 'c': List[int]}, {'a': 2, 'b': 'string', 'c': [1, 2, 3]}))
+
+        def func_three(a: Dict[int, str], b: str = "string", c: int = 1):
+            pass
+        with self.subTest("annotations and not all default values"):
+            self.assertEqual(type_annotations.get_func_annotations(func_three),
+                             ({'a': Dict[int, str], 'b': str, 'c': int}, {'b': 'string', 'c': 1}))
+
+
+if __name__ == '__main__':
+    unittest.main()

numba_typing/tests/test_type_annotations.py

@@ -1,5 +1,5 @@
 from inspect import signature
-from typing import get_type_hints
+from typing import get_type_hints, Union
 
 
 def get_func_annotations(func):
@@ -11,8 +11,11 @@ def get_func_annotations(func):
     for name, param in sig.parameters.items():
         if param.annotation == sig.empty:
             raise SyntaxError(f'Not found annotation for parameter {name}')
-
-        annotations[name] = param.annotation
+        annot = param.annotation
+        if get_args_union(annot):
+            annotations[name] = get_args_union(annot)
+        else:
+            annotations[name] = annot
         if param.default != sig.empty:
             defaults[name] = param.default
 
@@ -21,11 +24,20 @@ def get_func_annotations(func):
 
 def get_cls_annotations(cls):
     """Get annotations of class attributes."""
-    return get_type_hints(cls)
-
-
-if __name__ == '__main__':
-    def foo(a: int, b: int = 3):
-        pass
-
-    print(get_func_annotations(foo))
+    annotations = get_type_hints(cls)
+    for x in annotations:
+        if get_args_union(annotations[x]):
+            annotations[x] = get_args_union(annotations[x])
+    return annotations, {}
+
+
+def get_args_union(annot):
+    try:
+        annot.__origin__
+    except:
+        return None
+    else:
+        if annot.__origin__ is Union:
+            return list(annot.__args__)
+        else:
+            return None