Merge pull request #37 from fishtown-analytics/fix/py39-compat

(#34) Add support for Python 3.9
dbt-labs · Oct 8, 2020 · ae61097 · ae61097
2 parents d9c079a + ccc848e
commit ae61097
Showing 1 changed file with 9 additions and 5 deletions.
diff --git a/hologram/__init__.py b/hologram/__init__.py
@@ -35,6 +35,8 @@
 JsonEncodable = Union[int, float, str, bool, None]
 JsonDict = Dict[str, Any]
 
+OPTIONAL_TYPES = ["Union", "Optional"]
+
 
 class ValidationError(jsonschema.ValidationError):
     pass
@@ -83,7 +85,9 @@ def issubclass_safe(klass: Any, base: Type) -> bool:
 
 
 def is_optional(field: Any) -> bool:
-    if str(field).startswith("typing.Union"):
+    if str(field).startswith("typing.Union") or str(field).startswith(
+        "typing.Optional"
+    ):
         for arg in field.__args__:
             if isinstance(arg, type) and issubclass(arg, type(None)):
                 return True
@@ -333,7 +337,7 @@ def encoder(ft, v, __):
                 def encoder(_, v, __):
                     return v.value
 
-            elif field_type_name == "Union":
+            elif field_type_name in OPTIONAL_TYPES:
                 # Attempt to encode the field with each union variant.
                 # TODO: Find a more reliable method than this since in the case 'Union[List[str], Dict[str, int]]' this
                 # will just output the dict keys as a list
@@ -491,7 +495,7 @@ def decoder(_, ft, val):
                 def decoder(_, ft, val):
                     return ft.from_dict(val, validate=validate)
 
-            elif field_type_name == "Union":
+            elif field_type_name in OPTIONAL_TYPES:
                 # Attempt to decode the value using each decoder in turn
                 union_excs = (
                     AttributeError,
@@ -725,7 +729,7 @@ def _get_schema_for_type(
             field_schema.update(cls._encode_restrictions(restrictions))
 
         # if Union[..., None] or Optional[...]
-        elif type_name == "Union":
+        elif type_name in OPTIONAL_TYPES:
             field_schema = {
                 "oneOf": [
                     cls._get_field_schema(variant)[0]
@@ -828,7 +832,7 @@ def _get_field_definitions(cls, field_type: Any, definitions: JsonDict):
             cls._get_field_definitions(field_type.__args__[1], definitions)
         elif field_type_name == "PatternProperty":
             cls._get_field_definitions(field_type.TARGET_TYPE, definitions)
-        elif field_type_name == "Union":
+        elif field_type_name in OPTIONAL_TYPES:
             for variant in field_type.__args__:
                 cls._get_field_definitions(variant, definitions)
         elif cls._is_json_schema_subclass(field_type):