Add support for Union typing

Author: explorigin

pydanticrud/backends/dynamodb.py

@@ -84,7 +84,6 @@ def _to_epoch_float(dt):
     "boolean": lambda d: 1 if d else 0,
     "object": json.dumps,
     "array": json.dumps,
-    "anyOf": str,  # FIXME - this could be more complicated. This is a hacky fix.
 }
 
 
@@ -98,7 +97,6 @@ def do_nothing(x):
     "boolean": bool,
     "object": json.loads,
     "array": json.loads,
-    "anyOf": do_nothing,  # FIXME - this could be more complicated. This is a hacky fix.
 }
 
 
@@ -123,20 +121,28 @@ def __init__(self, schema):
         self.properties = schema["properties"]
         self.definitions = schema.get("definitions")
 
-    def _get_type(self, field_name):
-        if "$ref" in self.properties[field_name]:
-            def_name = self.properties[field_name]["$ref"].split('/')[-1]
-            return self.definitions[def_name].get("type")
-        return self.properties[field_name].get("type", "anyOf")
+    def _get_type_possibilities(self, field_name) -> Set[str]:
+        field_properties = self.properties[field_name]
+
+        possible_types = []
+        if "anyOf" in field_properties:
+            possible_types.extend([r.get("$ref", r.get("type")) for r in field_properties["anyOf"]])
+        else:
+            possible_types.append(field_properties.get("$ref", field_properties.get("type")))
+        return set([
+            (self.definitions[t.split('/')[-1]].get("type") if t.startswith('#/') else t)
+            for t in possible_types
+        ])
 
     def _serialize_field(self, field_name, value):
-        field_type = self._get_type(field_name)
-        try:
-            if value is not None:
-                return SERIALIZE_MAP[field_type](value)
-        except KeyError:
-            log.debug(f"No serializer for field_type {field_type}")
-            return value  # do nothing but log it.
+        field_types = self._get_type_possibilities(field_name)
+        if value is not None:
+            for t in field_types:
+                try:
+                    return SERIALIZE_MAP[t](value)
+                except (ValueError, TypeError, KeyError):
+                    pass
+        return value
 
     def serialize_record(self, data_dict) -> dict:
         """
@@ -148,13 +154,14 @@ def serialize_record(self, data_dict) -> dict:
         }
 
     def _deserialize_field(self, field_name, value):
-        field_type = self._get_type(field_name)
-        try:
-            if value is not None:
-                return DESERIALIZE_MAP[field_type](value)
-        except KeyError:
-            log.debug(f"No deserializer for field_type {field_type}")
-            return value  # do nothing but log it.
+        field_types = self._get_type_possibilities(field_name)
+        if value is not None:
+            for t in field_types:
+                try:
+                    return DESERIALIZE_MAP[t](value)
+                except (ValueError, TypeError, KeyError):
+                    pass
+        return value
 
     def deserialize_record(self, data_dict) -> dict:
         """
@@ -252,7 +259,7 @@ def initialize(self):
                 {
                     "AttributeName": attr,
                     "AttributeType": DYNAMO_TYPE_MAP.get(
-                        schema["properties"][attr].get("type", "anyOf"), "S"
+                        schema["properties"][attr].get("type"), "S"
                     ),
                 }
                 for attr in self.possible_keys

tests/test_dynamodb.py

@@ -1,4 +1,4 @@
-from typing import Dict, List, Optional
+from typing import Dict, List, Optional, Union
 from base64 import b64decode, b64encode
 from decimal import Decimal
 from datetime import datetime
@@ -80,11 +80,17 @@ class Ticket(PydanticBaseModel):
     number: str
 
 
+class SomethingElse(PydanticBaseModel):
+    herp: bool
+    derp: int
+
+
 class NestedModel(BaseModel):
     account: str
     sort_date_key: str
     expires: str
     ticket: Optional[Ticket]
+    other: Union[Ticket, SomethingElse]
 
     class Config:
         title = "NestedModelTitle123"
@@ -141,9 +147,20 @@ def nested_model_data_generator(include_ticket=True, **kwargs):
         expires=future_datetime(days=1, hours=random.randint(1, 12), minutes=random.randint(1, 58)).isoformat(),
         ticket={
             'created_time': random_datetime().isoformat(),
-            'number': random.randint(0, 1000)
+            'number': str(random.randint(0, 1000))
+
+        } if include_ticket else None,
+        other=random.choice([
+            {
+                'created_time': random_datetime().isoformat(),
+                'number': str(random.randint(0, 1000))
 
-        } if include_ticket else None
+            }, {
+                'herp': random.choice([True, False]),
+                'derp': random.randint(0, 1000)
+
+            }
+        ])
     )
     data.update(kwargs)
     return data
@@ -464,6 +481,7 @@ def test_query_with_nested_model(dynamo, nested_query_data):
         assert isinstance(m.ticket, Ticket)
         assert m.ticket.created_time is not None
         assert m.ticket.number is not None
+        assert isinstance(m.other, (Ticket, SomethingElse))
 
 
 def test_query_with_nested_model_optional(dynamo, nested_query_data_empty_ticket):
@@ -486,5 +504,3 @@ def test_get_alias_model_data(dynamo, alias_query_data):
     data = alias_model_data_generator()
     res = AliasKeyModel.get(alias_query_data[0]['name'])
     assert res.dict(by_alias=True) == alias_query_data[0]
-
-