xradio 0.0.28__py3-none-any.whl → 0.0.30__py3-none-any.whl

xradio/schema/check.py

@@ -5,13 +5,14 @@ import functools
 
 import xarray
 import numpy
+from typeguard import check_type, TypeCheckError
 
 from xradio.schema import (
     metamodel,
+    bases,
     xarray_dataclass_to_array_schema,
     xarray_dataclass_to_dataset_schema,
-    AsDataset,
-    AsDataArray,
+    xarray_dataclass_to_dict_schema,
 )
 
 
@@ -110,6 +111,9 @@ class SchemaIssues(Exception):
             issues_string = "\n * ".join(repr(issue) for issue in self.issues)
             return f"\n * {issues_string}"
 
+    def __repr__(self):
+        return f"SchemaIssues({str(self)})"
+
     def expect(
         self, elem: typing.Optional[str] = None, ix: typing.Optional[str] = None
     ):
@@ -129,12 +133,14 @@ class SchemaIssues(Exception):
             raise self
 
 
-def check_array(array: xarray.DataArray, schema: metamodel.ArraySchema) -> SchemaIssues:
+def check_array(
+    array: xarray.DataArray, schema: typing.Union[type, metamodel.ArraySchema]
+) -> SchemaIssues:
     """
     Check whether an xarray DataArray conforms to a schema
 
     :param array: DataArray to check
-    :param schema: Schema to check against
+    :param schema: Schema to check against (possibly as :py:class:`AsDataset`)
     :returns: List of :py:class:`SchemaIssue`s found
     """
 
@@ -143,7 +149,7 @@ def check_array(array: xarray.DataArray, schema: metamodel.ArraySchema) -> Schem
         raise TypeError(
             f"check_array: Expected xarray.DataArray, but got {type(array)}!"
         )
-    if type(schema) == type and issubclass(schema, AsDataArray):
+    if bases.is_dataarray_schema(schema):
         schema = xarray_dataclass_to_array_schema(schema)
     if not isinstance(schema, metamodel.ArraySchema):
         raise TypeError(f"check_array: Expected ArraySchema, but got {type(schema)}!")
@@ -164,13 +170,13 @@ def check_array(array: xarray.DataArray, schema: metamodel.ArraySchema) -> Schem
 
 
 def check_dataset(
-    dataset: xarray.Dataset, schema: metamodel.DatasetSchema
+    dataset: xarray.Dataset, schema: typing.Union[type, metamodel.DatasetSchema]
 ) -> SchemaIssues:
     """
     Check whether an xarray DataArray conforms to a schema
 
     :param array: DataArray to check
-    :param schema: Schema to check against (possibly as dataclass)
+    :param schema: Schema to check against (possibly as :py:class:`AsDataArray` dataclass)
     :returns: List of :py:class:`SchemaIssue`s found
     """
 
@@ -179,7 +185,7 @@ def check_dataset(
         raise TypeError(
             f"check_dataset: Expected xarray.Dataset, but got {type(dataset)}!"
         )
-    if isinstance(schema, AsDataset):
+    if bases.is_dataset_schema(schema):
         schema = xarray_dataclass_to_dataset_schema(schema)
     if not isinstance(schema, metamodel.DatasetSchema):
         raise TypeError(
@@ -247,7 +253,7 @@ def check_dimensions(
     hint_remove = [f"'{hint}'" for hint in dims_set - best]
     hint_add = [f"'{hint}'" for hint in best - dims_set]
     if hint_remove and hint_add:
-        message = f"Unexpected coordinates, replace {','.join(hint_add)} by {','.join(hint_remove)}?"
+        message = f"Unexpected coordinates, replace {','.join(hint_remove)} by {','.join(hint_add)}?"
     elif hint_remove:
         message = f"Superflous coordinate {','.join(hint_remove)}?"
     elif hint_add:
@@ -259,8 +265,8 @@ def check_dimensions(
             SchemaIssue(
                 path=[("dims", None)],
                 message=message,
-                found=dims,
-                expected=expected,
+                found=list(dims),
+                expected=list(expected),
             )
         ]
     )
@@ -276,7 +282,13 @@ def check_dtype(dtype: numpy.dtype, expected: [numpy.dtype]) -> SchemaIssues:
     """
 
     for exp_dtype in expected:
-        if dtype == exp_dtype:
+        # If the expected dtype has no size (e.g. "U", a.k.a. a string of
+        # arbitrary length), we don't check itemsize, only kind.
+        if (
+            dtype.kind == exp_dtype.kind
+            and exp_dtype.itemsize == 0
+            or exp_dtype == dtype
+        ):
             return SchemaIssues()
 
     # Not sure there's anything more helpful that we can do here? Any special
@@ -287,7 +299,7 @@ def check_dtype(dtype: numpy.dtype, expected: [numpy.dtype]) -> SchemaIssues:
                 path=[("dtype", None)],
                 message="Wrong numpy dtype",
                 found=dtype,
-                expected=expected,
+                expected=list(expected),
             )
         ]
     )
@@ -296,6 +308,7 @@ def check_dtype(dtype: numpy.dtype, expected: [numpy.dtype]) -> SchemaIssues:
 def check_attributes(
     attrs: typing.Dict[str, typing.Any],
     attrs_schema: typing.List[metamodel.AttrSchemaRef],
+    attr_kind: str = "attrs",
 ) -> SchemaIssues:
     """
     Check whether an attribute set conforms to a schema
@@ -307,23 +320,29 @@ def check_attributes(
 
     issues = SchemaIssues()
     for attr_schema in attrs_schema:
-
         # Attribute missing? Note that a value of "None" is equivalent for the
         # purpose of the check
         val = attrs.get(attr_schema.name)
         if val is None:
             if not attr_schema.optional:
+                # Get options
+                if typing.get_origin(attr_schema.typ) is typing.Union:
+                    options = typing.get_args(attr_schema.typ)
+                else:
+                    options = [attr_schema.typ]
+
                 issues.add(
                     SchemaIssue(
-                        path=[("attrs", attr_schema.name)],
+                        path=[(attr_kind, attr_schema.name)],
                         message=f"Required attribute {attr_schema.name} is missing!",
+                        expected=options,
                     )
                 )
             continue
 
         # Check attribute value
         issues += _check_value_union(val, attr_schema.typ).at_path(
-            "attrs", attr_schema.name
+            attr_kind, attr_schema.name
         )
 
     # Extra attributes are always okay
@@ -352,7 +371,6 @@ def check_data_vars(
 
     issues = SchemaIssues()
     for data_var_schema in data_vars_schema:
-
         # Data_Varinate missing?
         data_var = data_vars.get(data_var_schema.name)
         if data_var is None:
@@ -382,6 +400,21 @@ def check_data_vars(
     return issues
 
 
+def check_dict(
+    dct: dict, schema: typing.Union[type, metamodel.DictSchema]
+) -> SchemaIssues:
+    # Check that this is actually a dictionary
+    if not isinstance(dct, dict):
+        raise TypeError(f"check_dict: Expected dictionary, but got {type(dct)}!")
+    if bases.is_dict_schema(schema):
+        schema = xarray_dataclass_to_dict_schema(schema)
+    if not isinstance(schema, metamodel.DictSchema):
+        raise TypeError(f"check_dict: Expected DictSchema, but got {type(schema)}!")
+
+    # Check attributes
+    return check_attributes(dct, schema.attributes, attr_kind="")
+
+
 def _check_value(val, ann):
     """
     Check whether value satisfies annotation
@@ -394,46 +427,60 @@ def _check_value(val, ann):
     """
 
     # Is supposed to be a data array?
-    if type(ann) == type and issubclass(ann, AsDataArray):
+    if bases.is_dataarray_schema(ann):
+        # Attempt to convert dictionaries automatically
+        if isinstance(val, dict):
+            try:
+                val = xarray.DataArray.from_dict(val)
+            except ValueError as e:
+                return SchemaIssues(
+                    [
+                        SchemaIssue(
+                            path=[], message=str(e), expected=[ann], found=type(val)
+                        )
+                    ]
+                )
+
         if not isinstance(val, xarray.DataArray):
-            return SchemaIssues(
-                [
-                    SchemaIssue(
-                        path=[],
-                        message="Unexpected type",
-                        expected=[xarray.DataArray],
-                        found=type(val),
-                    )
-                ]
-            )
+            # Fall through to plain type check
+            ann = xarray.DataArray
         else:
             return check_array(val, ann)
 
     # Is supposed to be a dataset?
-    if type(ann) == type and issubclass(ann, AsDataset):
+    if bases.is_dataset_schema(ann):
+        # Attempt to convert dictionaries automatically
+        if isinstance(val, dict):
+            try:
+                val = xarray.Dataset.from_dict(val)
+            except ValueError as e:
+                return SchemaIssues(
+                    [
+                        SchemaIssue(
+                            path=[], message=str(t), expected=[ann], found=type(val)
+                        )
+                    ]
+                )
         if not isinstance(val, xarray.Dataset):
-            return SchemaIssues(
-                [
-                    SchemaIssue(
-                        path=[],
-                        message="Unexpected type",
-                        expected=[xarray.DataArray],
-                        found=type(val),
-                    )
-                ]
-            )
+            # Fall through to plain type check
+            ann = xarray.Dataset
         else:
             return check_dataset(val, ann)
 
-    # Otherwise straight type check (TODO - be more fancy, possibly by
-    # importing from Typeguard module? Don't want to overdo it...)
-    if not isinstance(val, ann):
+    # Is supposed to be a dictionary?
+    if bases.is_dict_schema(ann):
+        if not isinstance(val, dict):
+            # Fall through to plain type check
+            ann = dict
+        else:
+            return check_dict(val, ann)
+
+    # Otherwise straight type check using typeguard
+    try:
+        check_type(val, ann)
+    except TypeCheckError as t:
         return SchemaIssues(
-            [
-                SchemaIssue(
-                    path=[], message="Unexpected type", expected=[ann], found=type(val)
-                )
-            ]
+            [SchemaIssue(path=[], message=str(t), expected=[ann], found=type(val))]
         )
 
     return SchemaIssues()
@@ -471,13 +518,8 @@ def _check_value_union(val, ann):
         if args_issues is None:
             args_issues = arg_issues
 
-        # Fancy merging of expected options (for "unexpected type")
-        elif (
-            len(args_issues) == 1
-            and len(arg_issues) == 1
-            and args_issues[0].message == arg_issues[0].message
-        ):
-
+        # Crude merging of expected options (for "unexpected type")
+        elif len(args_issues) == 1 and len(arg_issues) == 1:
             args_issues[0].expected += arg_issues[0].expected
 
     # Return representative issues list
@@ -511,7 +553,6 @@ def schema_checked(fn, check_parameters: bool = True, check_return: bool = True)
 
     @functools.wraps(fn)
     def _check_fn(*args, **kwargs):
-
         # Hide this function in pytest tracebacks
         __tracebackhide__ = True
 

xradio/schema/dataclass.py

@@ -1,10 +1,13 @@
 from typing import get_type_hints, get_args
-from .typing import get_dims, get_dtypes, get_role, Role, get_annotated, is_optional
+from .typing import get_dims, get_types, get_role, Role, get_annotated, is_optional
 
 import typing
 import inspect
 import ast
 import dataclasses
+import numpy
+import itertools
+import textwrap
 
 from xradio.schema.metamodel import *
 
@@ -18,8 +21,11 @@ def extract_field_docstrings(klass):
     """
 
     # Parse class body
-    src = inspect.getsource(klass)
-    module = ast.parse(src)
+    try:
+        src = inspect.getsource(klass)
+    except OSError:
+        return {}
+    module = ast.parse(textwrap.dedent(src))
 
     # Expect module containing a class definition
     if not isinstance(module, ast.Module) or len(module.body) != 1:
@@ -31,7 +37,6 @@ def extract_field_docstrings(klass):
     # Go through body, collect dostrings
     docstrings = {}
     for i, assign in enumerate(cls.body):
-
         # Handle both annotated and unannotated case
         if isinstance(assign, ast.AnnAssign):
             if not isinstance(assign.target, ast.Name):
@@ -78,7 +83,6 @@ def extract_xarray_dataclass(klass):
     data_vars = []
     attributes = []
     for field in dataclasses.fields(klass):
-
         # Get field "role" (coordinate / data variable / attribute) from its
         # type hint
         typ = type_hints[field.name]
@@ -111,7 +115,6 @@ def extract_xarray_dataclass(klass):
         # Defined using a dataclass, i.e. Coordof/Dataof?
         dataclass = typing.get_args(get_annotated(typ))[0]
         if dataclasses.is_dataclass(dataclass):
-
             # Recursively get array schema for data class
             arr_schema = xarray_dataclass_to_array_schema(dataclass)
             arr_schema_fields = {
@@ -129,24 +132,51 @@ def extract_xarray_dataclass(klass):
             )
 
         else:
-
-            # Assume that it's an "inline" declaration using "Coord"/"Data"
-            schema_ref = ArraySchemaRef(
-                name=field.name,
-                optional=is_optional(typ),
-                default=field.default,
-                docstring=field_docstrings.get(field.name),
-                schema_name=f"{klass.__module__}.{klass.__qualname__}.{field.name}",
-                dimensions=get_dims(typ),
-                dtypes=get_dtypes(typ),
-                coordinates=[],
-                attributes=[],
-                class_docstring=None,
-                data_docstring=None,
-            )
+            # Get dimensions and dtypes
+            dims = get_dims(typ)
+            types = get_types(typ)
+
+            # Is types a (single) dataclass?
+            if len(types) == 1 and dataclasses.is_dataclass(types[0]):
+                # Recursively get array schema for data class
+                arr_schema = xarray_dataclass_to_array_schema(types[0])
+
+                # Prepend dimensions to array schema
+                combined_dimensions = [
+                    dims1 + dims2
+                    for dims1, dims2 in itertools.product(dims, arr_schema.dimensions)
+                ]
+
+                # Repackage as reference
+                arr_schema_fields = {
+                    f.name: getattr(arr_schema, f.name)
+                    for f in dataclasses.fields(ArraySchema)
+                }
+                arr_schema_fields["dimensions"] = combined_dimensions
+                schema_ref = ArraySchemaRef(
+                    name=field.name,
+                    optional=is_optional(typ),
+                    default=field.default,
+                    docstring=field_docstrings.get(field.name),
+                    **arr_schema_fields,
+                )
+            else:
+                # Assume that it's an "inline" declaration using "Coord"/"Data"
+                schema_ref = ArraySchemaRef(
+                    name=field.name,
+                    optional=is_optional(typ),
+                    default=field.default,
+                    docstring=field_docstrings.get(field.name),
+                    schema_name=f"{klass.__module__}.{klass.__qualname__}.{field.name}",
+                    dimensions=dims,
+                    dtypes=[numpy.dtype(typ) for typ in types],
+                    coordinates=[],
+                    attributes=[],
+                    class_docstring=None,
+                    data_docstring=None,
+                )
 
         if is_coord:
-
             # Make sure that it is valid to use as a coordinate - i.e. we don't
             # have "recursive" (?!) coordinate definitions
             if not schema_ref.is_coord():
@@ -171,23 +201,34 @@ def xarray_dataclass_to_array_schema(klass):
     refer to using CoordOf or DataOf
     """
 
+    # Cached?
+    if hasattr(klass, "__xradio_array_schema"):
+        return klass.__xradio_array_schema
+
     # Extract from data class
     coordinates, data_vars, attributes = extract_xarray_dataclass(klass)
 
     # For a dataclass there must be exactly one data variable
-    # (typically called "data", but we don't check that)
     if not data_vars:
         raise ValueError(
-            f"Found no data declaration in (supposed) data darray class {klass.__name__}!"
+            f"Found no data declaration in (supposed) data array class {klass.__name__}!"
         )
     if len(data_vars) > 1:
         raise ValueError(
             f"Found multiple data variables ({', '.join(v.name for v in data_vars)})"
-            f" in supposed data darray class {klass.__name__}!"
+            f" in supposed data array class {klass.__name__}!"
+        )
+
+    # Check that data variable is named "data". This is important for this to
+    # match up with parameters to xarray.DataArray() later (see bases.AsArray)
+    if data_vars[0].name != "data":
+        raise ValueError(
+            f"Data variable in data array class {klass.__name__} "
+            f'must be called "data", not {data_vars[0].name}!'
         )
 
     # Make class
-    return ArraySchema(
+    schema = ArraySchema(
         schema_name=f"{klass.__module__}.{klass.__qualname__}",
         dimensions=data_vars[0].dimensions,
         dtypes=data_vars[0].dtypes,
@@ -196,6 +237,8 @@ def xarray_dataclass_to_array_schema(klass):
         class_docstring=inspect.cleandoc(klass.__doc__),
         data_docstring=data_vars[0].docstring,
     )
+    klass.__xradio_array_schema = schema
+    return schema
 
 
 def xarray_dataclass_to_dataset_schema(klass):
@@ -206,6 +249,10 @@ def xarray_dataclass_to_dataset_schema(klass):
     refer to using CoordOf or DataOf
     """
 
+    # Cached?
+    if hasattr(klass, "__xradio_dataset_schema"):
+        return klass.__xradio_dataset_schema
+
     # Extract from data class
     coordinates, data_vars, attributes = extract_xarray_dataclass(klass)
 
@@ -242,7 +289,7 @@ def xarray_dataclass_to_dataset_schema(klass):
     dimensions = [[dim for dim in all_dimensions if dim in dims] for dims in dimensions]
 
     # Make class
-    return DatasetSchema(
+    schema = DatasetSchema(
         schema_name=f"{klass.__module__}.{klass.__qualname__}",
         dimensions=dimensions,
         coordinates=coordinates,
@@ -250,3 +297,52 @@ def xarray_dataclass_to_dataset_schema(klass):
         attributes=attributes,
         class_docstring=inspect.cleandoc(klass.__doc__),
     )
+    klass.__xradio_dataset_schema = schema
+    return schema
+
+
+def xarray_dataclass_to_dict_schema(klass):
+    """
+    Convert an xarray-dataclass style schema dataclass to an DictSchema
+
+    This should work on any class annotated with :py:func:`~xradio.schema.bases.dict_schema`
+    """
+
+    # Cached?
+    if hasattr(klass, "__xradio_dict_schema"):
+        return klass.__xradio_dict_schema
+
+    # Get docstrings and type hints
+    field_docstrings = extract_field_docstrings(klass)
+    type_hints = get_type_hints(klass, include_extras=True)
+    attributes = []
+    for field in dataclasses.fields(klass):
+        typ = type_hints[field.name]
+
+        # Handle optional value: Strip "None" from the types
+        optional = is_optional(typ)
+        if optional:
+            typs = [typ for typ in get_args(typ) if typ is not None.__class__]
+            if len(typs) == 1:
+                typ = typs[0]
+            else:
+                typ = typing.Union.__getitem__[tuple(typs)]
+
+        attributes.append(
+            AttrSchemaRef(
+                name=field.name,
+                typ=typ,
+                optional=optional,
+                default=field.default,
+                docstring=field_docstrings.get(field.name),
+            )
+        )
+
+    # Return
+    schema = DictSchema(
+        schema_name=f"{klass.__module__}.{klass.__qualname__}",
+        attributes=attributes,
+        class_docstring=inspect.cleandoc(klass.__doc__),
+    )
+    klass.__xradio_dict_schema = schema
+    return schema

xradio/schema/metamodel.py

@@ -6,10 +6,11 @@ __all__ = [
     "ArraySchema",
     "ArraySchemaRef",
     "DatasetSchema",
+    "DictSchema",
 ]
 
 
-@dataclass
+@dataclass(frozen=True)
 class AttrSchemaRef:
     """
     Schema information about an attribute as referenced from an array or
@@ -34,7 +35,7 @@ class AttrSchemaRef:
     """Documentation string of attribute reference"""
 
 
-@dataclass
+@dataclass(frozen=True)
 class ArraySchema:
     """
     Schema for xarray data array
@@ -83,7 +84,7 @@ class ArraySchema:
         return req_dims
 
 
-@dataclass
+@dataclass(frozen=True)
 class ArraySchemaRef(ArraySchema):
     """
     Schema for xarray data array as referenced from a dataset schema
@@ -102,7 +103,7 @@ class ArraySchemaRef(ArraySchema):
     """Documentation string of array reference"""
 
 
-@dataclass
+@dataclass(frozen=True)
 class DatasetSchema:
     """
     Schema for an xarray dataset
@@ -122,3 +123,19 @@ class DatasetSchema:
 
     class_docstring: typing.Optional[str]
     """Documentation string of class"""
+
+
+@dataclass(frozen=True)
+class DictSchema:
+    """
+    Schema for a simple dictionary
+    """
+
+    schema_name: str
+    """(Class) name of the schema"""
+
+    attributes: [AttrSchemaRef]
+    """List of attributes"""
+
+    class_docstring: typing.Optional[str]
+    """Documentation string of class"""

xradio/schema/typing.py

@@ -29,16 +29,24 @@ from typing import (
     Protocol,
 )
 
+from typing import Union
+
 try:
+    # Python 3.10 forward: TypeAlias, ParamSpec are standard, and there is the
+    # "a | b" UnionType alternative to "Union[a,b]"
     from typing import TypeAlias, ParamSpec
     from types import UnionType
+
+    HAVE_UNIONTYPE = True
 except ImportError:
-    # Required for Python 3.9
+    # Python 3.9: Get TypeAlias, ParamSpec from typing_extensions, no support
+    # for "a | b"
     from typing_extensions import (
         TypeAlias,
         ParamSpec,
     )
-    from typing import Union as UnionType
+
+    HAVE_UNIONTYPE = False
 import numpy as np
 from itertools import chain
 from enum import Enum
@@ -295,7 +303,9 @@ def get_dims(tp: Any) -> List[Dims]:
         raise TypeError(f"Could not find any dims in {tp!r}.")
 
     # List of allowed dtypes (might just be one)
-    if get_origin(dims) is UnionType:
+    if get_origin(dims) is Union:
+        dims_in = get_args(dims)
+    elif HAVE_UNIONTYPE and get_origin(dims) is UnionType:
         dims_in = get_args(dims)
     else:
         dims_in = [dims]
@@ -326,35 +336,40 @@ def get_dims(tp: Any) -> List[Dims]:
     return dims_out
 
 
-def get_dtypes(tp: Any) -> List[AnyDType]:
-    """Extract a NumPy data types (dtype)."""
+def get_types(tp: Any) -> List[AnyDType]:
+    """Extract data types from type annotation
+
+    E.g. Coord[..., Type1 | Type2 | ...] or Data[..., Type1 | Type2 | ...]
+
+    """
     try:
-        dtype = get_args(get_args(get_annotated(tp))[1])[0]
+        typ = get_args(get_args(get_annotated(tp))[1])[0]
     except TypeError:
         raise TypeError(f"Could not find any dtype in {tp!r}.")
 
     # List of allowed dtypes (might just be one)
-    if get_origin(dtype) is UnionType:
-        dtypes_in = get_args(dtype)
+    if get_origin(typ) is Union:
+        types_in = get_args(typ)
+    elif HAVE_UNIONTYPE and get_origin(typ) is UnionType:
+        types_in = get_args(typ)
     else:
-        dtypes_in = [dtype]
-
-    dtypes_out = []
-    for dt in dtypes_in:
+        types_in = [typ]
 
+    types_out = []
+    for dt in types_in:
         # Handle case that we want to allow "Any"
         if dt is Any or dt is type(None):
-            dtypes_out.append(None)
+            types_out.append(None)
             continue
 
-        # Allow specifying dtype as literal (e.g. string)
+        # Allow specifying type as literal (e.g. string)
         elif get_origin(dt) is Literal:
-            dtype = get_args(dt)[0]
+            dt = get_args(dt)[0]
 
-        # Construct numpy dtype
-        dtypes_out.append(np.dtype(dt))
+        # Return type
+        types_out.append(dt)
 
-    return dtypes_out
+    return types_out
 
 
 def get_name(tp: Any, default: Hashable = None) -> Hashable: