import sys
from typing import (
TypeVar,
Union,
)
from linkml.generators.common.array import ArrayRangeGenerator, ArrayRepresentation
from linkml_runtime.linkml_model.meta import ArrayExpression, DimensionExpression
if sys.version_info.minor < 12:
from typing_extensions import TypeAliasType
else:
from typing import TypeAliasType
from linkml.generators.pydanticgen.build import RangeResult
from linkml.generators.pydanticgen.template import ConditionalImport, Import, Imports, ObjectImport
_T = TypeVar("_T")
AnyShapeArray = TypeAliasType("AnyShapeArray", list[Union[_T, "AnyShapeArray[_T]"]], type_params=(_T,))
_AnyShapeArrayImports = (
Imports()
+ Import(
module="typing",
objects=[
ObjectImport(name="TypeVar"),
ObjectImport(name="Union"),
],
)
+ ConditionalImport(
condition="sys.version_info.minor >= 12",
module="typing",
objects=[ObjectImport(name="TypeAliasType")],
alternative=Import(module="typing_extensions", objects=[ObjectImport(name="TypeAliasType")]),
)
)
# annotated types are special and inspect.getsource() can't stringify them
_AnyShapeArrayInjects = [
'_T = TypeVar("_T")',
"""AnyShapeArray = TypeAliasType(
"AnyShapeArray", list[Union[_T, "AnyShapeArray[_T]"]], type_params=(_T,)
)""",
]
_ConListImports = Imports() + Import(module="pydantic", objects=[ObjectImport(name="conlist")])
[docs]
class ListOfListsArray(ArrayRangeGenerator):
"""
Represent arrays as lists of lists!
"""
REPR = ArrayRepresentation.LIST
@staticmethod
def _list_of_lists(dimensions: int, dtype: str) -> str:
return ("list[" * dimensions) + dtype + ("]" * dimensions)
@staticmethod
def _parameterized_dimension(dimension: DimensionExpression, dtype: str) -> RangeResult:
# TODO: Preserve label representation in some readable way! doing the MVP now of using conlist
if dimension.exact_cardinality:
dmin = dimension.exact_cardinality
dmax = dimension.exact_cardinality
elif dimension.minimum_cardinality or dimension.maximum_cardinality:
dmin = dimension.minimum_cardinality
dmax = dimension.maximum_cardinality
else:
# TODO: handle labels for labeled but unshaped arrays
return RangeResult(range="list[" + dtype + "]")
items = []
if dmin is not None:
items.append(f"min_length={dmin}")
if dmax is not None:
items.append(f"max_length={dmax}")
items.append(f"item_type={dtype}")
items = ", ".join(items)
range = f"conlist({items})"
return RangeResult(range=range, imports=_ConListImports)
def _any_shape(self, array: ArrayExpression | None = None, with_inner_union: bool = False) -> RangeResult:
"""
An AnyShaped array (using :class:`.AnyShapeArray` )
Args:
array (:class:`.ArrayExpression`): The array expression (not used)
with_inner_union (bool): If ``True`` , the innermost type is a ``Union`` of the ``AnyShapeArray`` class
and ``dtype`` (default: ``False`` )
"""
if self.dtype in ("Any", "AnyType"):
range = "AnyShapeArray"
else:
range = f"AnyShapeArray[{self.dtype}]"
if with_inner_union:
range = f"Union[{range}, {self.dtype}]"
return RangeResult(range=range, injected_classes=_AnyShapeArrayInjects, imports=_AnyShapeArrayImports)
def _bounded_dimensions(self, array: ArrayExpression) -> RangeResult:
"""
A nested series of ``list[]`` ranges with :attr:`.dtype` at the center.
When an array expression allows for a range of dimensions, each set of ``List`` s is joined by a ``Union`` .
"""
if array.exact_number_dimensions or (
array.minimum_number_dimensions
and array.maximum_number_dimensions
and array.minimum_number_dimensions == array.maximum_number_dimensions
):
exact_dims = array.exact_number_dimensions or array.minimum_number_dimensions
return RangeResult(range=self._list_of_lists(exact_dims, self.dtype))
elif not array.maximum_number_dimensions and (
array.minimum_number_dimensions is None or array.minimum_number_dimensions == 1
):
return self._any_shape()
elif array.maximum_number_dimensions:
# e.g., if min = 2, max = 3, range = Union[list[list[dtype]], list[list[list[dtype]]]]
min_dims = array.minimum_number_dimensions if array.minimum_number_dimensions is not None else 1
ranges = [self._list_of_lists(i, self.dtype) for i in range(min_dims, array.maximum_number_dimensions + 1)]
return RangeResult(range="Union[" + ", ".join(ranges) + "]")
else:
# min specified with no max
# e.g., if min = 3, range = list[list[AnyShapeArray[dtype]]]
return RangeResult(
range=self._list_of_lists(array.minimum_number_dimensions - 1, self._any_shape().range),
injected_classes=_AnyShapeArrayInjects,
imports=_AnyShapeArrayImports,
)
def _parameterized_dimensions(self, array: ArrayExpression) -> RangeResult:
"""
Constrained shapes using :func:`pydantic.conlist`
TODO:
- preservation of aliases
- (what other metadata is allowable on labeled dimensions?)
"""
# generate dimensions from inside out and then format
# e.g., if dimensions = [{min_card: 3}, {min_card: 2}],
# range = conlist(min_length=3, item_type=conlist(min_length=2, item_type=dtype))
range = self.dtype
for dimension in reversed(array.dimensions):
range = self._parameterized_dimension(dimension, range).range
return RangeResult(range=range, imports=_ConListImports)
def _complex_dimensions(self, array: ArrayExpression) -> RangeResult:
"""
Mixture of parameterized dimensions with a max or min (or both) shape for anonymous dimensions.
A mixture of ``List`` , :class:`.conlist` , and :class:`.AnyShapeArray` .
"""
res = None
# first process any unlabeled dimensions which must be the innermost level of the range,
# then wrap that with labeled dimensions
if array.exact_number_dimensions or (
array.minimum_number_dimensions
and array.maximum_number_dimensions
and array.minimum_number_dimensions == array.maximum_number_dimensions
):
exact_dims = array.exact_number_dimensions or array.minimum_number_dimensions
if exact_dims > len(array.dimensions):
res = RangeResult(range=self._list_of_lists(exact_dims - len(array.dimensions), self.dtype))
elif exact_dims == len(array.dimensions):
# equivalent to labeled shape
return self._parameterized_dimensions(array)
# else is invalid, see: ArrayValidator.array_consistent_n_dimensions
elif array.maximum_number_dimensions is not None and not array.maximum_number_dimensions:
# unlimited n dimensions, so innermost is AnyShape with dtype
res = self._any_shape(with_inner_union=True)
if array.minimum_number_dimensions:
# some minimum anonymous dimensions but unlimited max dimensions
# e.g., if min = 3, len(dim) = 2, then res.range = list[Union[AnyShapeArray[dtype], dtype]]
# res.range will be wrapped with the 2 labeled dimensions later
res.range = self._list_of_lists(array.minimum_number_dimensions - len(array.dimensions), res.range)
elif array.maximum_number_dimensions:
initial_min = array.minimum_number_dimensions if array.minimum_number_dimensions is not None else 0
dmin = max(len(array.dimensions), initial_min) - len(array.dimensions)
dmax = array.maximum_number_dimensions - len(array.dimensions)
res = self._bounded_dimensions(
ArrayExpression(minimum_number_dimensions=dmin, maximum_number_dimensions=dmax)
)
if res is None:
raise ValueError("Unsupported array specification! this is almost certainly a bug!") # pragma: no cover
# Wrap inner dimension with labeled dimension
# e.g., if dimensions = [{min_card: 3}, {min_card: 2}]
# and res.range = list[Union[AnyShapeArray[dtype], dtype]]
# (min 3 dims, no max dims)
# then the final range = conlist(
# min_length=3,
# item_type=conlist(
# min_length=2,
# item_type=list[Union[AnyShapeArray[dtype], dtype]]
# )
# )
for dim in reversed(array.dimensions):
res = res.merge(self._parameterized_dimension(dim, dtype=res.range))
return res
[docs]
class NumpydanticArray(ArrayRangeGenerator):
"""
Represent array range with :class:`numpydantic.NDArray` annotations,
allowing an abstract array specification to be used with many different array
libraries.
"""
REPR = ArrayRepresentation.NUMPYDANTIC
MIN_NUMPYDANTIC_VERSION = "1.6.1"
"""
Minimum numpydantic version needed to be installed in the environment using
the generated models
"""
IMPORTS = Imports() + Import(
module="numpydantic", objects=[ObjectImport(name="NDArray"), ObjectImport(name="Shape")]
)
INJECTS = [f'MIN_NUMPYDANTIC_VERSION = "{MIN_NUMPYDANTIC_VERSION}"']
[docs]
def make(self) -> RangeResult:
result = super().make()
result.imports = self.IMPORTS.model_copy()
result.injected_classes = self.INJECTS.copy()
return result
[docs]
@staticmethod
def ndarray_annotation(shape: list[int | str] | None = None, dtype: str | None = None) -> str:
"""
Make a stringified :class:`numpydantic.NDArray` annotation for a given shape
and dtype.
If either ``shape`` or ``dtype`` is ``None`` , use ``Any``
"""
if shape is None:
shape = "Any"
else:
shape_expression = ", ".join([str(i) for i in shape])
shape = f'Shape["{shape_expression}"]'
if dtype is None or dtype in ("Any", "AnyType"):
dtype = "Any"
if shape == "Any" and dtype == "Any":
return "NDArray"
else:
return f"NDArray[{shape}, {dtype}]"
@staticmethod
def _dimension_shape(dimension: DimensionExpression) -> str:
if dimension.exact_cardinality:
shape = str(dimension.exact_cardinality)
elif dimension.minimum_cardinality and not dimension.maximum_cardinality:
shape = f"{dimension.minimum_cardinality}-*"
elif dimension.maximum_cardinality and not dimension.minimum_cardinality:
shape = f"*-{dimension.maximum_cardinality}"
elif dimension.minimum_cardinality and dimension.maximum_cardinality:
shape = f"{dimension.minimum_cardinality}-{dimension.maximum_cardinality}"
else:
shape = "*"
return shape
@classmethod
def _parameterized_dimension(cls, dimension: DimensionExpression) -> str:
shape = cls._dimension_shape(dimension)
if dimension.alias is not None:
return f"{shape} {dimension.alias}"
else:
return shape
def _any_shape(self, array: ArrayRepresentation | None = None) -> RangeResult:
"""
Any shaped array, either an unparameterized :class:`numpydantic.NDArray`
if dtype is :class:`typing.Any` , or like ``NDArray[Any, {self.dtype}]``
otherwise.
"""
if self.dtype in ("Any", "AnyType"):
range = "NDArray"
else:
range = f"NDArray[Any, {self.dtype}]"
return RangeResult(range=range)
def _bounded_dimensions(self, array: ArrayExpression) -> RangeResult:
"""
Number of dimensions specified without shape
"""
if array.exact_number_dimensions or (
array.minimum_number_dimensions
and array.maximum_number_dimensions
and array.minimum_number_dimensions == array.maximum_number_dimensions
):
exact_dims = array.exact_number_dimensions or array.minimum_number_dimensions
return RangeResult(range=self.ndarray_annotation(["*"] * exact_dims, self.dtype))
elif not array.maximum_number_dimensions and (
array.minimum_number_dimensions is None or array.minimum_number_dimensions == 1
):
return self._any_shape()
elif array.maximum_number_dimensions:
# e.g., if min = 2, max = 3, range = Union[NDArray[Shape["*, *"], dtype], NDArray[Shape["*, *, *"], dtype]]
min_dims = array.minimum_number_dimensions if array.minimum_number_dimensions is not None else 1
ranges = [
self.ndarray_annotation(["*"] * i, self.dtype)
for i in range(min_dims, array.maximum_number_dimensions + 1)
]
return RangeResult(range="Union[" + ", ".join(ranges) + "]")
else:
# min specified with no max
# e.g., if min = 3, range = NDArray[Shape[*, *, *, ...], dtype]
shape_inner = ["*"] * array.minimum_number_dimensions
shape_inner.append("...")
return RangeResult(range=self.ndarray_annotation(shape_inner, self.dtype))
def _parameterized_dimensions(self, array: ArrayExpression) -> RangeResult:
"""
Arrays with constrained shapes or labels
"""
dims = [self._parameterized_dimension(d) for d in array.dimensions]
range = self.ndarray_annotation(dims, self.dtype)
return RangeResult(range=range)
def _complex_dimensions(self, array: ArrayExpression) -> RangeResult:
"""
Mixture of parameterized dimensions with a max or min (or both) shape for anonymous dimensions.
"""
dims = [self._parameterized_dimension(d) for d in array.dimensions]
res = None
if array.exact_number_dimensions or (
array.minimum_number_dimensions
and array.maximum_number_dimensions
and array.minimum_number_dimensions == array.maximum_number_dimensions
):
exact_dims = array.exact_number_dimensions or array.minimum_number_dimensions
if exact_dims > len(array.dimensions):
dims.extend(["*"] * (exact_dims - len(dims)))
res = self.ndarray_annotation(dims, self.dtype)
elif exact_dims == len(array.dimensions):
# equivalent to labeled shape
return self._parameterized_dimensions(array)
# else is invalid, see: ArrayValidator.array_consistent_n_dimensions(array)
elif array.maximum_number_dimensions is not None and not array.maximum_number_dimensions:
# unlimited n dimensions
if array.minimum_number_dimensions:
# some minimum anonymous dimensions but unlimited max dimensions
dims.extend(["*"] * (array.minimum_number_dimensions - len(dims)))
dims.append("...")
res = self.ndarray_annotation(dims, self.dtype)
elif array.maximum_number_dimensions:
# some res of anonymous dimensions
if array.minimum_number_dimensions:
min_dim = array.minimum_number_dimensions
else:
min_dim = len(dims)
dim_union = []
for i in range(min_dim, array.maximum_number_dimensions + 1):
this_dims = dims.copy()
this_dims.extend(["*"] * (i - len(dims)))
dim_union.append(self.ndarray_annotation(this_dims, self.dtype))
dim_union = ", ".join(dim_union)
res = f"Union[{dim_union}]"
if res is None:
raise ValueError(f"Unhandled range case! {array}")
return RangeResult(range=res)