def prettify_json_schema()

in schemas/generate_json_schema.py [0:0]

• 49 lines of code
• 16 McCabe index (conditional complexity)

def prettify_json_schema(schema: dict[str, Any]) -> dict[str, Any]:
    # Add a $schema field.
    pretty_schema = {
        "$schema": "https://json-schema.org/draft/2019-09/schema",
    }

    # Re-order the properties in the dict so that they are in a sensible order
    # for humans consuming these schemas.

    # Use this order for top-level keys.
    key_order = [
        "title",
        "description",
        "type",
        "properties",
        "required",
        "additionalProperties",
        "if",
        "then",
        "$defs",
    ]

    # If there are any other keys not listed above, splice them in before $defs.
    key_order = [
        *key_order[:-1],
        *(set(schema.keys()) - set(key_order)),
        key_order[-1],
    ]

    pretty_schema.update({key: schema[key] for key in key_order if key in schema})

    # Assert that the schemas have not structurally changed.
    #
    # We have to add the $schema field back to the original schema for comparison.
    schema["$schema"] = pretty_schema["$schema"]
    assert schema == pretty_schema

    # Next, lets walk the schema and remove attributes we don't care about.
    def _walk_objects(objs: Iterable[dict[str, Any]]):
        for obj in objs:
            _walk_object(obj)

    def _walk_object(obj: dict[str, Any], top_level: bool = False):
        # All but the top-level title will be auto-generated base on field names. They are
        # not useful.
        if not top_level:
            obj.pop("title", None)

        # We don't support defaults.
        obj.pop("default", None)

        # This is an OpenAPI extension and it leads to incorrect code generation in our
        # case (due to using a boolean discriminator).
        obj.pop("discriminator", None)

        # Strip newlines from descriptions.
        if description := obj.get("description"):
            obj["description"] = NEWLINES_RE.sub(" ", description)

        # Remove redundant enum entries for constants.
        if obj.get("const") is not None:
            obj.pop("enum", None)

        match obj.get("type"):
            case "object":
                if properties := obj.get("properties"):
                    _walk_objects(properties.values())

            case "array":
                if items := obj.get("items"):
                    _walk_object(items)

        for group_key in ("allOf", "anyOf", "oneOf"):
            if group := obj.get(group_key):
                _walk_objects(group)

    _walk_object(pretty_schema, top_level=True)
    if defs := pretty_schema.get("$defs"):
        _walk_objects(defs.values())

    return pretty_schema