Test JSON Schema validation; SI base units; cleanup

* Make sure that the JSON representation and JSON Schema work with a 3rd party validator `jsonschema`
* Remove unused code -- apparently WrapValidator is doing most of the job
* Convert to SI base units for JSON representation -- reduce variability of output

pyproject.toml

@@ -30,7 +30,7 @@ authors = [
 repository = "https://github.com/LiberTEM/LiberTEM-schema"
 
 [project.optional-dependencies]
-test = ["pytest", "pytest-cov"]
+test = ["pytest", "pytest-cov", "jsonschema"]
 
 [tool.setuptools.dynamic]
 version = {attr = "libertem_schema.__version__"}

src/libertem_schema/__init__.py

@@ -1,18 +1,15 @@
-from typing import Any, Tuple, Sequence
-from numbers import Number
+from typing import Any, Sequence
 
 from typing_extensions import Annotated
 from pydantic_core import core_schema
 from pydantic import (
     BaseModel,
     GetCoreSchemaHandler,
-    GetJsonSchemaHandler,
     WrapValidator,
     ValidationInfo,
     ValidatorFunctionWrapHandler,
 )
 
-from pydantic.json_schema import JsonSchemaValue
 import pint
 
 
@@ -31,44 +28,26 @@ class DimensionError(ValueError):
 ])
 
 
+def to_tuple(q: pint.Quantity):
+    base = q.to_base_units()
+    return (float(base.magnitude), str(base.units))
+
+
 class PintAnnotation:
     @classmethod
     def __get_pydantic_core_schema__(
         cls,
         _source_type: Any,
         _handler: GetCoreSchemaHandler,
     ) -> core_schema.CoreSchema:
-        def validate_from_tuple(value: Tuple[Number, str]) -> pint.Quantity:
-            m, u = value
-            return m * ureg(u)
-
-        from_tuple_schema = core_schema.chain_schema(
-            [
-                _pint_base_repr,
-                core_schema.no_info_plain_validator_function(validate_from_tuple),
-            ]
-        )
-
         return core_schema.json_or_python_schema(
-            json_schema=from_tuple_schema,
-            python_schema=core_schema.chain_schema(
-                [
-                    # check if it's an instance first before doing any further work
-                    core_schema.is_instance_schema(pint.Quantity),
-                ]
-            ),
+            json_schema=_pint_base_repr,
+            python_schema=core_schema.is_instance_schema(pint.Quantity),
             serialization=core_schema.plain_serializer_function_ser_schema(
-                lambda instance: (float(instance.m), str(instance.u))
+                to_tuple
             ),
         )
 
-    @classmethod
-    def __get_pydantic_json_schema__(
-        cls, _core_schema: core_schema.CoreSchema, handler: GetJsonSchemaHandler
-    ) -> JsonSchemaValue:
-        # Use the same schema that would be used for the tuple
-        return handler(_pint_base_repr)
-
 
 _length_dim = ureg.meter.dimensionality
 _angle_dim = ureg.radian.dimensionality
@@ -83,9 +62,9 @@ def is_matching(
         if isinstance(q, pint.Quantity):
             pass
         elif isinstance(q, Sequence):
-            m, u = q
+            magnitude, unit = q
             # Turn into Quantity: measure * unit
-            q = m * ureg(u)
+            q = magnitude * ureg(unit)
         else:
             raise ValueError(f"Don't know how to interpret type {type(q)}.")
         # Check dimension

tests/test_schemas.py

@@ -1,6 +1,9 @@
 import pprint
+import json
 
+import jsonschema.exceptions
 import pytest
+import jsonschema
 
 from pint import UnitRegistry, Quantity
 from pydantic_core import from_json
@@ -34,6 +37,31 @@ def test_smoke():
     assert res == params
 
 
+def test_angstrom():
+    params = Simple4DSTEMParams(
+        overfocus=0.0015 * ureg.meter,
+        scan_pixel_pitch=10000 * ureg.angstrom,
+        camera_length=0.15 * ureg.meter,
+        detector_pixel_pitch=0.000050 * ureg.meter,
+        semiconv=0.020 * ureg.radian,
+        scan_rotation=330. * ureg.degree,
+        flip_y=False,
+        cy=(32 - 2) * ureg.pixel,
+        cx=(32 - 2) * ureg.pixel,
+    )
+    as_json = params.model_dump_json()
+    pprint.pprint(("as json", as_json))
+    from_j = from_json(as_json)
+    pprint.pprint(("from json", from_j))
+    res = Simple4DSTEMParams.model_validate(from_j)
+    pprint.pprint(("validated", res))
+    assert isinstance(res.overfocus, Quantity)
+    # To JSON converts to SI base units
+    assert res.scan_pixel_pitch.units == 'meter'
+    assert isinstance(res.flip_y, bool)
+    assert res == params
+
+
 def test_missing():
     with pytest.raises(ValidationError):
         Simple4DSTEMParams(
@@ -141,3 +169,109 @@ def test_json_missing():
     del from_j['overfocus']
     with pytest.raises(ValidationError):
         Simple4DSTEMParams.model_validate(from_j)
+
+
+def test_json_schema_smoke():
+    params = Simple4DSTEMParams(
+        overfocus=1.5 * ureg.millimeter,
+        scan_pixel_pitch=0.000001 * ureg.meter,
+        camera_length=0.15 * ureg.meter,
+        detector_pixel_pitch=0.000050 * ureg.meter,
+        semiconv=0.020 * ureg.radian,
+        scan_rotation=330. * ureg.degree,
+        flip_y=False,
+        # Offset to avoid subchip gap in butted detectors
+        cy=(32 - 2) * ureg.pixel,
+        cx=(32 - 2) * ureg.pixel,
+    )
+    json_schema = params.model_json_schema()
+    pprint.pprint(json_schema)
+    as_json = params.model_dump_json()
+    pprint.pprint(as_json)
+    loaded = json.loads(as_json)
+    jsonschema.validate(
+        instance=loaded,
+        schema=json_schema
+    )
+    # JSON is in SI base units
+    assert tuple(loaded['overfocus']) == (0.0015, 'meter')
+
+
+def test_json_schema_repr():
+    params = Simple4DSTEMParams(
+        overfocus=0.0015 * ureg.meter,
+        scan_pixel_pitch=0.000001 * ureg.meter,
+        camera_length=0.15 * ureg.meter,
+        detector_pixel_pitch=0.000050 * ureg.meter,
+        semiconv=0.020 * ureg.radian,
+        scan_rotation=330. * ureg.degree,
+        flip_y=False,
+        # Offset to avoid subchip gap in butted detectors
+        cy=(32 - 2) * ureg.pixel,
+        cx=(32 - 2) * ureg.pixel,
+    )
+    json_schema = params.model_json_schema()
+    pprint.pprint(json_schema)
+    as_json = params.model_dump_json()
+    pprint.pprint(as_json)
+    loaded = json.loads(as_json)
+    loaded['overfocus'].append('3')
+    with pytest.raises(jsonschema.exceptions.ValidationError):
+        jsonschema.validate(
+            instance=loaded,
+            schema=json_schema
+        )
+
+
+def test_json_schema_missing():
+    params = Simple4DSTEMParams(
+        overfocus=0.0015 * ureg.meter,
+        scan_pixel_pitch=0.000001 * ureg.meter,
+        camera_length=0.15 * ureg.meter,
+        detector_pixel_pitch=0.000050 * ureg.meter,
+        semiconv=0.020 * ureg.radian,
+        scan_rotation=330. * ureg.degree,
+        flip_y=False,
+        # Offset to avoid subchip gap in butted detectors
+        cy=(32 - 2) * ureg.pixel,
+        cx=(32 - 2) * ureg.pixel,
+    )
+    json_schema = params.model_json_schema()
+    pprint.pprint(json_schema)
+    as_json = params.model_dump_json()
+    pprint.pprint(as_json)
+    loaded = json.loads(as_json)
+    del loaded['overfocus']
+    with pytest.raises(jsonschema.exceptions.ValidationError):
+        jsonschema.validate(
+            instance=loaded,
+            schema=json_schema
+        )
+
+
+# No dimensionality check in JSON Schema yet
+@pytest.mark.xfail
+def test_json_schema_dim():
+    params = Simple4DSTEMParams(
+        overfocus=0.0015 * ureg.meter,
+        scan_pixel_pitch=0.000001 * ureg.meter,
+        camera_length=0.15 * ureg.meter,
+        detector_pixel_pitch=0.000050 * ureg.meter,
+        semiconv=0.020 * ureg.radian,
+        scan_rotation=330. * ureg.degree,
+        flip_y=False,
+        # Offset to avoid subchip gap in butted detectors
+        cy=(32 - 2) * ureg.pixel,
+        cx=(32 - 2) * ureg.pixel,
+    )
+    json_schema = params.model_json_schema()
+    pprint.pprint(json_schema)
+    as_json = params.model_dump_json()
+    pprint.pprint(as_json)
+    loaded = json.loads(as_json)
+    loaded['overfocus'][1] = 'degree'
+    with pytest.raises(jsonschema.exceptions.ValidationError):
+        jsonschema.validate(
+            instance=loaded,
+            schema=json_schema
+        )