Tidy up JSON schema generation (#17)

Co-authored-by: Nezar Abdennur <nabdennur@gmail.com>

Author: manzt

pyproject.toml

@@ -15,8 +15,8 @@ dependencies = ["pydantic>=2.0", "rich>=13.0.0"]
 [project.urls]
 homepage = "https://github.com/higlass/higlass-schema"
 
-[tool.uv]
-dev-dependencies = ["black", "pytest", "ruff"]
+[dependency-groups]
+dev = ["black", "pytest", "ruff"]
 
 [project.scripts]
 higlass-schema = "higlass_schema.cli:main"

src/higlass_schema/cli.py

@@ -9,8 +9,8 @@
 console = Console()
 
 
-def export(_: argparse.Namespace) -> None:
-    console.print_json(schema_json(indent=2))
+def export(args: argparse.Namespace) -> None:
+    print(schema_json(indent=args.indent))
 
 
 def check(args: argparse.Namespace) -> None:
@@ -30,7 +30,8 @@ def check(args: argparse.Namespace) -> None:
             console.print_exception()
 
         console.print(
-            f"{msg} Run [white]`hgschema check --verbose`[/white] for more details.",
+            f"{msg} Run [white]`higlass-schema check --verbose`[/white] for "
+            "more details.",
             style="yellow",
         )
         sys.exit(1)
@@ -43,6 +44,7 @@ def main():
 
     # export
     parser_export = subparsers.add_parser("export")
+    parser_export.add_argument("--indent", type=int)
     parser_export.set_defaults(func=export)
 
     # check

src/higlass_schema/schema.py

@@ -1,5 +1,4 @@
 import json
-from collections import OrderedDict
 from typing import (
     Any,
     Dict,
@@ -14,18 +13,14 @@
 
 from pydantic import BaseModel as PydanticBaseModel
 from pydantic import ConfigDict, Field, RootModel, model_validator
-from pydantic.json_schema import GenerateJsonSchema
 from typing_extensions import Annotated, Literal, TypedDict
 
-from .utils import exclude_properties_titles, get_schema_of, simplify_enum_schema
+from .utils import _GenerateJsonSchema, get_schema_of
 
 
 # Override Basemodel
 class BaseModel(PydanticBaseModel):
-    model_config = ConfigDict(
-        validate_assignment=True,
-        json_schema_extra=lambda s, _: exclude_properties_titles(s),
-    )
+    model_config = ConfigDict(validate_assignment=True)
 
     # nice repr if printing with rich
     def __rich_repr__(self):
@@ -116,14 +111,13 @@ class Overlay(BaseModel):
 
 
 # We'd rather have tuples in our final model, because a
-# __root__ model is clunky from a python user perspective.
+# RootModel is clunky from a python user perspective.
 # We create this class to get validation for free in `root_validator`
 class _LockEntryModel(RootModel[LockEntry]):
     pass
 
 
 def _lock_schema_extra(schema: Dict[str, Any], _: Any) -> None:
-    exclude_properties_titles(schema)
     schema["additionalProperties"] = get_schema_of(LockEntry)
 
 
@@ -160,7 +154,6 @@ class _ValueScaleLockEntryModel(RootModel[ValueScaleLockEntry]):
 
 
 def _value_scale_lock_schema_extra(schema: Dict[str, Any], _: Any) -> None:
-    exclude_properties_titles(schema)
     schema["additionalProperties"] = get_schema_of(ValueScaleLockEntry)
 
 
@@ -191,14 +184,7 @@ def validate_locks(cls, values: Dict[str, Any]):
         return values
 
 
-def _axis_specific_lock_schema_extra(schema: Dict[str, Any], _: Any) -> None:
-    exclude_properties_titles(schema)
-    schema["properties"]["axis"] = simplify_enum_schema(schema["properties"]["axis"])
-
-
 class AxisSpecificLock(BaseModel):
-    model_config = ConfigDict(json_schema_extra=_axis_specific_lock_schema_extra)
-
     axis: Literal["x", "y"]
     lock: str
 
@@ -249,15 +235,8 @@ class Data(BaseModel):
     tiles: Optional[Tile] = None
 
 
-def _base_track_schema_extra(schema, _):
-    exclude_properties_titles(schema)
-    props = schema["properties"]
-    if "enum" in props["type"] or "allOf" in props["type"]:
-        props["type"] = simplify_enum_schema(props["type"])
-
-
 class BaseTrack(BaseModel, Generic[TrackTypeT]):
-    model_config = ConfigDict(extra="allow", json_schema_extra=_base_track_schema_extra)
+    model_config = ConfigDict(extra="allow")
 
     type: TrackTypeT
     uid: Optional[str] = None
@@ -500,11 +479,7 @@ class View(BaseModel, Generic[TrackT]):
 class Viewconf(BaseModel, Generic[ViewT]):
     """Root object describing a HiGlass visualization."""
 
-    model_config = ConfigDict(
-        extra="forbid",
-        title="HiGlass viewconf",
-        json_schema_extra=lambda s, _: exclude_properties_titles(s),
-    )
+    model_config = ConfigDict(extra="forbid")
 
     editable: Optional[bool] = True
     viewEditable: Optional[bool] = True
@@ -521,21 +496,10 @@ class Viewconf(BaseModel, Generic[ViewT]):
 
 
 def schema():
-    root = Viewconf.model_json_schema()
-
-    # remove titles in defintions
-    for d in root["$defs"].values():
-        d.pop("title", None)
-
-    # nice ordering, insert additional metadata
-    ordered_root = OrderedDict(
-        [
-            ("$schema", GenerateJsonSchema.schema_dialect),
-            *root.items(),
-        ]
-    )
-
-    return dict(ordered_root)
+    json_schema = Viewconf.model_json_schema(schema_generator=_GenerateJsonSchema)
+    json_schema["$schema"] = _GenerateJsonSchema.schema_dialect
+    json_schema["title"] = "HiGlass viewconf"
+    return json_schema
 
 
 def schema_json(**kwargs):

src/higlass_schema/utils.py

@@ -1,41 +1,101 @@
-from typing import Any, Dict, TypeVar
+from __future__ import annotations
 
+from typing import TYPE_CHECKING, Any, TypeVar, Union
+
+import pydantic_core.core_schema as core_schema
 from pydantic import BaseModel, TypeAdapter
+from pydantic.json_schema import GenerateJsonSchema, JsonSchemaMode, JsonSchemaValue
 
+if TYPE_CHECKING:
+    from typing import TypeGuard
 
-def simplify_schema(root_schema: Dict[str, Any]) -> Dict[str, Any]:
-    """Lift defintion reference to root if only definition"""
-    # type of root is not a reference to a definition
-    if "$ref" not in root_schema:
-        return root_schema
+### Vendored from pydantic._internal._core_utils
 
-    defs = list(root_schema["$defs"].values())
-    if len(defs) != 1:
-        return root_schema
+CoreSchemaField = Union[
+    core_schema.ModelField,
+    core_schema.DataclassField,
+    core_schema.TypedDictField,
+    core_schema.ComputedField,
+]
 
-    return defs[0]
+CoreSchemaOrField = Union[core_schema.CoreSchema, CoreSchemaField]
 
+_CORE_SCHEMA_FIELD_TYPES = {
+    "typed-dict-field",
+    "dataclass-field",
+    "model-field",
+    "computed-field",
+}
 
-# Schema modifiers
-ModelT = TypeVar("ModelT", bound=BaseModel)
 
+def is_core_schema(
+    schema: CoreSchemaOrField,
+) -> TypeGuard[core_schema.CoreSchema]:
+    return schema["type"] not in _CORE_SCHEMA_FIELD_TYPES
+
+
+### End vendored code
+
+
+class _GenerateJsonSchema(GenerateJsonSchema):
+    def field_title_should_be_set(self, schema: CoreSchemaOrField) -> bool:
+        """Check if the title should be set for a field.
+
+        Override the default implementation to not set the title for core
+        schemas. Makes the final schema more readable by removing
+        redundant titles. Explicit Field(title=...) can still be used.
+        """
+        return_value = super().field_title_should_be_set(schema)
+        if return_value and is_core_schema(schema):
+            return False
+        return return_value
+
+    def nullable_schema(self, schema: core_schema.NullableSchema) -> JsonSchemaValue:
+        """Generate a JSON schema for a nullable schema.
 
-def exclude_properties_titles(schema: Dict[str, Any]) -> None:
-    """Remove automatically generated tiles for pydantic classes."""
-    for prop in schema.get("properties", {}).values():
-        prop.pop("title", None)
+        This overrides the default implementation to ignore the nullable
+        and generate a more simple schema. All the Optional[T] fields
+        are converted to T (instead of the
+        default {"anyOf": [{"type": "null"}, {"type": "T"}]}).
+        """
+        return self.generate_inner(schema["schema"])
+
+    def default_schema(self, schema: core_schema.WithDefaultSchema) -> JsonSchemaValue:
+        """Generate a JSON schema for a schema with a default value.
+
+        Similar to above, this overrides the default implementation to
+        not explicity set {"default": null} in the schema when the field
+        is Optional[T] = None.
+        """
+        if schema.get("default") is None:
+            return self.generate_inner(schema["schema"])
+        return super().default_schema(schema)
+
+    def generate(
+        self, schema: core_schema.CoreSchema, mode: JsonSchemaMode = "validation"
+    ) -> JsonSchemaValue:
+        """Generate a JSON schema.
+
+        This overrides the default implementation to remove the titles
+        from the definitions. This makes the final schema more readable.
+        """
+
+        json_schema = super().generate(schema, mode=mode)
+        # clear the titles from the definitions
+        for d in json_schema.get("$defs", {}).values():
+            d.pop("title", None)
+        return json_schema
+
+
+# Schema modifiers
+ModelT = TypeVar("ModelT", bound=BaseModel)
 
 
-def get_schema_of(type_: Any):
-    schema = TypeAdapter(type_).json_schema()
-    schema = simplify_schema(schema)
-    exclude_properties_titles(schema)
-    # remove autogenerated title
-    schema.pop("title", None)
-    return schema
+def get_schema_of(type_: object):
+    return TypeAdapter(type_).json_schema(schema_generator=_GenerateJsonSchema)
 
 
-def simplify_enum_schema(schema: Dict[str, Any]):
+def simplify_enum_schema(schema: dict[str, Any]):
     # reduce union of enums into single enum
     if "anyOf" in schema:
         enum = []