Generate a JSON Schema specification from Python classes

Question

Introduction

Hello everyone !

I am trying to develop multi-agent models in Python3. So my approach is to create basic classes and derive them to more concrete and specific ones. For instance, a class Bike inherits from Vehicle, itself inheriting from a basic Agent class.

Problem

I want to offer a clear specification of my classes init parameters using JSON Schema (and also use them for validation), but I am struggling to automate their generation. Let's look at an example:

class Agent:
    
    SCHEMA = {
        "properties": {
            "agent_id": {
                "type": "string",
                "description": "unique identifier"
            },
            "network": {
                "type": "string",
                "description": "road network used by the agent to move"
            },
            "origin": {
                "type": "integer",
                "description": "origin position id",
            },
            "icon": {
                "type": "string",
                "description": "display icon"
            }
        },
        "required": ["agent_id", "network", "origin", "icon"]
    }
    
    def __init__(self, agent_id, network, origin, icon):
        self.id = agent_id
        self.network = network
        self.position = origin
        self.icon = icon
        
    def move(self, position):
        self.position = position


class User(Agent):
    SCHEMA = {
        # that's what i want in the end, but i don't want to duplicate the common properties
        "properties": {
            "agent_id": {
                "type": "string",
                "description": "unique identifier"
            },
            # notice that there is no "network" property
            "origin": {
                "type": "integer",
                "description": "origin position id",
            },
            "destination": {
                "type": "integer",
                "description": "destination position id",
            },
            "icon": {
                "type": "string",
                "description": "display icon"
            }
        },
        "required": ["agent_id", "origin", "destination", "icon"]
    }

    def __init__(self, agent_id, origin, destination, icon="user"):
        super().__init__(agent_id, "walk", origin, icon)
        self.destination = destination


class Vehicle(Agent):
    
    SCHEMA = {
        # one way to inherit the schema could be like this, but it has its flaws
        **super().SCHEMA,
        "seats": {
            "type": "integer",
            "description": "capacity of the vehicle"
        }
    }
    
    def __init__(self, agent_id, network, origin, seats, icon):
        super().__init__(agent_id, network, origin, icon)
        self.seats = seats
      
        
class Bike(Vehicle):

    # i want a schema here too, but without the "seats" prop
    # and maybe specify the default value for "icon" ?

    def __init__(self, agent_id, origin, icon="bike"):
        super().__init__(agent_id, "bike", origin, icon, 1)
        
        
class Car(Vehicle):
    
    # quite same question here

    def __init__(self, agent_id, origin, seats, icon="car"):
        super().__init__(agent_id, "drive", origin, icon, seats)

As you can see, I want to write and add specifications as I add new parameters to my classes. I would like to re-use the schemas from higher classes in order to reduce code duplication, but it's hard. I have some ideas, for instance what I proposed above, but it does not allow stripping a parameter schema from the parent class for instance.. Maybe using a method would give me more control on how the schemas are built.

The question

I would like to know if there is a library that allows doing this, or, otherwise, I would be glad to have some advice on how to accomplish this.

Answer 1

Pydantic can help you achieve this:

Example from the Pydantic documentation

from enum import Enum
from pydantic import BaseModel, Field


class FooBar(BaseModel):
    count: int
    size: float = None


class Gender(str, Enum):
    male = 'male'
    female = 'female'
    other = 'other'
    not_given = 'not_given'


class MainModel(BaseModel):
    """
    This is the description of the main model
    """

    foo_bar: FooBar = Field(...)
    gender: Gender = Field(None, alias='Gender')
    snap: int = Field(
        42,
        title='The Snap',
        description='this is the value of snap',
        gt=30,
        lt=50,
    )

    class Config:
        title = 'Main'


# this is equivalent to json.dumps(MainModel.schema(), indent=2):
print(MainModel.schema_json(indent=2))

Output:

# this is equivalent to json.dumps(MainModel.schema(), indent=2):
print(MainModel.schema_json(indent=2))

{
  "title": "Main",
  "description": "This is the description of the main model",
  "type": "object",
  "properties": {
    "foo_bar": {
      "$ref": "#/definitions/FooBar"
    },
    "Gender": {
      "$ref": "#/definitions/Gender"
    },
    "snap": {
      "title": "The Snap",
      "description": "this is the value of snap",
      "default": 42,
      "exclusiveMinimum": 30,
      "exclusiveMaximum": 50,
      "type": "integer"
    }
  },
  "required": [
    "foo_bar"
  ],
  "definitions": {
    "FooBar": {
      "title": "FooBar",
      "type": "object",
      "properties": {
        "count": {
          "title": "Count",
          "type": "integer"
        },
        "size": {
          "title": "Size",
          "type": "number"
        }
      },
      "required": [
        "count"
      ]
    },
    "Gender": {
      "title": "Gender",
      "description": "An enumeration.",
      "enum": [
        "male",
        "female",
        "other",
        "not_given"
      ],
      "type": "string"
    }
  }
}