Tidy up JSON schema generation

src/higlass_schema/cli.py

@@ -10,7 +10,7 @@
 
 
 def export(_: argparse.Namespace) -> None:
-    console.print_json(schema_json(indent=2))
+    print(schema_json())
 
 
 def check(args: argparse.Namespace) -> None:

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,7 @@ 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)
+    return Viewconf.model_json_schema(schema_generator=_GenerateJsonSchema)
 
 
 def schema_json(**kwargs):

src/higlass_schema/utils.py

@@ -1,37 +1,52 @@
 from typing import Any, Dict, TypeVar
 
+import pydantic_core.core_schema as core_schema
 from pydantic import BaseModel, TypeAdapter
+from pydantic._internal._core_utils import CoreSchemaOrField, is_core_schema
+from pydantic.json_schema import GenerateJsonSchema, JsonSchemaMode, JsonSchemaValue
 
 
-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
+class _GenerateJsonSchema(GenerateJsonSchema):
+    def field_title_should_be_set(self, schema: CoreSchemaOrField) -> bool:
+        return_value = super().field_title_should_be_set(schema)
+        if return_value and is_core_schema(schema):
+            return False
+        return return_value
 
-    defs = list(root_schema["$defs"].values())
-    if len(defs) != 1:
-        return root_schema
+    def nullable_schema(self, schema: core_schema.NullableSchema) -> JsonSchemaValue:
+        inner_json_schema = self.generate_inner(schema["schema"])
+        # ignore the nullable
+        return inner_json_schema
 
-    return defs[0]
+    def default_schema(self, schema: core_schema.WithDefaultSchema) -> JsonSchemaValue:
+        if schema.get("default") is None:
+            return self.generate_inner(schema["schema"])
+        return super().default_schema(schema)
 
+    def union_schema(self, schema: core_schema.UnionSchema) -> JsonSchemaValue:
+        return super().union_schema(schema)
 
-# Schema modifiers
-ModelT = TypeVar("ModelT", bound=BaseModel)
+    def generate(
+        self, schema: core_schema.CoreSchema, mode: JsonSchemaMode = "validation"
+    ) -> JsonSchemaValue:
+        json_schema = super().generate(schema, mode=mode)
+        json_schema["$schema"] = self.schema_dialect
+        json_schema["title"] = "HiGlass viewconf"
+        for d in json_schema.get("$defs", {}).values():
+            d.pop("title", None)
+
+        return json_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)
+# 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
+def get_schema_of(type_: object):
+    schema = TypeAdapter(type_).json_schema(schema_generator=_GenerateJsonSchema)
+    # remove the title and $schema fields
     schema.pop("title", None)
+    schema.pop("$schema", None)
     return schema