CODEWITHSUNDEEP
pythonenumspydantic
Dhruv Saraswat
Dhruv Saraswat

Reputation: 21

Export Pydantic Model, with all ENUM values expanded

I have a pydantic model ModelX, with member Enums defined. The script should run as is.

from enum import Enum
from pydantic import BaseModel,Field,schema_json_of
from typing import Optional


class FirstEnumList(str, Enum):
   VAL1="VAL1"
   VAL2="VAL2"
   VAL3="VAL3"

class SecondEnumList(str, Enum):
   VAL4="VAL4"
   VAL5="VAL5"

class ModelX(BaseModel):
   V: Optional[FirstEnumList]=Field(FirstEnumList.VAL1, alias='value1')
   A: Optional[str]
   B: Optional[int] 
   P: Optional[SecondEnumList]=Field(SecondEnumList.VAL4, alias='valueP')


print (schema_json_of(ModelX))

'''Output below for schema_json_of
{"title": "ParsingModel[ModelX]", "$ref": "#/definitions/ModelX", 
"definitions": 
    {
        "FirstEnumList": 
        {
            "title": "FirstEnumList", 
            "description": "An enumeration.", 
            "enum": ["VAL1", "VAL2", "VAL3"], 
            "type": "string"
        }, 
        "SecondEnumList": 
        {
            "title": "SecondEnumList", 
            "description": "An enumeration.", 
            "enum": ["VAL4", "VAL5"], 
            "type": "string"
        }, 
        "ModelX": 
        {
            "title": "ModelX", 
            "type": "object", 
            "properties": 
            {
                "value1": 
                {
                    "default": "VAL1", 
                    "allOf": [{"$ref": "#/definitions/FirstEnumList"}]
                }, 
                "A": {
                    "title": "A", 
                    "type": "string"
                    }, 
                "B": {
                    "title": "B", 
                    "type": "integer"
                    }, 
                "valueP": {
                    "default": "VAL4", 
                    "allOf": [{"$ref": "#/definitions/SecondEnumList"}]}
            }
        }
    }
}
'''

I want to export this model as a JSON with the list of all Enum values. Something like


{
  "value1": {
    "allOf": [
      "VAL1",
      "VAL2",
      "VAL3"
    ]
  },
  "A": {
    "type": "string"
  },
  "B": {
    "type": "integer"
  },
  "valueP": {
    "allOf": [
      "VAL4",
      "VAL5"
    ]
  }
}

I am looking to return this Json as a response of my api call, all Enum values, it can be under "allOf" key or some other better way. This way the response has all the possible enum values and the consumer of this response can use it to pre-validate requests, or even show them on a UI screen.

Let me know in case more details or explanation is required.

Upvotes: 1

Views: 724

Answers (1)

Yaakov Bressler
Yaakov Bressler

Reputation: 12088

Pydantic is configured to export json schema compliant with the following specifications: JSON Schema Core, JSON Schema Validation, OpenAPI.

These specs follow the design principle of reducing repeated elements. Child models are referenced with ref to avoid unnecessarily repeating model definitions.

You can override some elements of this by modifying a "sub-model" in the Field class by adding a custom title or default value. Only the overloaded values however will be included.

In your case, you will need to employ a workaround to manipulate the final state:

my_schema = ModelX.schema()

for k, v in my_schema["properties"].items():
    if ('allOf' in v and
        len(v['allOf']) == 1 and
        isinstance(v['allOf'][0], dict) and
        v['allOf'][0].get('$ref') is not None
    ):
        # get the values
        enum_key = v['allOf'][0]['$ref'].replace('#/definitions/', '')
        enum_vals = my_schema['definitions'][enum_key]['enum']

        # Set the values
        my_schema["properties"][k]['allOf'] = enum_vals

# Get enum values
print(my_schema)

{
  "value1": {
    "title": "value1",
    "default": "VAL1",
    "allOf": [
      "VAL1",
      "VAL2",
      "VAL3"
    ]
  },
  "A": {
    "title": "A",
    "type": "string"
  },
  "B": {
    "title": "B",
    "type": "integer"
  },
  "valueP": {
    "default": "VAL4",
    "allOf": [
      "VAL4",
      "VAL5"
    ]
  }
}

Upvotes: 1

Related Questions