Repository: s-knibbs/dataclasses-jsonschema (python)

Issue:
Support discriminators when used with Union types
It should be possible to use discriminators when decoding union fields. Currently discriminators only work when calling `from_dict` on the base class.

Example:

```python
@dataclass
class Base(JsonSchemaMixin, discriminator=True):
    pass


@dataclass
class Foo(Base):
    common_field: int


@dataclass
class Bar(Base):
    common_field: int
    other_field: Optional[int] = None


@dataclass
class Baz(JsonSchemaMixin):
    my_foo_bar: Union[Bar, Foo]
```

This doesn't currently work since, since data can be valid under both the `Foo` and `Bar` schemas and the discriminator field is not checked.

Pull request:
Support using discriminator to determine the Union variant
Support overlapping Union types with `anyOf`

Fixes #170 

Patch:
diff --git a/dataclasses_jsonschema/__init__.py b/dataclasses_jsonschema/__init__.py
index b6eb68d..b1c779f 100644
--- a/dataclasses_jsonschema/__init__.py
+++ b/dataclasses_jsonschema/__init__.py
@@ -391,7 +391,13 @@ class JsonSchemaMixin:
             cls.__mapped_fields = _get_fields_uncached()
         return cls.__mapped_fields  # type: ignore
 
-    def to_dict(self, omit_none: bool = True, validate: bool = False, validate_enums: bool = True) -> JsonDict:
+    def to_dict(
+        self,
+        omit_none: bool = True,
+        validate: bool = False,
+        validate_enums: bool = True,
+        schema_type: SchemaType = DEFAULT_SCHEMA_TYPE,
+    ) -> JsonDict:
         """Converts the dataclass instance to a JSON encodable dict, with optional JSON schema validation.
 
         If omit_none (default True) is specified, any items with value None are removed
@@ -414,7 +420,7 @@ class JsonSchemaMixin:
             data[self.__discriminator_name] = self.__class__.__name__
 
         if validate:
-            self._validate(data, validate_enums)
+            self._validate(data, validate_enums, schema_type)
         return data
 
     @classmethod
@@ -484,11 +490,14 @@ class JsonSchemaMixin:
         return decoder(field, field_type, value)
 
     @classmethod
-    def _validate(cls, data: JsonDict, validate_enums: bool = True):
+    def _validate(cls, data: JsonDict, validate_enums: bool = True, schema_type: SchemaType = DEFAULT_SCHEMA_TYPE):
+        if schema_type == SchemaType.OPENAPI_3 or schema_type == SchemaType.SWAGGER_V2:
+            warnings.warn("Only draft-04, draft-06 and draft-07 schema types are supported for validation")
+            schema_type = DEFAULT_SCHEMA_TYPE
+
         try:
             if fast_validation:
-                # TODO: Support validating with other schema types
-                schema_validator = cls.__compiled_schema.get(SchemaOptions(DEFAULT_SCHEMA_TYPE, validate_enums))
+                schema_validator = cls.__compiled_schema.get(SchemaOptions(schema_type, validate_enums))
                 if schema_validator is None:
                     formats = {}
                     for encoder in cls._field_encoders.values():
@@ -497,17 +506,23 @@ class JsonSchemaMixin:
                             formats[schema['format']] = schema['pattern']
 
                     schema_validator = fastjsonschema.compile(
-                        cls.json_schema(validate_enums=validate_enums), formats=formats
+                        cls.json_schema(schema_type=schema_type, validate_enums=validate_enums), formats=formats
                     )
-                    cls.__compiled_schema[SchemaOptions(DEFAULT_SCHEMA_TYPE, validate_enums)] = schema_validator
+                    cls.__compiled_schema[SchemaOptions(schema_type, validate_enums)] = schema_validator
                 schema_validator(data)
             else:
-                validate_func(data, cls.json_schema(validate_enums=validate_enums))
+                validate_func(data, cls.json_schema(schema_type=schema_type, validate_enums=validate_enums))
         except JsonSchemaValidationError as e:
             raise ValidationError(str(e)) from e
 
     @classmethod
-    def from_dict(cls: Type[T], data: JsonDict, validate=True, validate_enums: bool = True) -> T:
+    def from_dict(
+        cls: Type[T],
+        data: JsonDict,
+        validate: bool = True,
+        validate_enums: bool = True,
+        schema_type: SchemaType = DEFAULT_SCHEMA_TYPE,
+    ) -> T:
         """Returns a dataclass instance with all nested classes converted from the dict given"""
         if cls is JsonSchemaMixin:
             raise NotImplementedError
@@ -517,11 +532,14 @@ class JsonSchemaMixin:
                 for subclass in cls.__subclasses__():
                     if subclass.__name__ == data[cls.__discriminator_name]:
                         return subclass.from_dict(data, validate)
+                raise TypeError(
+                    f"Class '{cls.__name__}' does not match discriminator '{data[cls.__discriminator_name]}'"
+                )
 
         init_values: Dict[str, Any] = {}
         non_init_values: Dict[str, Any] = {}
         if validate:
-            cls._validate(data, validate_enums)
+            cls._validate(data, validate_enums, schema_type)
 
         for f in cls._get_fields():
             values = init_values if f.field.init else non_init_values
@@ -680,8 +698,8 @@ class JsonSchemaMixin:
             elif field_type_name == 'Union':
                 if schema_options.schema_type == SchemaType.SWAGGER_V2:
                     raise TypeError('Type unions unsupported in Swagger 2.0')
-                field_schema = {'oneOf': [cls._get_field_schema(variant, schema_options)[0] for variant in field_args]}
-                field_schema['oneOf'].sort(key=lambda item: item.get('type', ''))
+                field_schema = {'anyOf': [cls._get_field_schema(variant, schema_options)[0] for variant in field_args]}
+                field_schema['anyOf'].sort(key=lambda item: item.get('type', ''))
             elif field_type_name in MAPPING_TYPES:
                 field_schema = {'type': 'object'}
                 if field_args[1] is not Any:


Test patch:
diff --git a/tests/test_core.py b/tests/test_core.py
index 41daa72..9f390d7 100644
--- a/tests/test_core.py
+++ b/tests/test_core.py
@@ -132,7 +132,7 @@ BAR_SCHEMA = {
     'description': "Type with union field",
     'properties': {
         'a': {
-            'oneOf': [
+            'anyOf': [
                 {'$ref': '#/definitions/Point'},
                 {'type': 'string', 'enum': ['Monday', 'Tuesday', 'Wednesday', 'Thursday', 'Friday']},
             ]
@@ -621,7 +621,7 @@ def test_optional_union():
         "description": "Class with optional union",
         "type": "object",
         "properties": {
-            "a": {"oneOf": [{"type": "integer"}, {"type": "string"}]}
+            "a": {"anyOf": [{"type": "integer"}, {"type": "string"}]}
         }
     })
     assert Baz.json_schema() == expected_schema
@@ -1028,3 +1028,27 @@ def test_integer_dict_keys():
     c = Config({1: 'foo', 2: 'bar'})
     data = c.to_json()
     assert c == Config.from_json(data)
+
+
+def test_union_with_discriminator():
+    @dataclass
+    class Pet(JsonSchemaMixin, discriminator=True):
+        pass
+
+    @dataclass
+    class Cat(Pet):
+        breed: str
+
+    @dataclass
+    class Dog(Pet):
+        breed: str
+        walk_distance: Optional[float] = None
+
+    @dataclass
+    class Person(JsonSchemaMixin):
+        name: str
+        pet: Union[Cat, Dog]
+
+    data = Person(name="Joe", pet=Dog(breed="Pug")).to_dict()
+    p = Person.from_dict(data)
+    assert p.pet == Dog(breed="Pug")
diff --git a/tests/test_peps.py b/tests/test_peps.py
index 83d092d..c522615 100644
--- a/tests/test_peps.py
+++ b/tests/test_peps.py
@@ -16,10 +16,10 @@ def test_pep_604_types():
 
     schema = Post.json_schema()
     assert schema['properties']['tags'] == {
-        'oneOf': [{'type': 'array', 'items': {'type': 'string'}}, {'type': 'string'}]
+        'anyOf': [{'type': 'array', 'items': {'type': 'string'}}, {'type': 'string'}]
     }
     assert schema['properties']['metadata'] == {
-        'type': 'array', 'items': {'oneOf': [{'type': 'integer'}, {'type': 'string'}]}
+        'type': 'array', 'items': {'anyOf': [{'type': 'integer'}, {'type': 'string'}]}
     }
     assert schema['required'] == ['body', 'metadata']