diff --git a/python-spec/src/somacore/options.py b/python-spec/src/somacore/options.py index b1661de2..5ff7efdc 100644 --- a/python-spec/src/somacore/options.py +++ b/python-spec/src/somacore/options.py @@ -5,7 +5,7 @@ """ import enum -from typing import Any, Mapping, Optional, Sequence, Union +from typing import Any, Mapping, Optional, Sequence, TypeVar, Union import attrs import numpy as np @@ -110,36 +110,42 @@ class ResultOrder(enum.Enum): ResultOrderStr = Union[ResultOrder, Literal["auto", "row-major", "column-major"]] """A ResultOrder, or the str representing it.""" - -DenseCoord = Union[None, int, types.Slice[int]] +# Coordinate types are specified as TypeVars instead of Unions to ensure that +# sequences and slices are homogeneous: +# +# BadCoord = Union[int, str] +# BadCoords = Union[BadCoord, Sequence[BadCoord]] +# # ["this", 1] is bad, but is a valid as a BadCoords value +# +# GoodCoord = TypeVar("GoodCoord", int, str) +# GoodCoords = Union[GoodCoord, Sequence[GoodCoord]] +# # ["this", 1] is bad, and is *not* valid as a GoodCoords value +# # ["this", "that"] is a valid as a GoodCoords value +# # [1, 2] is valid as a GoodCoords value + +DenseCoord = TypeVar("DenseCoord", None, int, types.Slice[int]) """A single coordinate range for reading dense data. ``None`` indicates the entire domain of a dimension; values of this type are not ``Optional``, but may be ``None``. """ + + DenseNDCoords = Sequence[DenseCoord] """A sequence of ranges to read dense data.""" -# TODO: Add support for non-integer types. +SparseCoord = TypeVar( + "SparseCoord", bytes, float, int, slice, str, np.datetime64, pa.TimestampType +) +"""Types that can be put into a single sparse coordinate.""" + # NOTE: Keep this in sync with the types accepted in `_canonicalize_coord` # in ./query/axis.py. -# https://github.com/single-cell-data/TileDB-SOMA/issues/960 SparseDFCoord = Union[ DenseCoord, - float, - np.datetime64, - pa.TimestampType, - Sequence[int], - Sequence[float], - Sequence[str], - Sequence[bytes], - Sequence[np.datetime64], - Sequence[pa.TimestampType], - types.Slice[float], - types.Slice[str], - types.Slice[bytes], - types.Slice[np.datetime64], - types.Slice[pa.TimestampType], + SparseCoord, + Sequence[SparseCoord], + types.Slice[SparseCoord], pa.Array, pa.ChunkedArray, npt.NDArray[np.integer], diff --git a/python-spec/src/somacore/query/axis.py b/python-spec/src/somacore/query/axis.py index 1698a8d8..e52fb1c5 100644 --- a/python-spec/src/somacore/query/axis.py +++ b/python-spec/src/somacore/query/axis.py @@ -34,6 +34,7 @@ def _canonicalize_coord(coord: options.SparseDFCoord) -> options.SparseDFCoord: if coord is None or isinstance( coord, ( + bool, bytes, float, int, diff --git a/python-spec/testing/test_query_axis.py b/python-spec/testing/test_query_axis.py index f5b22a89..a235698d 100644 --- a/python-spec/testing/test_query_axis.py +++ b/python-spec/testing/test_query_axis.py @@ -22,6 +22,7 @@ (slice(np.datetime64(946684802, "s"), np.datetime64(946684803, "s")),), ), (("string-coord", [b"lo", b"hi"]), ("string-coord", (b"lo", b"hi"))), + ((slice(4, 5), True, None), (slice(4, 5), True, None)), ], ) def test_canonicalization(coords: Any, want: Tuple[options.SparseDFCoord, ...]) -> None: