"""
Functions to import and export :py:class:`xradio.schema.metamodel.DatasetSchema`
as JSON representation. This can be used to externalise schema checks, or
generate documentation from schemas in JSON representation.
"""
import dataclasses
import json
from xradio.schema import (
bases,
metamodel,
xarray_dataclass_to_array_schema,
xarray_dataclass_to_dataset_schema,
xarray_dataclass_to_dict_schema,
)
__all__ = ["export_schema_json_file", "import_schema_json_file"]
CLASS_ATTR = "$class"
class DataclassEncoder(json.JSONEncoder):
"""
General-purpose encoder that represents data classes as
dictionaries, omitting defaults and annotating the original class
as a ``'$class'`` attribute.
"""
def default(self, o):
if dataclasses.is_dataclass(o):
res = {CLASS_ATTR: o.__class__.__name__}
for fld in dataclasses.fields(type(o)):
if (
getattr(o, fld.name) is not fld.default
and getattr(o, fld.name) is not dataclasses.MISSING
):
res[fld.name] = getattr(o, fld.name)
return res
return super().default(o)
DATACLASS_MAP = {
cls.__name__: cls
for cls in [
metamodel.DictSchema,
metamodel.ValueSchema,
metamodel.AttrSchemaRef,
metamodel.ArraySchema,
metamodel.ArraySchemaRef,
metamodel.DatasetSchema,
]
}
class DataclassDecoder(json.JSONDecoder):
"""
General-purpose decoder that reads JSON as generated by
:py:class:`DataclassEncoder`.
"""
def __init__(self, dataclass_map, *args, **kwargs):
self._dataclass_map = dataclass_map
super().__init__(*args, object_hook=self.object_hook, **kwargs)
def object_hook(self, obj):
# Detect dictionaries with '$class' annotation
if isinstance(obj, dict) and CLASS_ATTR in obj:
# Identify the class
cls_name = obj[CLASS_ATTR]
cls = self._dataclass_map.get(cls_name)
if not cls:
raise ValueError(
f"Unknown $dataclass encountered while decoding JSON: {cls_name}"
)
# Instantiate
del obj[CLASS_ATTR]
obj = cls(**obj)
return obj
[docs]
def export_schema_json_file(schema: "DatasetSchema", fname: str):
"""
Exports given schema as a JSON file
:param schema: Dataset schema. Dataclasses will be converted automatically.
:param fname: File name to write serialised schema to
"""
# Check that this is actually a Dataset
if bases.is_dataset_schema(schema):
schema = xarray_dataclass_to_dataset_schema(schema)
if not isinstance(schema, metamodel.DatasetSchema):
raise TypeError(
f"export_schema_json_file: Expected DatasetSchema, but got {type(schema)}!"
)
# Perform export
with open(fname, "w", encoding="utf8") as f:
json.dump(schema, f, cls=DataclassEncoder, ensure_ascii=False, indent=" ")
[docs]
def import_schema_json_file(fname: str):
"""
Imports a schema from a JSON file
For JSON files generated by
:py:func:`export_schema_json_file`, this will return a
:py:class:`~xradio.schema.metamodel.DatasetSchema`.
:param fname: File name to load
:returns: Deserialised object
"""
with open(fname, "r", encoding="utf8") as f:
return json.load(f, cls=DataclassDecoder, dataclass_map=DATACLASS_MAP)