Skip to content

Commit

Permalink
(feat): support ellipsis indexing
Browse files Browse the repository at this point in the history
  • Loading branch information
ilan-gold committed Oct 22, 2024
1 parent 6dfe780 commit d2b859d
Show file tree
Hide file tree
Showing 3 changed files with 16 additions and 4 deletions.
9 changes: 6 additions & 3 deletions src/anndata/_core/index.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
from collections.abc import Iterable, Sequence
from functools import singledispatch
from itertools import repeat
from types import EllipsisType
from typing import TYPE_CHECKING

import h5py
Expand Down Expand Up @@ -47,7 +48,8 @@ def _normalize_index(
| str
| Sequence[bool | int | np.integer]
| np.ndarray
| pd.Index,
| pd.Index
| EllipsisType,
index: pd.Index,
) -> slice | int | np.ndarray: # ndarray of int or bool
if not isinstance(index, pd.RangeIndex):
Expand All @@ -72,6 +74,8 @@ def name_idx(i):
return slice(start, stop, step)
elif isinstance(indexer, np.integer | int):
return indexer
elif isinstance(indexer, EllipsisType):
return slice(None)
elif isinstance(indexer, str):
return index.get_loc(indexer) # int
elif isinstance(
Expand Down Expand Up @@ -107,8 +111,7 @@ def name_idx(i):
"are not valid obs/ var names or indices."
)
return positions # np.ndarray[int]
else:
raise IndexError(f"Unknown indexer {indexer!r} of type {type(indexer)}")
raise IndexError(f"Unknown indexer {indexer!r} of type {type(indexer)}")


def _fix_slice_bounds(s: slice, length: int) -> slice:
Expand Down
3 changes: 2 additions & 1 deletion src/anndata/compat/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@
from importlib.util import find_spec
from inspect import Parameter, signature
from pathlib import Path
from types import EllipsisType
from typing import TYPE_CHECKING, TypeVar
from warnings import warn

Expand Down Expand Up @@ -46,7 +47,7 @@ class Empty:
pass


Index1D = slice | int | str | np.int64 | np.ndarray
Index1D = slice | int | str | np.int64 | np.ndarray | EllipsisType
Index = Index1D | tuple[Index1D, Index1D] | scipy.sparse.spmatrix | SpArray
H5Group = h5py.Group
H5Array = h5py.Dataset
Expand Down
8 changes: 8 additions & 0 deletions tests/test_views.py
Original file line number Diff line number Diff line change
Expand Up @@ -786,6 +786,14 @@ def test_dataframe_view_index_setting():
assert a2.obs.index.values.tolist() == ["a", "b"]


def test_ellipsis_index(adata, subset_func, matrix_type):
adata = gen_adata((10, 10), X_type=matrix_type, **GEN_ADATA_DASK_ARGS)
subset_obs_names = subset_func(adata.obs_names)
subset_ellipsis = adata[subset_obs_names, ...]
subset = adata[subset_obs_names, :]
assert_equal(subset_ellipsis, subset)


# @pytest.mark.parametrize("dim", ["obs", "var"])
# @pytest.mark.parametrize(
# ("idx", "pat"),
Expand Down

0 comments on commit d2b859d

Please sign in to comment.