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

Issue:
Support for swagger / openapi discriminators
Swagger 2.0 and OpenAPI 3 have support for polymorphism via the `discriminator` keyword.

Support could be added for generating this automatically via a discriminator argument on the base class:

```python

@dataclass
class Pet(JsonSchemaMixin, discriminator=True)
    name: str

@dataclass
class Cat(Pet):
    hunting_skill: Literal["clueless", "lazy", "adventurous", "aggressive"] = "lazy"
```

This would generate the following schema in OpenAPI 3:
```json
      "Pet": {
        "type": "object",
        "discriminator": {
          "propertyName": "petType"
        },
        "properties": {
          "name": {
            "type": "string"
          },
          "petType": {
            "type": "string"
          }
        },
        "required": [
          "name",
          "petType"
        ]
      },
      "Cat": {
        "allOf": [
          {
            "$ref": "#/components/schemas/Pet"
          },
          {
            "type": "object",
            "properties": {
              "hunting_skill": {
                "type": "string",
                "default": "lazy",
                "enum": [
                  "clueless",
                  "lazy",
                  "adventurous",
                  "aggressive"
                ]
              }
            },
            "required": [
              "huntingSkill"
            ]
          }
        ]
      },
```

When serialising to json via `to_dict`, the `petType` property name will be automatically set to the name of the class.

Optionally, the discriminator property name that gets created could be overriden by passing a string as the discriminator argument to the base class:

```python
@dataclass
class Pet(JsonSchemaMixin, discriminator="propertyName"):
    ...
```

This issue is blocked until #40 is resolved.

Pull request description:
Add support for OpenAPI 3 discriminators
 Closes #69 

Patch:
diff --git a/dataclasses_jsonschema/__init__.py b/dataclasses_jsonschema/__init__.py
index 391302f..b17098d 100644
--- a/dataclasses_jsonschema/__init__.py
+++ b/dataclasses_jsonschema/__init__.py
@@ -179,8 +179,15 @@ class JsonSchemaMixin:
     __encode_cache: Dict[Any, _ValueEncoder]
     __decode_cache: Dict[Any, _ValueDecoder]
     __mapped_fields: List[Tuple[Field, str]]
+    __discriminator_name: Optional[str]
+    # True if __discriminator_name is inherited from the base class
+    __discriminator_inherited: bool
 
-    def __init_subclass__(cls):
+    @classmethod
+    def _discriminator(cls) -> Optional[str]:
+        return cls.__discriminator_name
+
+    def __init_subclass__(cls, discriminator: Optional[Union[str, bool]] = None):
         # Initialise caches
         cls.__schema = {}
         cls.__compiled_schema = {}
@@ -188,6 +195,24 @@ class JsonSchemaMixin:
         cls.__encode_cache = {}
         cls.__decode_cache = {}
         cls.__mapped_fields = []
+        cls.__discriminator_inherited = False
+        if discriminator is not None:
+            cls.__discriminator_name = discriminator if isinstance(discriminator, str) else f"{cls.__name__}Type"
+        else:
+            dataclass_bases = [
+                klass for klass in cls.__bases__ if is_dataclass(klass) and issubclass(klass, JsonSchemaMixin)
+            ]
+            if len(dataclass_bases) > 0:
+                base_discriminators = [
+                    base._discriminator() for base in dataclass_bases if base._discriminator() is not None
+                ]
+                if len(base_discriminators):
+                    if len(base_discriminators) > 1:
+                        raise TypeError("Multiple bases with discriminators is unsupported")
+                    cls.__discriminator_name = base_discriminators[0]
+                    cls.__discriminator_inherited = True
+            else:
+                cls.__discriminator_name = None
 
     @classmethod
     def field_mapping(cls) -> Dict[str, str]:
@@ -306,6 +331,10 @@ class JsonSchemaMixin:
             if value is NULL:
                 value = None
             data[target_field] = value
+
+        if self.__discriminator_name is not None:
+            data[self.__discriminator_name] = self.__class__.__name__
+
         if validate:
             self._validate(data)
         return data
@@ -392,6 +421,12 @@ class JsonSchemaMixin:
         if cls is JsonSchemaMixin:
             raise NotImplementedError
 
+        if cls.__discriminator_name is not None and cls.__discriminator_name in data:
+            if data[cls.__discriminator_name] != cls.__name__:
+                for subclass in cls.__subclasses__():
+                    if subclass.__name__ == data[cls.__discriminator_name]:
+                        return subclass.from_dict(data, validate)
+
         init_values: Dict[str, Any] = {}
         non_init_values: Dict[str, Any] = {}
         if validate:
@@ -650,6 +685,12 @@ class JsonSchemaMixin:
                 'properties': properties
             }
 
+            if cls.__discriminator_name is not None and \
+                    schema_type == SchemaType.OPENAPI_3 and \
+                    not cls.__discriminator_inherited:
+                schema['discriminator'] = {'propertyName': cls.__discriminator_name}
+                required.append(cls.__discriminator_name)
+
             # Needed for Draft 04 backwards compatibility
             if len(required) == 0:
                 del schema["required"]


Test patch:
diff --git a/tests/test_core.py b/tests/test_core.py
index a6aeb3d..793e572 100644
--- a/tests/test_core.py
+++ b/tests/test_core.py
@@ -640,3 +640,45 @@ def test_underscore_fields():
     expected_obj = Album(_id=5, name="Foo")
     assert expected_data == expected_obj.to_dict()
     assert expected_obj == Album.from_dict(expected_data)
+
+
+def test_discriminators():
+    @dataclass
+    class Pet(JsonSchemaMixin, discriminator=True):
+        """A generic pet"""
+        name: str
+
+    @dataclass
+    class Dog(Pet):
+        """A dog"""
+        breed: str
+
+    expected_dog_schema = {
+        "Dog": {
+            "description": "A dog",
+            "allOf": [
+                {"$ref": "#/components/schemas/Pet"},
+                {
+                    "type": "object",
+                    "properties": {"breed": {"type": "string"}},
+                    "required": ["breed"]
+                },
+            ]
+        },
+        "Pet": {
+            "description": "A generic pet",
+            "type": "object",
+            "properties": {"name": {"type": "string"}},
+            "required": ["name", "PetType"],
+            "discriminator": {"propertyName": "PetType"}
+        }
+    }
+
+    assert Dog.json_schema(embeddable=True, schema_type=SchemaType.OPENAPI_3) == expected_dog_schema
+    expected_dog = {
+        "PetType": "Dog",
+        "name": "Fido",
+        "breed": "Dalmation"
+    }
+    assert Dog(name="Fido", breed="Dalmation").to_dict() == expected_dog
+    assert Dog(name="Fido", breed="Dalmation") == Pet.from_dict(expected_dog)