1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
|
from __future__ import annotations
from collections.abc import Generator
from functools import partial
from importlib.metadata import version
from typing import TYPE_CHECKING
import joblib
import pytest
from dask.base import normalize_token, tokenize
from packaging.version import Version
from anndata.compat import is_zarr_v2
if Version(version("dask")) < Version("2024.8.0"):
from dask.base import normalize_seq
else:
from dask.tokenize import normalize_seq
from filelock import FileLock
from scipy import sparse
import anndata as ad
from anndata.tests.helpers import subset_func # noqa: F401
if TYPE_CHECKING:
from collections.abc import Generator
from pathlib import Path
from types import EllipsisType
from typing import Literal
@pytest.fixture
def backing_h5ad(tmp_path: Path) -> Path:
return tmp_path / "test.h5ad"
@pytest.fixture(
params=[
("h5ad", None),
("zarr", 2),
pytest.param(
("zarr", 3),
marks=pytest.mark.skipif(
is_zarr_v2(), reason="zarr v3 file format not supported with v2 package"
),
),
],
ids=["h5ad", "zarr2", "zarr3"],
)
def diskfmt(
request: pytest.FixtureRequest,
) -> Generator[Literal["h5ad", "zarr"], None, None]:
if (fmt := request.param[0]) == "h5ad":
yield fmt
else:
with ad.settings.override(zarr_write_format=request.param[1]):
yield fmt
@pytest.fixture
def diskfmt2(
diskfmt: Literal["h5ad", "zarr"],
) -> Generator[Literal["zarr", "h5ad"], None, None]:
if diskfmt == "h5ad":
with ad.settings.override(zarr_write_format=2):
yield "zarr"
else:
yield "h5ad"
@pytest.fixture(
params=[
pytest.param((..., (slice(None), slice(None))), id="ellipsis"),
pytest.param(((...,), (slice(None), slice(None))), id="ellipsis_tuple"),
pytest.param(
((..., slice(0, 10)), (slice(None), slice(0, 10))), id="obs-ellipsis"
),
pytest.param(
((slice(0, 10), ...), (slice(0, 10), slice(None))), id="var-ellipsis"
),
pytest.param(
((slice(0, 10), slice(0, 10), ...), (slice(0, 10), slice(0, 10))),
id="obs-var-ellipsis",
),
pytest.param(
((..., slice(0, 10), slice(0, 10)), (slice(0, 10), slice(0, 10))),
id="ellipsis-obs-var",
),
pytest.param(
((slice(0, 10), ..., slice(0, 10)), (slice(0, 10), slice(0, 10))),
id="obs-ellipsis-var",
),
]
)
def ellipsis_index_with_equivalent(
request,
) -> tuple[tuple[EllipsisType | slice, ...] | EllipsisType, tuple[slice, slice]]:
return request.param
@pytest.fixture
def ellipsis_index(
ellipsis_index_with_equivalent: tuple[
tuple[EllipsisType | slice, ...] | EllipsisType, tuple[slice, slice]
],
) -> tuple[EllipsisType | slice, ...] | EllipsisType:
return ellipsis_index_with_equivalent[0]
@pytest.fixture
def equivalent_ellipsis_index(
ellipsis_index_with_equivalent: tuple[
tuple[EllipsisType | slice, ...] | EllipsisType, tuple[slice, slice]
],
) -> tuple[slice, slice]:
return ellipsis_index_with_equivalent[1]
@pytest.fixture(scope="session")
def local_cluster_addr(
tmp_path_factory: pytest.TempPathFactory, worker_id: str
) -> Generator[str, None, None]:
# Adapted from https://pytest-xdist.readthedocs.io/en/latest/how-to.html#making-session-scoped-fixtures-execute-only-once
import dask.distributed as dd
def make_cluster(worker_id: str) -> dd.LocalCluster:
# If we're not using multiple pytest-xdist workers, let the cluster have multiple workers.
return dd.LocalCluster(
n_workers=1 if worker_id != "master" else 2, threads_per_worker=1
)
if worker_id == "master":
with make_cluster(worker_id) as cluster:
yield cluster.scheduler_address
return
# get the temp directory shared by all workers
root_tmp_dir = tmp_path_factory.getbasetemp().parent
fn = root_tmp_dir / "dask_scheduler_address.txt"
lock = FileLock(str(fn) + ".lock")
lock.acquire() # can’t use context manager, because we need to release the lock before yielding
address = fn.read_text() if fn.is_file() else None
if address:
lock.release()
yield address
return
with make_cluster(worker_id) as cluster:
fn.write_text(cluster.scheduler_address)
lock.release()
yield cluster.scheduler_address
#####################
# Dask tokenization #
#####################
# TODO: Should we be exporting this?
# sparray classes don't have tokenize defined yet, see: https://github.com/dask/dask/issues/10375
def normalize_sparse_matrix(x, attrs):
return (
type(x).__name__,
normalize_seq(normalize_token(getattr(x, key)) for key in attrs),
)
for cls, attrs in [
(sparse.dia_array, ("data", "offsets", "shape")),
(sparse.bsr_array, ("data", "indices", "indptr", "blocksize", "shape")),
(sparse.coo_array, ("data", "row", "col", "shape")),
(sparse.csr_array, ("data", "indices", "indptr", "shape")),
(sparse.csc_array, ("data", "indices", "indptr", "shape")),
(sparse.lil_array, ("data", "rows", "shape")),
]:
normalize_token.register(cls, partial(normalize_sparse_matrix, attrs=attrs))
@normalize_token.register(sparse.dok_array)
def normalize_dok_matrix(x):
return type(x).__name__, normalize_token(sorted(x.items()))
@normalize_token.register(ad.AnnData)
def tokenize_anndata(adata: ad.AnnData):
res = []
if adata.X is not None:
res.append(tokenize(adata.X))
res.extend([tokenize(adata.obs), tokenize(adata.var)])
for attr in ["obsm", "varm", "obsp", "varp", "layers"]:
elem = getattr(adata, attr)
res.append(tokenize(list(dict(elem).items())))
res.append(joblib.hash(adata.uns))
if adata.raw is not None:
res.append(tokenize(adata.raw.to_adata()))
return tuple(res)
|
