[python] Add support for timestamp types in axis query (#139)

python-spec/src/somacore/options.py

@@ -130,12 +130,15 @@ class ResultOrder(enum.Enum):
     Sequence[float],
     Sequence[str],
     Sequence[bytes],
+    Sequence[npt.NDArray[np.datetime64]],
     types.Slice[float],
     types.Slice[str],
     types.Slice[bytes],
+    types.Slice[npt.NDArray[np.datetime64]],
     pa.Array,
     pa.ChunkedArray,
     npt.NDArray[np.integer],
+    npt.NDArray[np.datetime64],
 ]
 """A single coordinate range for one dimension of a sparse dataframe."""
 SparseDFCoords = Sequence[SparseDFCoord]

python-spec/src/somacore/query/axis.py

@@ -33,7 +33,18 @@ def _canonicalize_coord(coord: options.SparseDFCoord) -> options.SparseDFCoord:
     # NOTE: Keep this in sync with the `SparseDFCoord` type.
     if coord is None or isinstance(
         coord,
-        (bytes, float, int, slice, str, pa.Array, pa.ChunkedArray, np.ndarray),
+        (
+            bytes,
+            float,
+            int,
+            slice,
+            str,
+            pa.Array,
+            pa.ChunkedArray,
+            np.ndarray,
+            np.datetime64,
+            pa.lib.TimestampType,
+        ),
     ):
         return coord
     if isinstance(coord, Sequence):

python-spec/testing/test_query_axis.py

@@ -17,6 +17,10 @@
         ([1, 1.5, 2], (1, 1.5, 2)),
         ((slice(None), [0, 88, 1001]), (slice(None), (0, 88, 1001))),
         ((slice(2.5, 3.5),), (slice(2.5, 3.5),)),
+        (
+            (slice(np.datetime64(946684802, "s"), np.datetime64(946684803, "s")),),
+            (slice(np.datetime64(946684802, "s"), np.datetime64(946684803, "s")),),
+        ),
         (("string-coord", [b"lo", b"hi"]), ("string-coord", (b"lo", b"hi"))),
     ],
 )