subset, recast, split, concat: write beside the output, rename into place

docs/COMMANDS.md

@@ -83,6 +83,13 @@ adata subset data.h5ad --inplace --obs barcodes.txt
 `raw/` is carried over and matched against its **own** var axis, which usually
 holds more genes than the main object.
 
+The selection is resolved before anything is written, and the output is
+written beside its destination (`out.adata-tmp.h5ad`) and moved into place
+only once complete. A query that fails therefore leaves nothing behind, and an
+existing output (or, with `--inplace`, the source) is replaced only once its
+successor is finished. The output path may not name the input; use
+`--inplace`.
+
 ## `convert`
 
 Change a store's format. The output's extension picks the format; `--to`
@@ -225,8 +232,9 @@ bigger than it started. `indices` and `indptr` keep their own widths, which
 can differ: a narrow matrix with more than 2^31 nonzeros needs int64 offsets
 over int32 coordinates, and both are range-checked before writing.
 
-The output path may not name the input; use `--inplace`, which writes to a
-temporary file and swaps it in only once the conversion has finished.
+The output path may not name the input; use `--inplace`. Either way the new
+store is written beside its destination and moved into place only once the
+conversion has finished, so a failure leaves the original untouched.
 
 Transposing streams by default and works on matrices too large to load, at
 the cost of two extra passes over the nonzeros. `--in-memory` is faster when
@@ -256,6 +264,10 @@ adata split data.h5ad --by cluster -o out/ --min-size 50
 | `--manifest / --no-manifest` | Write a CSV listing the outputs (default on) |
 
 Labels are sanitised for use as filenames; collisions get a numeric suffix.
+Every store is written beside its destination and moved into place only once
+all of them are complete, and the manifest is written last, so a split that
+fails leaves no partial stores. An output that would land on the input is
+refused.
 
 ## `concat`
 

src/adata/commands/convert.py

@@ -20,16 +20,10 @@
     detect_format,
     output_format,
 )
+from adata.util.path import same_path
 from adata.util.staging import staged_outputs
 
 
-def _same_path(left: Path, right: Path) -> bool:
-    try:
-        return left.resolve() == right.resolve()
-    except OSError:
-        return False
-
-
 def convert_store(
     file: Path,
     output: Path,
@@ -56,7 +50,7 @@ def convert_store(
 
     plan = converter.plan(file, output, src_format, dst_format, options)
     for out in plan.outputs:
-        if _same_path(file, out):
+        if same_path(file, out):
             # Writing a Zarr store over its own source empties it before a
             # single element has been read.
             raise ValueError(f"Output path is the input: {out}")

src/adata/commands/recast.py

@@ -2,13 +2,14 @@
 
 Writes a whole new store rather than editing one in place, for the same
 reason `subset` does: a conversion that fails half way leaves the original
-untouched, and `--inplace` becomes an atomic swap of a finished file rather
-than a partial edit of a live one.
+untouched. The new store is written beside its destination and renamed into
+place only once complete (see `adata.util.staging`), so `--inplace` replaces a
+finished file rather than editing a live one, and the original is deleted only
+after its successor is in place.
 """
 
 from __future__ import annotations
 
-import shutil
 from pathlib import Path
 from typing import Any, List, Optional
 
@@ -25,15 +26,16 @@
 from adata.elements import spec
 from adata.elements.write import ensure_anndata_skeleton
 from adata.storage import copy_tree, detect_backend, is_group, open_store
-
-#: What `--all` means. The matrices an AnnData object is built from, not
-#: every 2-D array in the file: obsm/varm hold embeddings, whose dtype is
-#: rarely what anyone is trying to shrink. Those still convert by path.
-ALL_PREFIXES = ("X", "layers", "raw/X")
-
+from adata.util.path import same_path
+from adata.util.staging import staged_outputs
 
 def discover_matrices(root: Any) -> List[str]:
-    """Entry paths `--all` resolves to, in a stable order."""
+    """Entry paths `--all` resolves to, in a stable order.
+
+    The matrices an AnnData object is built from, not every 2-D array in the
+    file: obsm/varm hold embeddings, whose dtype is rarely what anyone is
+    trying to shrink. Those still convert by path.
+    """
     found: List[str] = []
     if "X" in root:
         found.append("X")
@@ -44,14 +46,6 @@ def discover_matrices(root: Any) -> List[str]:
     return found
 
 
-def _same_path(left: Path, right: Path) -> bool:
-    """Do these name the same store, through symlinks and `..` alike?"""
-    try:
-        return left.resolve() == right.resolve()
-    except OSError:  # pragma: no cover - unresolvable path
-        return left.absolute() == right.absolute()
-
-
 def _resolve(root: Any, path: str) -> Any:
     obj = root
     for part in path.split("/"):
@@ -96,7 +90,7 @@ def recast_store(
         else None
     )
 
-    if output is not None and not inplace and _same_path(file, output):
+    if output is not None and not inplace and same_path(file, output):
         # Opening the destination "w" clears it while the source is still
         # being read from it. HDF5 refuses; Zarr does not, and quietly
         # produced an empty store where the data used to be.
@@ -105,69 +99,54 @@ def recast_store(
             "it, which writes to a temporary file first."
         )
 
-    if inplace:
-        backend = detect_backend(file)
-        if backend == "zarr":
-            base = file.stem if file.suffix else file.name
-            dst_path = file.with_name(f"{base}.recast-tmp.zarr")
-        else:
-            dst_path = file.with_name(f"{file.name}.recast-tmp")
-        if dst_path.exists():
-            raise FileExistsError(f"Temporary path already exists: {dst_path}")
-    else:
-        dst_path = output
+    dst_path = file if inplace else Path(output)
 
     if zarr_format is None and detect_backend(file) == "zarr":
         with open_store(file, "r") as probe:
             zarr_format = probe.zarr_format
 
-    with open_store(file, "r") as src_store:
-        src = src_store.root
-        targets = discover_matrices(src) if convert_all else list(entries or [])
-        if not targets:
-            raise ValueError(
-                "No matrices selected. Name one (e.g. `X`) or pass --all."
-            )
-        # Every check first, before the destination exists. A refusal must
-        # not leave behind a store holding a copy of obs and var and nothing
-        # else -- the caller cannot tell that from a finished conversion.
-        plans = {
-            path: plan_conversion(
-                _resolve(src, path),
-                path,
-                dtype=data_dtype,
-                index_dtype=index_dtype,
-                layout=layout,
-                chunk=chunk,
-                force=force,
-                console=console,
+    # The source is opened inside the staging block, so it is closed before
+    # `--inplace` renames it aside.
+    with staged_outputs([dst_path]) as (temp,):
+        with open_store(file, "r") as src_store:
+            src = src_store.root
+            targets = discover_matrices(src) if convert_all else list(entries or [])
+            if not targets:
+                raise ValueError(
+                    "No matrices selected. Name one (e.g. `X`) or pass --all."
+                )
+            # Every check first, before the destination exists. A refusal must
+            # not leave behind a store holding a copy of obs and var and nothing
+            # else -- the caller cannot tell that from a finished conversion.
+            plans = {
+                path: plan_conversion(
+                    _resolve(src, path),
+                    path,
+                    dtype=data_dtype,
+                    index_dtype=index_dtype,
+                    layout=layout,
+                    chunk=chunk,
+                    force=force,
+                    console=console,
+                )
+                for path in targets
+            }
+
+            console.print(
+                f"[cyan]Converting {len(targets)} matrix/matrices:[/] "
+                + ", ".join(targets)
             )
-            for path in targets
-        }
-
-        console.print(
-            f"[cyan]Converting {len(targets)} matrix/matrices:[/] "
-            + ", ".join(targets)
-        )
 
-        with open_store(dst_path, "w", zarr_format=zarr_format) as dst_store:
-            dst = dst_store.root
-            _write(
-                src, dst, plans,
-                chunk=chunk, chunk_rows=chunk_rows, in_memory=in_memory,
-                console=console,
-            )
-            ensure_anndata_skeleton(dst)
+            with open_store(temp, "w", zarr_format=zarr_format) as dst_store:
+                dst = dst_store.root
+                _write(
+                    src, dst, plans,
+                    chunk=chunk, chunk_rows=chunk_rows, in_memory=in_memory,
+                    console=console,
+                )
+                ensure_anndata_skeleton(dst)
 
     if inplace:
-        if file.is_dir():
-            shutil.rmtree(file)
-        elif file.exists():
-            file.unlink()
-        if dst_path.is_dir():
-            shutil.move(str(dst_path), str(file))
-        else:
-            dst_path.replace(file)
         console.print(f"[green]Converted[/] {file}")
     else:
         console.print(f"[green]Wrote[/] {dst_path}")

src/adata/commands/split.py

@@ -4,6 +4,11 @@
 streaming: the source is never loaded into memory, only the column being split
 on is read, and each output is produced by the same subset machinery used by
 `adata subset`, fanned out so every matrix is read once for all groups.
+
+Every output is written beside its destination and renamed into place only
+once all of them are complete, so a split that fails leaves no partial stores
+(and existing outputs as they were). The manifest is written last, so it
+exists only when everything it lists does.
 """
 
 from __future__ import annotations
@@ -19,6 +24,8 @@
 from adata.core.subset import split_h5ad
 from adata.storage import detect_backend, open_store
 from adata.util.names import unique_names
+from adata.util.path import same_path
+from adata.util.staging import staged_outputs
 
 
 def split_store(
@@ -78,6 +85,13 @@ def split_store(
     for label, out_path, count in planned:
         console.print(f"  {label!r} -> {out_path} ({count} {axis})")
 
+    for _, out_path, _ in planned:
+        if same_path(file, out_path):
+            raise ValueError(
+                f"Output path is the input: {out_path}. Choose another "
+                "--output-dir."
+            )
+
     if dry_run:
         console.print("[yellow]Dry run: nothing written.[/]")
         return planned
@@ -86,20 +100,21 @@ def split_store(
 
     # One call for every group, so each matrix is read once rather than once
     # per group.
-    split_h5ad(
-        file,
-        [
-            (
-                out_path,
-                np.sort(groups[label]) if axis == "obs" else None,
-                np.sort(groups[label]) if axis == "var" else None,
-            )
-            for label, out_path, _ in planned
-        ],
-        chunk_rows=chunk_rows,
-        console=console,
-        zarr_format=zarr_format,
-    )
+    with staged_outputs([out_path for _, out_path, _ in planned]) as temps:
+        split_h5ad(
+            file,
+            [
+                (
+                    temp,
+                    np.sort(groups[label]) if axis == "obs" else None,
+                    np.sort(groups[label]) if axis == "var" else None,
+                )
+                for (label, _, _), temp in zip(planned, temps)
+            ],
+            chunk_rows=chunk_rows,
+            console=console,
+            zarr_format=zarr_format,
+        )
 
     if manifest:
         _write_manifest(file, output_dir, column, planned, console)

src/adata/core/concat.py

@@ -46,6 +46,7 @@
     is_group,
     open_store,
 )
+from adata.util.staging import staged_outputs
 
 MERGE_STRATEGIES = ("same", "unique", "first", "only")
 #: What the CLI accepts. "drop" is the default and maps to no merge at all.
@@ -906,7 +907,11 @@ def concat_on_disk(
         # half-written store behind.
         check_matrix_encodings(roots, console)
 
-        with open_store(output, "w", zarr_format=zarr_format) as dst_store:
+        # Written beside the output and renamed into place once complete, so
+        # a failure part way leaves nothing that reads as a finished store.
+        with staged_outputs([output]) as (temp,), open_store(
+            temp, "w", zarr_format=zarr_format
+        ) as dst_store:
             dst = dst_store.root
             _write_obs(
                 dst, roots, target_obs, obs_counts, join, label, input_keys, console