fix comments

Author: KsanaKozlova

numba_typing/tests/test_type_annotations.py

@@ -9,37 +9,42 @@ 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)
+        with self.assertRaises(SyntaxError) as raises:
+            type_annotations.get_func_annotations(foo)
+        self.assertIn('Not found annotation for parameter b', str(raises.exception))
 
     def test_get_cls_annotations(self):
-        class MyClass(object):
+        class TestClass(object):
             x: int = 3
             y: str = "string"
 
-            def __init__(self, x: str, y: int):
+            def __init__(self, x, y):
                 self.x = x
                 self.y = y
 
-        self.assertEqual(type_annotations.get_cls_annotations(MyClass), ({'x': int, 'y': str}, {}))
+        result = type_annotations.get_cls_annotations(TestClass)
+        expected = ({'x': [int], 'y': [str]}, {})
+        self.assertEqual(result, expected)
 
     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}))
+
+        expected_results = {
+            func_one: ({'a': [int], 'b': [int, float], 'c': [str]}, {}),
+            func_two: ({'a': [int], 'b': [str], 'c': [List[int]]}, {'a': 2, 'b': 'string', 'c': [1, 2, 3]}),
+            func_three: ({'a': [Dict[int, str]], 'b': [str], 'c': [int]}, {'b': 'string', 'c': 1}),
+        }
+        for f, expected in expected_results.items():
+            with self.subTest(func=f.__name__):
+                self.assertEqual(type_annotations.get_func_annotations(f), expected)
 
 
 if __name__ == '__main__':

numba_typing/type_annotations.py

@@ -11,11 +11,8 @@ 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}')
-        annot = param.annotation
-        if get_args_union(annot):
-            annotations[name] = get_args_union(annot)
-        else:
-            annotations[name] = annot
+
+        annotations[name] = get_annotation_types(param.annotation)
         if param.default != sig.empty:
             defaults[name] = param.default
 
@@ -26,18 +23,16 @@ def get_cls_annotations(cls):
     """Get annotations of class attributes."""
     annotations = get_type_hints(cls)
     for x in annotations:
-        if get_args_union(annotations[x]):
-            annotations[x] = get_args_union(annotations[x])
+        annotations[x] = get_annotation_types(annotations[x])
     return annotations, {}
 
 
-def get_args_union(annot):
+def get_annotation_types(annotation):
+    """Get types of passed annotation."""
     try:
-        annot.__origin__
-    except:
-        return None
-    else:
-        if annot.__origin__ is Union:
-            return list(annot.__args__)
-        else:
-            return None
+        if annotation.__origin__ is Union:
+            return list(annotation.__args__)
+    except AttributeError:
+        pass
+
+    return [annotation, ]