pyadata-cli 0.5.0__py3-none-any.whl

adata/elements/write.py

@@ -0,0 +1,255 @@
+"""Writing AnnData elements in the current on-disk spec.
+
+Every element this tool creates is tagged with its `encoding-type` and
+`encoding-version` here, so a store the CLI writes is indistinguishable from
+one anndata wrote. Text always goes out as variable-length UTF-8 -- fixed-width
+bytes are readable by anndata but are not what `string-array` specifies, and on
+Zarr v3 they produce a dtype other libraries cannot read.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Iterable, Optional, Sequence
+
+import numpy as np
+
+from adata.elements import spec
+from adata.storage import create_dataset, is_zarr_group
+
+
+def _backend_of(parent: Any) -> str:
+    return "zarr" if is_zarr_group(parent) else "hdf5"
+
+
+def _replace(parent: Any, name: str) -> None:
+    """Remove an existing member so it can be rewritten."""
+    if name in parent:
+        del parent[name]
+
+
+def write_mapping(parent: Any, name: str, replace: bool = False) -> Any:
+    """Create (or fetch) a group tagged as a `dict` mapping."""
+    if replace:
+        _replace(parent, name)
+    group = parent[name] if name in parent else parent.create_group(name)
+    spec.set_encoding(group, spec.DICT)
+    return group
+
+
+def write_string_array(
+    parent: Any, name: str, values: Iterable[Any], replace: bool = False
+) -> Any:
+    """Write a `string-array`: variable-length UTF-8 on either backend.
+
+    Multi-dimensional input keeps its shape; `list()`-ing it first would
+    silently flatten a 2-D array of labels into a vector.
+    """
+    if replace:
+        _replace(parent, name)
+    data = (
+        values
+        if isinstance(values, np.ndarray)
+        else np.asarray(list(values), dtype=object)
+    )
+    ds = create_dataset(parent, name, data=data)
+    spec.set_encoding(ds, spec.STRING_ARRAY)
+    return ds
+
+
+def write_dense(
+    parent: Any, name: str, data: Any, replace: bool = False, **kwargs: Any
+) -> Any:
+    """Write a dense numeric `array`."""
+    if replace:
+        _replace(parent, name)
+    ds = create_dataset(parent, name, data=np.asarray(data), **kwargs)
+    spec.set_encoding(ds, spec.ARRAY)
+    return ds
+
+
+def write_scalar(parent: Any, name: str, value: Any, replace: bool = False) -> Any:
+    """Write a 0-d scalar as `string` or `numeric-scalar`, per its type."""
+    if replace:
+        _replace(parent, name)
+
+    if isinstance(value, (str, bytes, np.str_, np.bytes_)):
+        text = value.decode("utf-8") if isinstance(value, bytes) else str(value)
+        ds = create_dataset(parent, name, data=np.asarray(text, dtype=object))
+        spec.set_encoding(ds, spec.STRING)
+        return ds
+
+    arr = np.asarray(value)
+    ds = create_dataset(parent, name, data=arr)
+    spec.set_encoding(ds, spec.NUMERIC_SCALAR)
+    return ds
+
+
+def write_null(parent: Any, name: str, replace: bool = False) -> Any:
+    """Write an explicit null, matching how anndata >= 0.12 encodes `None`.
+
+    HDF5 uses a null dataspace (`h5py.Empty`); Zarr uses a 0-d boolean array.
+    Both carry `encoding-type: null`, so `None` round-trips instead of needing
+    an invented marker attribute.
+    """
+    if replace:
+        _replace(parent, name)
+
+    if is_zarr_group(parent):
+        ds = parent.create_array(name, shape=(), dtype=bool)
+    else:
+        import h5py
+
+        ds = parent.create_dataset(name, data=h5py.Empty("f4"))
+    spec.set_encoding(ds, spec.NULL)
+    return ds
+
+
+def write_categorical(
+    parent: Any,
+    name: str,
+    codes: Any,
+    categories: Sequence[Any],
+    ordered: bool = False,
+    replace: bool = False,
+) -> Any:
+    """Write a `categorical` group of `categories` + `codes`.
+
+    A code of -1 denotes a missing value, as the spec requires.
+    """
+    if replace:
+        _replace(parent, name)
+    group = parent.create_group(name)
+    spec.set_encoding(group, spec.CATEGORICAL)
+    group.attrs["ordered"] = bool(ordered)
+
+    write_string_array(group, "categories", categories)
+    write_dense(group, "codes", np.asarray(codes, dtype=_codes_dtype(len(categories))))
+    return group
+
+
+def _codes_dtype(n_categories: int) -> Any:
+    """Smallest signed integer dtype that can index `n_categories` plus -1."""
+    if n_categories < 128:
+        return np.int8
+    if n_categories < 32_768:
+        return np.int16
+    return np.int32
+
+
+def write_masked(
+    parent: Any,
+    name: str,
+    values: Any,
+    mask: Any,
+    enc_type: str,
+    na_value: Optional[str] = None,
+    replace: bool = False,
+) -> Any:
+    """Write a nullable element as a `values` + `mask` group.
+
+    A True entry in `mask` marks a missing value.
+    """
+    if enc_type not in spec.MASKED_TYPES:
+        raise ValueError(f"{enc_type!r} is not a masked encoding.")
+    if replace:
+        _replace(parent, name)
+
+    group = parent.create_group(name)
+    spec.set_encoding(group, enc_type)
+    if na_value is not None:
+        group.attrs["na-value"] = na_value
+
+    if enc_type == spec.NULLABLE_STRING_ARRAY:
+        write_string_array(group, "values", values)
+    else:
+        write_dense(group, "values", values)
+    write_dense(group, "mask", np.asarray(mask, dtype=bool))
+    return group
+
+
+def write_sparse(
+    parent: Any,
+    name: str,
+    data: Any,
+    indices: Any,
+    indptr: Any,
+    shape: Sequence[int],
+    enc_type: str = spec.CSR_MATRIX,
+    replace: bool = False,
+) -> Any:
+    """Write a CSR/CSC sparse matrix group."""
+    if enc_type not in spec.SPARSE_TYPES:
+        raise ValueError(f"{enc_type!r} is not a sparse encoding.")
+    if replace:
+        _replace(parent, name)
+
+    group = parent.create_group(name)
+    spec.set_encoding(group, enc_type)
+    set_shape_attr(group, shape)
+
+    create_dataset(group, "data", data=np.asarray(data))
+    create_dataset(group, "indices", data=np.asarray(indices))
+    create_dataset(group, "indptr", data=np.asarray(indptr))
+    return group
+
+
+def set_shape_attr(group: Any, shape: Sequence[int]) -> None:
+    """Set a sparse group's `shape`, in the form the backend's attrs accept.
+
+    Zarr attributes must be JSON-serialisable, so a numpy array cannot be
+    stored there; HDF5 conventionally holds an integer array.
+    """
+    dims = [int(d) for d in shape]
+    if is_zarr_group(group):
+        group.attrs["shape"] = dims
+    else:
+        group.attrs["shape"] = np.array(dims, dtype=np.int64)
+
+
+def write_dataframe_header(
+    parent: Any,
+    name: str,
+    index_values: Iterable[Any],
+    column_order: Sequence[str],
+    index_name: str = "_index",
+    replace: bool = True,
+) -> Any:
+    """Create a `dataframe` group with its index and declared column order.
+
+    Columns themselves are written afterwards by the caller; `column-order`
+    records the authored order so readers do not fall back to the backend's
+    own (alphabetical, on HDF5) enumeration.
+    """
+    if replace:
+        _replace(parent, name)
+
+    group = parent.create_group(name)
+    spec.set_encoding(group, spec.DATAFRAME)
+    group.attrs["_index"] = index_name
+    set_column_order(group, column_order)
+    write_string_array(group, index_name, index_values)
+    return group
+
+
+def set_column_order(group: Any, columns: Sequence[str]) -> None:
+    """Record a dataframe's column order in the form the backend accepts.
+
+    Zarr attributes must be JSON, so a plain list is used. HDF5 needs an
+    explicit variable-length UTF-8 dtype -- an object array has no native
+    HDF5 equivalent, and an empty one cannot be inferred at all.
+    """
+    names = [str(c) for c in columns]
+    if is_zarr_group(group):
+        group.attrs["column-order"] = names
+    else:
+        import h5py
+
+        group.attrs["column-order"] = np.array(
+            names, dtype=h5py.string_dtype(encoding="utf-8")
+        )
+
+
+def ensure_anndata_skeleton(root: Any) -> None:
+    """Create the optional mapping groups an AnnData store is expected to have."""
+    for key in ("layers", "obsm", "obsp", "varm", "varp", "uns"):
+        write_mapping(root, key)

adata/formats/__init__.py

@@ -0,0 +1 @@
+"""Format-specific import/export helpers."""

adata/formats/array.py

@@ -0,0 +1,109 @@
+from __future__ import annotations
+
+from pathlib import Path
+import sys
+from typing import Any, Optional
+
+import numpy as np
+from rich.console import Console
+
+from adata.formats.common import _get_encoding_type, _resolve
+from adata.formats.validate import validate_dimensions
+from adata.elements.write import write_dense
+from adata.storage import create_dataset, is_dataset, is_group
+from adata.util.path import norm_path
+
+
+def export_npy(
+    root: Any,
+    obj: str,
+    out: Optional[Path],
+    chunk_elements: int,
+    console: Console,
+) -> None:
+    """Write a dense array to a .npy file, or to stdout when `out` is None.
+
+    Writing to a file streams in chunks through a memory-mapped output. Stdout
+    is not seekable, so that path materialises the array first and is only
+    suitable for arrays that fit in memory.
+    """
+    h5obj = _resolve(root, obj)
+
+    if is_group(h5obj):
+        enc = _get_encoding_type(h5obj)
+        if enc in ("nullable-integer", "nullable-boolean", "nullable-string-array"):
+            if "values" not in h5obj:
+                raise ValueError(f"Encoded group '{obj}' is missing 'values' dataset.")
+            ds = h5obj["values"]
+            console.print(f"[dim]Exporting nullable array values from '{obj}'[/]")
+        else:
+            raise ValueError(
+                f"Target '{obj}' is a group with encoding '{enc}'; cannot export as .npy directly."
+            )
+    elif is_dataset(h5obj):
+        ds = h5obj
+    else:
+        raise ValueError("Target is not an array-like object.")
+
+    if out is None or str(out) == "-":
+        np.save(sys.stdout.buffer, np.asarray(ds[...]), allow_pickle=False)
+        sys.stdout.buffer.flush()
+        return
+
+    out.parent.mkdir(parents=True, exist_ok=True)
+    mm = np.lib.format.open_memmap(out, mode="w+", dtype=ds.dtype, shape=ds.shape)
+    try:
+        if ds.shape == ():
+            mm[...] = ds[()]
+            console.print(f"[green]Wrote[/] {out}")
+            return
+
+        if ds.ndim == 1:
+            n = int(ds.shape[0])
+            step = max(1, int(chunk_elements))
+            with console.status(f"[magenta]Exporting {obj} to {out}...[/]") as status:
+                for start in range(0, n, step):
+                    end = min(start + step, n)
+                    status.update(
+                        f"[magenta]Exporting {obj}: {start}-{end} of {n}...[/]"
+                    )
+                    mm[start:end] = ds[start:end]
+            console.print(f"[green]Wrote[/] {out}")
+            return
+
+        n0 = int(ds.shape[0])
+        row_elems = int(np.prod(ds.shape[1:])) if ds.ndim > 1 else 1
+        step0 = max(1, int(chunk_elements) // max(1, row_elems))
+        with console.status(f"[magenta]Exporting {obj} to {out}...[/]") as status:
+            for start in range(0, n0, step0):
+                end = min(start + step0, n0)
+                status.update(
+                    f"[magenta]Exporting {obj}: {start}-{end} of {n0}...[/]"
+                )
+                mm[start:end, ...] = ds[start:end, ...]
+        console.print(f"[green]Wrote[/] {out}")
+    finally:
+        del mm
+
+
+def import_npy(
+    root: Any,
+    obj: str,
+    input_file: Path,
+    console: Console,
+) -> None:
+    obj = norm_path(obj)
+    arr = np.load(input_file)
+
+    validate_dimensions(root, obj, arr.shape, console)
+
+    parts = obj.split("/")
+    parent = root
+    for part in parts[:-1]:
+        parent = parent[part] if part in parent else parent.create_group(part)
+    name = parts[-1]
+
+    write_dense(parent, name, arr, replace=True)
+
+    shape_str = "×".join(str(d) for d in arr.shape)
+    console.print(f"[green]Imported[/] {shape_str} array into '{obj}'")

adata/formats/common.py

@@ -0,0 +1,51 @@
+from __future__ import annotations
+
+from typing import Any
+
+import numpy as np
+
+from adata.storage import is_dataset, is_group
+from adata.util.path import norm_path
+
+
+def _get_encoding_type(group: Any) -> str:
+    enc = group.attrs.get("encoding-type", "")
+    if isinstance(enc, bytes):
+        enc = enc.decode("utf-8")
+    return str(enc)
+
+
+def _resolve(root: Any, obj: str) -> Any:
+    obj = norm_path(obj)
+    if obj not in root:
+        raise KeyError(f"'{obj}' not found in the file.")
+    return root[obj]
+
+
+def _check_json_exportable(h5obj: Any, max_elements: int, path: str = "") -> None:
+    if is_dataset(h5obj):
+        if h5obj.shape == ():
+            return
+        n = int(np.prod(h5obj.shape)) if h5obj.shape else 0
+        if n > max_elements:
+            obj_name = getattr(h5obj, "name", "<object>")
+            raise ValueError(
+                f"Cannot export to JSON: '{path or obj_name}' has {n} elements "
+                f"(max {max_elements}). Use --max-elements to increase limit."
+            )
+        return
+
+    if is_group(h5obj):
+        enc = _get_encoding_type(h5obj)
+        if enc in ("csr_matrix", "csc_matrix"):
+            obj_name = getattr(h5obj, "name", "<object>")
+            raise ValueError(
+                f"Cannot export to JSON: '{path or obj_name}' is a sparse matrix. "
+                "Export it as .mtx instead."
+            )
+
+        for key in h5obj.keys():
+            child = h5obj[key]
+            child_path = f"{path}/{key}" if path else key
+            if is_group(child) or is_dataset(child):
+                _check_json_exportable(child, max_elements=max_elements, path=child_path)

adata/formats/dataframe.py

@@ -0,0 +1,234 @@
+from __future__ import annotations
+
+import csv
+import sys
+from contextlib import nullcontext
+from pathlib import Path
+from typing import Any, List, Optional, Tuple
+
+import numpy as np
+from rich.console import Console
+
+from adata.core.read import col_chunk_as_strings
+from adata.formats.common import _resolve
+from adata.util.path import norm_path
+from adata.elements.read import dataframe_columns, element_len, resolve_index
+from adata.formats.validate import validate_dimensions
+from adata.elements.write import (
+    write_categorical,
+    write_mapping,
+    write_dataframe_header,
+    write_dense,
+    write_string_array,
+)
+from adata.storage import create_dataset, is_group, is_zarr_group
+
+
+def export_dataframe(
+    root: Any,
+    axis: str,
+    columns: Optional[List[str]],
+    out: Optional[Path],
+    chunk_rows: int,
+    head: Optional[int],
+    console: Console,
+) -> None:
+    """Stream a dataframe group out as CSV.
+
+    `axis` is any path to a dataframe-encoded group, not only "obs" or "var" --
+    a store may hold dataframes in obsm, varm, uns or raw as well.
+    """
+    group = _resolve(root, axis)
+    if not is_group(group):
+        raise ValueError(f"'{axis}' is not a group and cannot be exported as CSV.")
+
+    try:
+        index, index_name = resolve_index(group)
+    except KeyError as exc:
+        raise ValueError(
+            f"'{axis}' is not a dataframe: it has no index. "
+            "CSV export needs a dataframe-encoded group such as 'obs', 'var' "
+            "or 'raw/var'."
+        ) from exc
+    n_rows = element_len(index)
+
+    if isinstance(index_name, bytes):
+        index_name = index_name.decode("utf-8")
+
+    if columns:
+        col_names = list(columns)
+    else:
+        col_names = dataframe_columns(group, index_name)
+
+    if index_name not in col_names:
+        col_names.insert(0, index_name)
+    else:
+        col_names = [index_name] + [c for c in col_names if c != index_name]
+
+    if head is not None and head > 0:
+        n_rows = min(n_rows, head)
+
+    if out is None or str(out) == "-":
+        out_fh = sys.stdout
+    else:
+        out_fh = open(out, "w", newline="", encoding="utf-8")
+    writer = csv.writer(out_fh)
+
+    try:
+        writer.writerow(col_names)
+        cat_cache = {}
+
+        use_status = out_fh is not sys.stdout
+        status_ctx = (
+            console.status(f"[magenta]Exporting {axis} table to {out}...[/]")
+            if use_status
+            else nullcontext()
+        )
+
+        with status_ctx as status:
+            for start in range(0, n_rows, chunk_rows):
+                end = min(start + chunk_rows, n_rows)
+                if use_status and status:
+                    status.update(
+                        f"[magenta]Exporting rows {start}-{end} of {n_rows}...[/]"
+                    )
+                cols_data: List[List[str]] = []
+                for col in col_names:
+                    cols_data.append(
+                        col_chunk_as_strings(group, col, start, end, cat_cache)
+                    )
+                for row_idx in range(end - start):
+                    row = [
+                        cols_data[col_idx][row_idx]
+                        for col_idx in range(len(col_names))
+                    ]
+                    writer.writerow(row)
+    finally:
+        if out_fh is not sys.stdout:
+            out_fh.close()
+
+
+def _read_csv(
+    input_file: Path,
+    index_column: Optional[str],
+) -> Tuple[List[dict], List[str], List[str], str]:
+    with open(input_file, "r", encoding="utf-8", newline="") as f:
+        reader = csv.DictReader(f)
+        if reader.fieldnames is None:
+            raise ValueError("CSV file has no header.")
+        fieldnames = list(reader.fieldnames)
+
+        if index_column:
+            if index_column not in fieldnames:
+                raise ValueError(
+                    f"Index column '{index_column}' not found in CSV. "
+                    f"Available columns: {', '.join(fieldnames)}"
+                )
+            idx_col = index_column
+        else:
+            idx_col = fieldnames[0]
+
+        rows = list(reader)
+
+    index_values = [row[idx_col] for row in rows]
+    data_columns = [c for c in fieldnames if c != idx_col]
+
+    return rows, data_columns, index_values, idx_col
+
+
+def _looks_categorical(values: List[str], n_rows: int) -> bool:
+    """Heuristic: few enough repeated labels that a category is the better fit.
+
+    Mirrors the usual pandas rule of thumb -- a column whose distinct values
+    number well under half its rows is almost always a label, not free text.
+    Callers can always override explicitly.
+    """
+    n_distinct = len(set(values))
+    return 0 < n_distinct <= min(1000, max(1, n_rows // 2))
+
+
+def _write_column(
+    group: Any,
+    name: str,
+    values: List[str],
+    n_rows: int,
+    categorical: bool,
+) -> None:
+    """Write one CSV column, choosing the narrowest faithful encoding."""
+    try:
+        write_dense(group, name, np.array(values, dtype=np.int64))
+        return
+    except (ValueError, TypeError, OverflowError):
+        pass
+
+    try:
+        write_dense(group, name, np.array(values, dtype=np.float64))
+        return
+    except (ValueError, TypeError):
+        pass
+
+    if categorical:
+        categories = sorted(set(values))
+        lookup = {c: i for i, c in enumerate(categories)}
+        write_categorical(
+            group, name, [lookup[v] for v in values], categories, ordered=False
+        )
+        return
+
+    write_string_array(group, name, values)
+
+
+def import_dataframe(
+    root: Any,
+    obj: str,
+    input_file: Path,
+    index_column: Optional[str],
+    console: Console,
+    categorical: Optional[List[str]] = None,
+    auto_categorical: bool = True,
+) -> None:
+    """Replace a dataframe at `obj` with the contents of a CSV file.
+
+    `obj` is any dataframe path, not only "obs" or "var". Columns that parse
+    cleanly as integers or floats become numeric arrays; the rest become
+    either a `categorical` or a `string-array`. `categorical` names columns to
+    force, and `auto_categorical` applies the heuristic to the remainder.
+    """
+    obj = norm_path(obj)
+    rows, data_columns, index_values, _ = _read_csv(input_file, index_column)
+    n_rows = len(rows)
+
+    validate_dimensions(root, obj, (n_rows,), console)
+
+    parts = obj.split("/")
+    parent = root
+    for part in parts[:-1]:
+        parent = parent[part] if part in parent else write_mapping(parent, part)
+
+    group = write_dataframe_header(
+        parent, parts[-1], index_values, data_columns, index_name="_index"
+    )
+
+    forced = set(categorical or ())
+    n_categorical = 0
+    for col in data_columns:
+        values = [row[col] for row in rows]
+        as_categorical = col in forced or (
+            auto_categorical and _looks_categorical(values, n_rows)
+        )
+        _write_column(group, col, values, n_rows, as_categorical)
+        if as_categorical and not _is_numeric(group, col):
+            n_categorical += 1
+
+    detail = f" ({n_categorical} categorical)" if n_categorical else ""
+    console.print(
+        f"[green]Imported[/] {n_rows} rows x {len(data_columns)} columns "
+        f"into '{obj}'{detail}"
+    )
+
+
+def _is_numeric(group: Any, col: str) -> bool:
+    """True when the written column ended up as a numeric array."""
+    from adata.elements import spec
+
+    return spec.encoding_type(group[col]) == spec.ARRAY