Add is_union_annotation() to helpers.py

Author: Rémy Lavainne

src/pyckson/helpers.py

@@ -2,7 +2,7 @@
 import sys
 from enum import Enum
 
-from typing import List, Set, Dict
+from typing import List, Set, Dict, Union
 
 try:
     from typing import _ForwardRef as ForwardRef
@@ -127,6 +127,12 @@ def is_typing_dict_annotation(annotation):
             return False
 
 
+def is_union_annotation(annotation) -> bool:
+    if hasattr(annotation, '__origin__'):
+        return annotation.__origin__ == Union
+    return False
+
+
 def using(attr):
     def class_decorator(cls):
         set_cls_attr(cls, PYCKSON_RULE_ATTR, attr)

tests/test_helpers.py

@@ -1,6 +1,7 @@
+from typing import Union
 from unittest import TestCase
 
-from pyckson.helpers import camel_case_name
+from pyckson.helpers import camel_case_name, is_union_annotation
 
 
 class CamelCaseNameTest(TestCase):
@@ -9,3 +10,11 @@ def test_should_do_nothing_on_simple_names(self):
 
     def test_should_camel_case_when_there_is_an_underscore(self):
         self.assertEqual(camel_case_name('foo_bar'), 'fooBar')
+
+
+def test_should_detect_union_annotation():
+    annotation = Union[str, int]
+
+    is_union = is_union_annotation(annotation)
+
+    assert is_union