96def copy(
 97    input: SourceLike,
 98    output: Optional[SourceLike] = None,
 99    *,
100    block_shape: Optional[Tuple[int, ...]] = None,
101    job_type: str = "local",
102    job_config: Optional[RunnerConfig] = None,
103    num_workers: int = 1,
104    mask: Optional[SourceLike] = None,
105    block_ids: Optional[Sequence[int]] = None,
106    resume_from: Optional[str] = None,
107) -> SourceLike:
108    """Copy a source into an output, block-wise.
109
110    The data is read from ``input`` and written into ``output`` one block at a time. The input may
111    be any source, including a read-only ``FileSource`` (e.g. a tiff stack) or a ``wrapper`` source
112    whose transformation is computed on read; copying it materializes the transformed data to the
113    output. The data is written into the output as-is, so the output array's dtype governs and a
114    cast is applied on assignment when it differs from the input dtype.
115
116    Args:
117        input: The input data to copy (a numpy/zarr/n5 array or a `Source`).
118        output: The output array to write into. Optional for local execution — a numpy array
119            matching the input shape and dtype is allocated and returned if omitted; **required**
120            for distributed execution, where it must be a writable, file-backed (zarr/n5) array.
121        block_shape: Shape of the processing blocks. Defaults to the input chunk shape; required
122            for unchunked data.
123        job_type: Execution backend: one of ``"local"``, ``"subprocess"`` or ``"slurm"``.
124        job_config: Backend configuration (a `RunnerConfig` / `SlurmConfig`).
125        num_workers: Number of parallel workers (threads for ``local``, tasks for distributed
126            backends).
127        mask: Optional binary mask; only voxels within the mask are copied (out-of-mask output
128            voxels are left unchanged).
129        block_ids: Restrict processing to these block ids (e.g. to re-run previously failed blocks
130            into the existing ``output``). Mutually exclusive with ``resume_from``.
131        resume_from: Distributed only; the preserved temp folder of a failed run to resume (see
132            ``runner.run``); the missing blocks are written into ``output``. Mutually exclusive
133            with ``block_ids``.
134
135    Returns:
136        The output array (the provided ``output``, or a newly allocated numpy array).
137    """
138    return _copy_source(input, output, block_shape=block_shape, job_type=job_type,
139                        job_config=job_config, num_workers=num_workers, mask=mask, name="copy",
140                        block_ids=block_ids, resume_from=resume_from)

25def downsample(
26    input: SourceLike,
27    scale_factor: ScaleFactor,
28    output: Optional[SourceLike] = None,
29    *,
30    order: int = 0,
31    anti_aliasing: bool = False,
32    block_shape: Optional[Tuple[int, ...]] = None,
33    job_type: str = "local",
34    job_config: Optional[RunnerConfig] = None,
35    num_workers: int = 1,
36    mask: Optional[SourceLike] = None,
37    block_ids: Optional[Sequence[int]] = None,
38    resume_from: Optional[str] = None,
39) -> SourceLike:
40    """Downsample a source by an integer factor, block-wise.
41
42    The input is wrapped in a :class:`~bioimage_py.wrapper.ResizedSource` at the downscaled shape
43    (computed with :func:`~bioimage_py.util.downscale_shape`, ceil mode) and copied into the output.
44
45    The defaults (``order=0`` nearest, ``anti_aliasing=False``) are label-safe — they preserve the
46    input values and are appropriate for segmentations. For intensity / image data pass ``order=1``
47    (or higher) and ``anti_aliasing=True`` for a smooth, alias-free downsample.
48
49    Args:
50        input: The input data to downsample (a numpy/zarr/n5 array or a `Source`). 2D or 3D.
51        scale_factor: The downscaling factor: a single int (isotropic) or a per-axis sequence of
52            ints. Each factor must be ``>= 1`` (1 leaves that axis unchanged).
53        output: The output array to write into. Optional for local execution — a numpy array of the
54            downscaled shape and the input dtype is allocated and returned if omitted; **required**
55            for distributed execution, where it (and the input) must be file-backed (zarr/n5).
56        order: The interpolation order (0 to 5). Use ``0`` (nearest) for label data.
57        anti_aliasing: Whether to Gaussian pre-smooth before sampling to avoid aliasing.
58            Recommended for image data; leave ``False`` for labels.
59        block_shape: Shape of the processing blocks (in the downscaled output space). Defaults to
60            the resized source's chunk shape; required for unchunked data.
61        job_type: Execution backend: one of ``"local"``, ``"subprocess"`` or ``"slurm"``.
62        job_config: Backend configuration (a `RunnerConfig` / `SlurmConfig`).
63        num_workers: Number of parallel workers (threads for ``local``, tasks for distributed
64            backends).
65        mask: Optional binary mask in the downscaled output space; only voxels within the mask are
66            written (out-of-mask output voxels are left unchanged).
67        block_ids: Restrict processing to these block ids (in the downscaled output space), e.g. to
68            re-run previously failed blocks into the existing ``output``. Mutually exclusive with
69            ``resume_from``.
70        resume_from: Distributed only; the preserved temp folder of a failed run to resume (see
71            ``runner.run``); the missing blocks are written into ``output``. Mutually exclusive
72            with ``block_ids``.
73
74    Returns:
75        The output array (the provided ``output``, or a newly allocated numpy array).
76    """
77    src = as_source(input)
78    ndim = src.ndim
79    if isinstance(scale_factor, Integral):
80        factors: Tuple[int, ...] = (int(scale_factor),) * ndim
81    else:
82        factors = tuple(int(f) for f in scale_factor)
83        if len(factors) != ndim:
84            raise ValueError(
85                f"scale_factor {scale_factor} does not match the input dimensionality {ndim}."
86            )
87    if any(f < 1 for f in factors):
88        raise ValueError(
89            f"downsample requires scale factors >= 1, got {factors}; "
90            "use a ResizedSource directly to upsample."
91        )
92
93    target_shape = downscale_shape(src.shape, factors)
94    wrapped = ResizedSource(src, target_shape, order=order, anti_aliasing=anti_aliasing)
95    return _copy_source(wrapped, output, block_shape=block_shape, job_type=job_type,
96                        job_config=job_config, num_workers=num_workers, mask=mask, name="downsample",
97                        block_ids=block_ids, resume_from=resume_from)

18def get_runner(job_type: str, config: Optional[RunnerConfig] = None) -> Runner:
19    """Return a runner for the given job type.
20
21    Args:
22        job_type: One of ``"local"``, ``"subprocess"`` or ``"slurm"``.
23        config: Optional runner configuration.
24
25    Returns:
26        A :class:`~bioimage_py.runner.base.Runner` instance.
27
28    Raises:
29        ValueError: If ``job_type`` is unknown.
30    """
31    try:
32        cls = _RUNNERS[job_type.lower()]
33    except KeyError:
34        raise ValueError(f"Unknown job_type {job_type!r}; expected one of {sorted(_RUNNERS)}.")
35    return cls(config)

 46@dataclass
 47class SlurmConfig(RunnerConfig):
 48    """Configuration for the slurm runner.
 49
 50    Inherits ``poll_interval``, ``tmp_root`` and ``python_executable`` from
 51    :class:`RunnerConfig`. For slurm, ``tmp_root`` is **required** and must point at a
 52    shared filesystem visible to all compute nodes (not node-local ``/tmp``), and
 53    ``num_workers`` (passed to the op / ``run``) is interpreted as the array throttle — the
 54    maximum number of tasks allowed to run concurrently — independently of how many tasks
 55    the work is partitioned into (the partition granularity is set by ``tasks_per_worker``;
 56    see :class:`RunnerConfig`).
 57
 58    Cluster-specific values (``partition``, ``account``, ``constraint``, ``tmp_root``, ...)
 59    can be stored once in a user config file and reused as defaults; see
 60    :meth:`load` and :func:`write_slurm_config`.
 61
 62    Attributes:
 63        partition: The slurm partition to submit to.
 64        time: The per-task time limit (slurm time format, e.g. ``"01:00:00"``).
 65        mem: The per-task memory limit (e.g. ``"8G"``).
 66        cpus_per_task: Number of CPUs requested per task.
 67        gpus: Number of GPUs requested per task (emitted as ``--gpus`` only when > 0).
 68        account: The accounting project to charge.
 69        qos: The quality-of-service to request.
 70        constraint: A node feature constraint.
 71        shebang: Optional environment setup for the generated job script. If given, its
 72            first line must be an interpreter line (starting with ``#!``) which is placed at
 73            the top of the script; any remaining lines are emitted as an activation preamble
 74            *after* the ``#SBATCH`` directives (so the directives are still honoured). The
 75            preamble is for making the package importable on the node (e.g. ``module load``
 76            / ``LD_LIBRARY_PATH`` exports), not for choosing the interpreter: the worker is
 77            always launched with the absolute ``python_executable`` (defaulting to the
 78            submitting ``sys.executable``). ``None`` uses ``#!/bin/bash`` and that absolute
 79            interpreter, which needs no activation when the env lives on a shared
 80            filesystem. Example::
 81
 82                shebang = "#!/bin/bash\\nmodule load gcc\\nexport LD_LIBRARY_PATH=...:$LD_LIBRARY_PATH"
 83
 84        max_array_size: Override for the maximum number of array tasks per job. ``None``
 85            queries the cluster's ``MaxArraySize`` (falling back to a safe default). A run
 86            partitioned into more tasks than this is rejected up front with a clear error.
 87        latency_wait: Seconds to wait for a finished task's ``.success`` sentinel to become
 88            visible on a shared (NFS) filesystem before giving up on it. A task that the
 89            scheduler reports ``COMPLETED`` wrote its sentinel, but the orchestrating node's
 90            attribute cache can lag the compute node by up to the mount's ``acdirmax``
 91            (typically 60 s); this must comfortably exceed that. It only bounds the wait on a
 92            ``COMPLETED``-but-not-yet-visible task — a task is resolved the moment its
 93            sentinel appears, so a generous value does not slow down successful runs.
 94    """
 95
 96    partition: Optional[str] = None
 97    time: Optional[str] = None
 98    mem: Optional[str] = None
 99    cpus_per_task: int = 1
100    gpus: int = 0
101    account: Optional[str] = None
102    qos: Optional[str] = None
103    constraint: Optional[str] = None
104    shebang: Optional[str] = None
105    max_array_size: Optional[int] = None
106    latency_wait: float = 120.0
107
108    @classmethod
109    def load(cls, path: Optional[str] = None, **overrides: Any) -> "SlurmConfig":
110        """Build a config from the user config file, with explicit overrides taking precedence.
111
112        Precedence is ``overrides`` > config file ``[slurm]`` section > dataclass defaults.
113        This is the way to combine the stored user defaults with per-run tweaks; constructing
114        ``SlurmConfig(...)`` directly does **not** consult the file (an explicitly built config
115        is used verbatim).
116
117        Args:
118            path: Path to the config file. ``None`` resolves the default location (see
119                :func:`config_file_path`). A missing file is treated as empty.
120            **overrides: Field values that override the file defaults. Each name must be a
121                valid ``SlurmConfig`` field.
122
123        Returns:
124            A :class:`SlurmConfig` with file defaults filled in and overrides applied.
125
126        Raises:
127            ValueError: If the file or ``overrides`` contain an unknown field name.
128        """
129        _validate_keys(overrides, "load() overrides")
130        merged: Dict[str, Any] = dict(_read_slurm_defaults(path))
131        merged.update(overrides)
132        return cls(**merged)

135def config_file_path(path: Optional[str] = None) -> Path:
136    """Resolve the path to the user config file.
137
138    Resolution order: an explicit ``path`` argument, then the ``BIOIMAGE_PY_CONFIG``
139    environment variable, then ``$XDG_CONFIG_HOME/bioimage-py/config.toml`` (falling back to
140    ``~/.config/bioimage-py/config.toml``).
141
142    Args:
143        path: An explicit path that short-circuits the resolution. ``None`` resolves the
144            default location.
145
146    Returns:
147        The resolved path (not guaranteed to exist).
148    """
149    if path is not None:
150        return Path(path).expanduser()
151    env = os.environ.get("BIOIMAGE_PY_CONFIG")
152    if env:
153        return Path(env).expanduser()
154    base = os.environ.get("XDG_CONFIG_HOME") or os.path.join(os.path.expanduser("~"), ".config")
155    return Path(base) / "bioimage-py" / "config.toml"

197def write_slurm_config(path: Optional[str] = None, *, replace: bool = False, **fields: Any) -> str:
198    """Create or update the user config file with default slurm settings.
199
200    This is the supported way to set up cluster-specific defaults (partition, account,
201    constraint, ``tmp_root``, ...) instead of editing the file by hand. Provided fields are
202    merged into the existing ``[slurm]`` table by default (so the file can be built up over
203    several calls); ``None`` values are skipped, and any other top-level tables in the file
204    (reserved for future named profiles) are preserved.
205
206    Args:
207        path: Path to write to. ``None`` resolves the default location (see
208            :func:`config_file_path`); the parent directory is created if needed.
209        replace: If ``True``, replace the whole ``[slurm]`` table instead of merging into it.
210        **fields: Default field values to store. Each name must be a valid ``SlurmConfig``
211            field.
212
213    Returns:
214        The path that was written.
215
216    Raises:
217        ValueError: If ``fields`` contains an unknown field name.
218    """
219    _validate_keys(fields, "write_slurm_config()")
220    provided = {k: v for k, v in fields.items() if v is not None}
221    fp = config_file_path(path)
222    data = _parse_toml(fp)
223    section = {} if replace else dict(data.get("slurm", {}))
224    section.update(provided)
225    data["slurm"] = section
226
227    import tomli_w  # local import: only the writer needs the (optional-at-runtime) dependency.
228
229    fp.parent.mkdir(parents=True, exist_ok=True)
230    with open(fp, "wb") as f:
231        tomli_w.dump(data, f)
232    return str(fp)

30def as_source(obj: "SourceLike") -> Source:
31    """Convert a supported object into a :class:`Source`.
32
33    Idempotent on :class:`Source` inputs. numpy / zarr / z5py arrays are wrapped in an
34    :class:`ArraySource`. Bare paths are intentionally not supported (see the design doc).
35
36    Args:
37        obj: The object to convert.
38
39    Returns:
40        A :class:`Source`.
41
42    Raises:
43        TypeError: If the object cannot be converted (e.g. a string path).
44    """
45    if isinstance(obj, Source):
46        return obj
47    if isinstance(obj, (str, bytes)):
48        raise TypeError(
49            "Passing strings / file paths as a source is not supported. Open the array "
50            "yourself (e.g. with zarr or z5py) and pass the handle."
51        )
52    for predicate, converter in _CONVERTERS:
53        if predicate(obj):
54            return converter(obj)
55    # numpy and any array-like with shape/dtype fall back to ArraySource.
56    if isinstance(obj, np.ndarray) or (hasattr(obj, "shape") and hasattr(obj, "dtype")):
57        return ArraySource(obj)
58    raise TypeError(f"Cannot convert object of type {type(obj)!r} to a Source.")

103def open_source(
104    path: PathLike,
105    internal_path: Optional[str] = None,
106    format: Optional[str] = None,
107    mode: str = "r",
108    **kwargs: Any,
109) -> FileSource:
110    """Open a file-backed array as a :class:`Source`.
111
112    The format is inferred from the path extension (overridable via ``format``). ``internal_path``
113    selects the array inside a container; when omitted it defaults to the format's natural key
114    (e.g. ``"data"`` for mrc/nifti, ``"mag1"`` for knossos, ``""`` for a single image stack), and is
115    required for multi-array containers (hdf5/zarr/n5).
116
117    Args:
118        path: Path to the file or folder to open.
119        internal_path: Key of the array inside the container; format-dependent default if omitted.
120        format: Force a registered format name, overriding extension inference.
121        mode: Open mode. ``"r"`` (default) is read-only; write modes (``"a"``/``"r+"``/``"w"``)
122            are only honored for writable formats (zarr/n5/hdf5).
123        kwargs: Extra keyword arguments forwarded to the backend constructor.
124
125    Returns:
126        A :class:`FileSource` with a reopenable ``kind="file"`` spec.
127    """
128    fmt = format if format is not None else infer_format(path)
129    # Validate the format is installed up front (raises a clear error otherwise).
130    constructor_for_format(fmt)
131
132    handle = open_file(path, mode=mode, format=fmt, **kwargs)
133    key = internal_path if internal_path is not None else getattr(handle, "default_key", None)
134    dataset, recorded_key = _resolve_dataset(handle, key)
135
136    writable = is_writable_format(fmt) and mode != "r"
137    return FileSource(
138        dataset,
139        path=path,
140        internal_path=recorded_key,
141        format=fmt,
142        mode=mode,
143        open_kwargs=kwargs,
144        writable=writable,
145    )

130def open_cloudvolume(
131    cloudpath: str,
132    mip: int = 0,
133    fill_missing: bool = False,
134    bounded: bool = True,
135    cache: bool = False,
136    non_aligned_writes: bool = True,
137    offset: Optional[Tuple[int, int, int]] = None,
138    size: Optional[Tuple[int, int, int]] = None,
139    **kwargs: Any,
140) -> CloudVolumeSource:
141    """Open a CloudVolume (precomputed) layer as a writable ZYX :class:`Source`.
142
143    Args:
144        cloudpath: The CloudVolume cloudpath (e.g. ``"precomputed://..."`` or ``"file://..."``).
145        mip: The resolution (mip) level to open.
146        fill_missing: Whether to zero-fill missing chunks instead of raising. For a *writable*
147            output whose volume size is not a multiple of the chunk size, set this to ``True`` so
148            the partial boundary chunks can be read-modify-written into a fresh layer.
149        bounded: Whether reads/writes are restricted to the volume bounds.
150        cache: Whether to enable CloudVolume's local cache.
151        non_aligned_writes: Whether to allow writes that are not chunk-aligned (needed for the
152            partial blocks at the volume boundary in block-wise writes).
153        offset: Optional absolute XYZ origin of the view; defaults to the layer's ``voxel_offset``.
154        size: Optional XYZ size of the view; defaults to the layer's ``volume_size``.
155        kwargs: Extra keyword arguments forwarded to ``CloudVolume``.
156
157    Returns:
158        A :class:`CloudVolumeSource`.
159    """
160    from cloudvolume import CloudVolume
161
162    open_params: Dict[str, Any] = dict(
163        mip=mip,
164        fill_missing=fill_missing,
165        bounded=bounded,
166        cache=cache,
167        non_aligned_writes=non_aligned_writes,
168    )
169    open_params.update(kwargs)
170    volume = CloudVolume(cloudpath, progress=False, **open_params)
171    return CloudVolumeSource(volume, offset=offset, size=size, open_params=open_params)

 35class TensorStorePrecomputedSource(Source):
 36    """A local-ZYX :class:`Source` over a TensorStore neuroglancer-precomputed layer.
 37
 38    TensorStore exposes precomputed data in absolute ``(x, y, z, channel)`` coordinates. This
 39    adapter requires exactly those dimension labels and a single channel, and presents the full
 40    spatial domain as a local, zero-based ``(z, y, x)`` array. Reads and writes are transposed and
 41    translated to the absolute TensorStore domain internally.
 42
 43    The TensorStore read-chunk shape is exposed as :attr:`chunks`; the write-chunk shape (a shard
 44    for sharded precomputed data) is exposed as :attr:`shards`. This lets the runner keep all
 45    blocks that touch the same TensorStore write chunk on one worker and avoid concurrent shard
 46    updates.
 47
 48    Optional geometry and dtype arguments are expectations, not an arbitrary ROI definition. If
 49    supplied, they must match the layer metadata. This keeps the local origin aligned with the
 50    TensorStore chunk grids, which is required for safe shard routing.
 51
 52    Args:
 53        path: Filesystem path of the neuroglancer-precomputed layer.
 54        mode: ``"r"`` (default) for read-only access or ``"r+"`` for synchronous writes.
 55        offset_xyz: Optional expected absolute XYZ origin.
 56        size_zyx: Optional expected spatial shape in ZYX order.
 57        chunks_zyx: Optional expected read-chunk shape in ZYX order.
 58        dtype: Optional expected numpy dtype.
 59    """
 60
 61    def __init__(
 62        self,
 63        path: PathLike,
 64        mode: str = "r",
 65        *,
 66        offset_xyz: Optional[Sequence[int]] = None,
 67        size_zyx: Optional[Sequence[int]] = None,
 68        chunks_zyx: Optional[Sequence[int]] = None,
 69        dtype: Optional[Union[str, np.dtype, type]] = None,
 70    ) -> None:
 71        if mode not in ("r", "r+"):
 72            raise ValueError(
 73                f"TensorStorePrecomputedSource mode must be 'r' or 'r+', got {mode!r}."
 74            )
 75
 76        # TensorStore is an optional dependency; importing bioimage_py must not require it.
 77        import tensorstore as ts
 78
 79        self._path = os.path.abspath(os.fspath(path))
 80        self._mode = mode
 81        self._store = ts.open(
 82            {
 83                "driver": "neuroglancer_precomputed",
 84                "kvstore": {"driver": "file", "path": self._path},
 85            },
 86            open=True,
 87            read=True,
 88            write=mode == "r+",
 89        ).result()
 90
 91        domain = self._store.domain
 92        if domain.rank != 4:
 93            raise ValueError(
 94                "TensorStorePrecomputedSource requires a rank-4 (x, y, z, channel) layer, "
 95                f"got rank {domain.rank}."
 96            )
 97        labels = tuple(domain.labels)
 98        if labels != ("x", "y", "z", "channel"):
 99            raise ValueError(
100                "TensorStorePrecomputedSource requires dimension labels "
101                f"('x', 'y', 'z', 'channel'), got {labels!r}."
102            )
103
104        origin_xyzc = tuple(int(value) for value in domain.inclusive_min)
105        shape_xyzc = tuple(int(value) for value in domain.shape)
106        if shape_xyzc[3] != 1:
107            raise ValueError(
108                "TensorStorePrecomputedSource supports single-channel layers only, "
109                f"got {shape_xyzc[3]} channels."
110            )
111
112        read_chunk_xyzc = tuple(int(value) for value in self._store.chunk_layout.read_chunk.shape)
113        write_chunk_xyzc = tuple(int(value) for value in self._store.chunk_layout.write_chunk.shape)
114        grid_origin_xyzc = tuple(int(value) for value in self._store.chunk_layout.grid_origin)
115        if read_chunk_xyzc[3] != 1 or write_chunk_xyzc[3] != 1:
116            raise ValueError(
117                "TensorStorePrecomputedSource requires read and write chunks with one channel, "
118                f"got {read_chunk_xyzc!r} and {write_chunk_xyzc!r}."
119            )
120        if grid_origin_xyzc != origin_xyzc:
121            raise ValueError(
122                "TensorStorePrecomputedSource requires the chunk grid to start at the layer "
123                f"domain origin for safe shard routing, got grid origin {grid_origin_xyzc!r} "
124                f"and domain origin {origin_xyzc!r}."
125            )
126
127        self._offset_xyz = origin_xyzc[:3]
128        self._channel_origin = origin_xyzc[3]
129        self._shape = shape_xyzc[:3][::-1]
130        self._chunks = read_chunk_xyzc[:3][::-1]
131        self._shards = write_chunk_xyzc[:3][::-1]
132        self._dtype = np.dtype(self._store.dtype.numpy_dtype)
133
134        _validate_expected("offset_xyz", offset_xyz, self._offset_xyz)
135        _validate_expected("size_zyx", size_zyx, self._shape)
136        _validate_expected("chunks_zyx", chunks_zyx, self._chunks)
137        if dtype is not None and np.dtype(dtype) != self._dtype:
138            raise ValueError(
139                "TensorStore precomputed dtype does not match the layer: "
140                f"expected {np.dtype(dtype)}, found {self._dtype}."
141            )
142
143    @property
144    def store(self):
145        """The wrapped TensorStore handle."""
146        return self._store
147
148    @property
149    def shape(self) -> Tuple[int, ...]:
150        """The spatial shape in local ZYX order."""
151        return self._shape
152
153    @property
154    def dtype(self) -> np.dtype:
155        """The numpy dtype of the layer."""
156        return self._dtype
157
158    @property
159    def chunks(self) -> Tuple[int, ...]:
160        """The TensorStore read-chunk shape in ZYX order."""
161        return self._chunks
162
163    @property
164    def shards(self) -> Tuple[int, ...]:
165        """The TensorStore write-chunk (shard) shape in ZYX order."""
166        return self._shards
167
168    @property
169    def writable(self) -> bool:
170        """Whether the layer was opened in ``"r+"`` mode."""
171        return self._mode == "r+"
172
173    def _concrete_roi(self, roi: Tuple[slice, ...]) -> Tuple[slice, ...]:
174        """Resolve normalized local slices to concrete, in-bounds ZYX coordinates."""
175        concrete = []
176        for axis_slice, size in zip(roi, self._shape):
177            start, stop, step = axis_slice.indices(size)
178            if step != 1:
179                raise ValueError("TensorStorePrecomputedSource only supports a step of 1.")
180            concrete.append(slice(start, stop))
181        return tuple(concrete)
182
183    def _absolute_index(self, roi: Tuple[slice, ...]) -> Tuple[slice, slice, slice, int]:
184        """Translate a concrete local ZYX roi to an absolute TensorStore XYZC index."""
185        z, y, x = roi
186        ox, oy, oz = self._offset_xyz
187        return (
188            slice(ox + x.start, ox + x.stop),
189            slice(oy + y.start, oy + y.stop),
190            slice(oz + z.start, oz + z.stop),
191            self._channel_origin,
192        )
193
194    def _getitem(self, roi: Tuple[slice, ...]) -> np.ndarray:
195        roi = self._concrete_roi(roi)
196        value_xyz = self._store[self._absolute_index(roi)].read().result()
197        return np.asarray(value_xyz).transpose(2, 1, 0)
198
199    def _setitem(self, roi: Tuple[slice, ...], value: np.ndarray) -> None:
200        if not self.writable:
201            raise TypeError(
202                f"TensorStorePrecomputedSource opened in mode {self._mode!r} is read-only."
203            )
204        roi = self._concrete_roi(roi)
205        roi_shape = tuple(axis_slice.stop - axis_slice.start for axis_slice in roi)
206        value_zyx = np.broadcast_to(np.asarray(value, dtype=self._dtype), roi_shape)
207        value_xyz = value_zyx.transpose(2, 1, 0)
208        # Waiting on the future is deliberate: a completed block means its output is durable, and
209        # write failures are attributed to the task that issued them.
210        self._store[self._absolute_index(roi)].write(value_xyz).result()
211
212    def to_spec(self) -> SourceSpec:
213        """Return a serializable spec containing the verified layer metadata."""
214        return SourceSpec(
215            kind="tensorstore_precomputed",
216            path=self._path,
217            params={
218                "mode": self._mode,
219                "offset_xyz": list(self._offset_xyz),
220                "size_zyx": list(self._shape),
221                "chunks_zyx": list(self._chunks),
222                "dtype": self._dtype.str,
223            },
224        )
225
226    @staticmethod
227    def reopen(spec: SourceSpec) -> "TensorStorePrecomputedSource":
228        """Reopen a source from current or legacy adapter specs.
229
230        Specs written by the earlier local shim did not record a mode; those reopen read-only.
231        """
232        params = dict(spec.params)
233        params.setdefault("mode", "r")
234        return TensorStorePrecomputedSource(spec.path, **params)

237def open_tensorstore_precomputed(
238    path: PathLike,
239    mode: str = "r",
240    *,
241    offset_xyz: Optional[Sequence[int]] = None,
242    size_zyx: Optional[Sequence[int]] = None,
243    chunks_zyx: Optional[Sequence[int]] = None,
244    dtype: Optional[Union[str, np.dtype, type]] = None,
245) -> TensorStorePrecomputedSource:
246    """Open a neuroglancer-precomputed layer as a local-ZYX :class:`Source`.
247
248    Metadata arguments are optional consistency checks. The layer metadata is always authoritative;
249    passing a mismatching expectation raises instead of silently producing incorrect coordinates.
250    """
251    return TensorStorePrecomputedSource(
252        path,
253        mode=mode,
254        offset_xyz=offset_xyz,
255        size_zyx=size_zyx,
256        chunks_zyx=chunks_zyx,
257        dtype=dtype,
258    )