pyadata-cli 0.5.0__py3-none-any.whl

adata/elements/__init__.py

@@ -0,0 +1,10 @@
+"""AnnData on-disk element handling: the spec, and reading/writing each element.
+
+`spec` holds the encoding constants, `read` turns any on-disk layout into
+usable values, `write` emits the current spec, and `strings` reconciles the two
+backends' incompatible spellings of text.
+"""
+
+from adata.elements import read, spec, strings, write
+
+__all__ = ["read", "spec", "strings", "write"]

adata/elements/read.py

@@ -0,0 +1,310 @@
+"""Reading AnnData elements, in every layout the spec has ever used.
+
+Elements that logically hold one column of text are stored four different ways
+depending on which anndata wrote the file:
+
+* a plain dataset of bytes or vlen str (`string-array`, or untagged in 0.7.x),
+* a group of `categories` + `codes` (`categorical`),
+* a dataset of codes plus a `categories` attribute (legacy 0.7.x categorical),
+* a group of `values` + `mask` (`nullable-string-array`, what anndata >= 0.11
+  writes for every string column *including the dataframe index*).
+
+The last case is why this module exists: code that assumed an index was always
+a dataset fails outright on any recent file.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Dict, List, Optional, Tuple
+
+import numpy as np
+
+from adata.elements import spec
+from adata.storage import is_dataset, is_group, is_hdf5_dataset
+
+
+def decode_str_array(array: np.ndarray) -> np.ndarray:
+    """Decode an array of bytes/objects to a numpy array of str."""
+    arr = np.asarray(array)
+
+    if np.issubdtype(arr.dtype, np.bytes_):
+        flat = arr.reshape(-1)
+        decoded = [
+            (
+                b.decode("utf-8", errors="replace")
+                if isinstance(b, (bytes, np.bytes_))
+                else str(b)
+            )
+            for b in flat
+        ]
+        return np.asarray(decoded, dtype=str).reshape(arr.shape)
+
+    # "T" is numpy's StringDType, which zarr-python 3 reports for both
+    # VariableLengthUTF8 (v3) and VLenUTF8 (v2); it refuses a direct astype(str).
+    if arr.dtype.kind in ("O", "T"):
+        flat = arr.reshape(-1)
+        decoded = [
+            (
+                v.decode("utf-8", errors="replace")
+                if isinstance(v, (bytes, np.bytes_))
+                else str(v)
+            )
+            for v in flat
+        ]
+        return np.asarray(decoded, dtype=str).reshape(arr.shape)
+
+    return arr.astype(str)
+
+
+def element_len(obj: Any) -> int:
+    """Length along the first axis, whether `obj` is a dataset or a group.
+
+    Handles the masked (`values`/`mask`) and categorical group layouts, so
+    callers never need to know how a column happens to be stored.
+    """
+    if is_dataset(obj):
+        shape = getattr(obj, "shape", None)
+        if not shape:
+            raise ValueError("Element is a scalar and has no length.")
+        return int(shape[0])
+
+    if is_group(obj):
+        enc = spec.encoding_type(obj)
+        if enc in spec.MASKED_TYPES or "values" in obj:
+            return int(obj["values"].shape[0])
+        if enc == spec.CATEGORICAL or "codes" in obj:
+            return int(obj["codes"].shape[0])
+        raise TypeError(
+            f"Cannot determine length of group with encoding {enc!r}: "
+            "expected 'values' or 'codes'."
+        )
+
+    raise TypeError(f"Unsupported element type: {type(obj)}")
+
+
+def resolve_index(group: Any, axis: Optional[str] = None) -> Tuple[Any, str]:
+    """Return the ``(index_element, index_name)`` of a dataframe group.
+
+    Honours the declared ``_index`` attribute first and only then falls back to
+    the ``obs_names``/``var_names`` convention -- the reverse of that order
+    picks the wrong column on a file carrying both.
+    """
+    index_name = spec.decode_attr(group.attrs.get("_index", None))
+
+    candidates: List[str] = []
+    if index_name:
+        candidates.append(str(index_name))
+    if axis == "obs":
+        candidates.append("obs_names")
+    elif axis == "var":
+        candidates.append("var_names")
+    else:
+        candidates.extend(("obs_names", "var_names"))
+    candidates.append("_index")
+
+    for name in candidates:
+        if name in group:
+            return group[name], name
+
+    raise KeyError(
+        f"Could not find an index in group {getattr(group, 'name', '?')!r}; "
+        f"tried {candidates}."
+    )
+
+
+def dataframe_columns(group: Any, index_name: Optional[str] = None) -> List[str]:
+    """Return a dataframe group's data columns in their authored order.
+
+    The spec records original column order in the `column-order` attribute;
+    without it the order would be whatever the backend enumerates, which is
+    alphabetical on HDF5. Columns present on disk but absent from the
+    attribute are appended so nothing is ever dropped.
+    """
+    if index_name is None:
+        try:
+            _, index_name = resolve_index(group)
+        except KeyError:
+            index_name = None
+
+    skip = set(spec.RESERVED_DATAFRAME_KEYS)
+    if index_name:
+        skip.add(index_name)
+
+    present = [k for k in group.keys() if k not in skip]
+
+    declared = group.attrs.get("column-order", None)
+    if declared is None:
+        return present
+
+    ordered = [
+        str(spec.decode_attr(c))
+        for c in np.asarray(declared).reshape(-1).tolist()
+    ]
+    seen = set()
+    out = [c for c in ordered if c in present and not (c in seen or seen.add(c))]
+    out.extend(c for c in present if c not in seen)
+    return out
+
+
+def _categorical_cache_key(col: Any, parent_group: Any | None = None) -> str:
+    col_name = getattr(col, "name", None)
+    if isinstance(col_name, str) and col_name:
+        return col_name
+
+    if parent_group is not None:
+        parent_name = getattr(parent_group, "name", "")
+        rel_name = getattr(col, "path", "")
+        if parent_name or rel_name:
+            return f"{parent_name}/{rel_name}"
+
+    return repr(col)
+
+
+def read_categories(col: Any, parent_group: Any | None = None) -> np.ndarray:
+    """Read the category labels of a categorical column, modern or legacy."""
+    if is_group(col):
+        return np.asarray(decode_str_array(col["categories"][...]), dtype=str)
+
+    cats_ref = col.attrs.get("categories", None)
+    if cats_ref is not None and is_hdf5_dataset(col):
+        # Legacy 0.7.x HDF5: an object reference into __categories.
+        return np.asarray(decode_str_array(col.file[cats_ref][...]), dtype=str)
+
+    if parent_group is not None and "__categories" in parent_group:
+        col_name = str(getattr(col, "name", "")).split("/")[-1]
+        cats_grp = parent_group["__categories"]
+        if col_name in cats_grp:
+            return np.asarray(decode_str_array(cats_grp[col_name][...]), dtype=str)
+
+    if cats_ref is not None and isinstance(cats_ref, (str, bytes)):
+        # Legacy Zarr: `categories` is an absolute path within the store.
+        path = spec.decode_attr(cats_ref).lstrip("/")
+        root = getattr(col, "store_path", None)
+        if parent_group is not None and path.split("/")[-1] in parent_group:
+            return np.asarray(
+                decode_str_array(parent_group[path.split("/")[-1]][...]), dtype=str
+            )
+        del root
+
+    raise KeyError(
+        f"Cannot find categories for categorical column "
+        f"{getattr(col, 'name', '?')!r}."
+    )
+
+
+def is_ordered(col: Any) -> bool:
+    """Whether a categorical column declares an order over its categories."""
+    return bool(spec.decode_attr(col.attrs.get("ordered", False)))
+
+
+def read_categorical_column(
+    col: Any,
+    start: int,
+    end: int,
+    cache: Dict[str, np.ndarray],
+    parent_group: Any | None = None,
+) -> List[str]:
+    """Read rows [start, end) of a categorical column as strings.
+
+    Codes outside the category range -- notably -1 -- render as empty, which is
+    how the spec denotes a missing value.
+    """
+    key = _categorical_cache_key(col, parent_group)
+    if key not in cache:
+        cache[key] = read_categories(col, parent_group)
+    cats = cache[key]
+
+    codes_ds = col["codes"] if is_group(col) else col
+    codes = np.asarray(codes_ds[start:end], dtype=np.int64)
+    return [cats[c] if 0 <= c < len(cats) else "" for c in codes]
+
+
+def _format_values(values: np.ndarray) -> List[str]:
+    """Render a numeric or text array as strings without a lossy detour."""
+    arr = np.asarray(values)
+    if arr.dtype.kind in ("O", "S", "U", "T"):
+        return decode_str_array(arr).tolist()
+    return [str(v) for v in arr.tolist()]
+
+
+def read_masked_column(
+    col: Any, start: int, end: int, na_repr: str = ""
+) -> List[str]:
+    """Read rows [start, end) of a `values`/`mask` group as strings.
+
+    A True mask marks a missing entry, rendered as `na_repr`.
+    """
+    values = col["values"][start:end]
+    mask = np.asarray(col["mask"][start:end], dtype=bool)
+    rendered = _format_values(values)
+    return [na_repr if m else v for v, m in zip(rendered, mask)]
+
+
+def read_str_chunk(
+    obj: Any,
+    start: int,
+    end: int,
+    cache: Optional[Dict[str, np.ndarray]] = None,
+    parent_group: Any | None = None,
+    na_repr: str = "",
+) -> List[str]:
+    """Read rows [start, end) of any column-like element as strings.
+
+    This is the one entry point for turning a column into text, whatever
+    layout it uses on disk.
+    """
+    if cache is None:
+        cache = {}
+
+    if is_dataset(obj):
+        if "categories" in obj.attrs:
+            return read_categorical_column(obj, start, end, cache, parent_group)
+        chunk = obj[start:end]
+        arr = np.asarray(chunk)
+        if arr.ndim != 1:
+            arr = arr.reshape(-1)
+        return _format_values(arr)
+
+    if is_group(obj):
+        enc = spec.encoding_type(obj)
+
+        if enc == spec.CATEGORICAL or (enc is None and "codes" in obj):
+            return read_categorical_column(obj, start, end, cache, parent_group)
+
+        if enc in spec.MASKED_TYPES or (enc is None and "values" in obj and "mask" in obj):
+            return read_masked_column(obj, start, end, na_repr=na_repr)
+
+        raise ValueError(
+            f"Unsupported group encoding {enc!r} for element "
+            f"{getattr(obj, 'name', '?')!r}."
+        )
+
+    raise TypeError(f"Unsupported element type: {type(obj)}")
+
+
+def read_str_all(obj: Any, chunk_size: int = 200_000, **kwargs: Any) -> List[str]:
+    """Read an entire column-like element as strings, in chunks."""
+    n = element_len(obj)
+    cache: Dict[str, np.ndarray] = {}
+    out: List[str] = []
+    for start in range(0, n, chunk_size):
+        out.extend(
+            read_str_chunk(obj, start, min(start + chunk_size, n), cache, **kwargs)
+        )
+    return out
+
+
+def col_chunk_as_strings(
+    group: Any,
+    col_name: str,
+    start: int,
+    end: int,
+    cat_cache: Dict[str, np.ndarray],
+) -> List[str]:
+    """Read rows [start, end) of a named column within a dataframe group."""
+    if col_name not in group:
+        raise RuntimeError(
+            f"Column {col_name!r} not found in group "
+            f"{getattr(group, 'name', '?')!r}"
+        )
+    return read_str_chunk(group[col_name], start, end, cat_cache, parent_group=group)

adata/elements/spec.py

@@ -0,0 +1,125 @@
+"""AnnData on-disk encoding constants and attribute access.
+
+This is the single source of truth for `encoding-type` / `encoding-version`
+pairs. Every read that inspects an encoding and every write that stamps one
+goes through here, so the spec version the tool targets is stated in exactly
+one place.
+
+See docs/ELEMENTS_h5ad.md and docs/ELEMENTS_zarr.md for the full spec.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Dict, Optional, Tuple
+
+ANNDATA = "anndata"
+RAW = "raw"
+DICT = "dict"
+DATAFRAME = "dataframe"
+ARRAY = "array"
+STRING_ARRAY = "string-array"
+CATEGORICAL = "categorical"
+CSR_MATRIX = "csr_matrix"
+CSC_MATRIX = "csc_matrix"
+NULLABLE_INTEGER = "nullable-integer"
+NULLABLE_BOOLEAN = "nullable-boolean"
+NULLABLE_STRING_ARRAY = "nullable-string-array"
+NUMERIC_SCALAR = "numeric-scalar"
+STRING = "string"
+AWKWARD_ARRAY = "awkward-array"
+NULL = "null"
+
+#: Encoding type -> the version this tool writes.
+CURRENT_VERSION = {
+    ANNDATA: "0.1.0",
+    RAW: "0.1.0",
+    DICT: "0.1.0",
+    DATAFRAME: "0.2.0",
+    ARRAY: "0.2.0",
+    STRING_ARRAY: "0.2.0",
+    CATEGORICAL: "0.2.0",
+    CSR_MATRIX: "0.1.0",
+    CSC_MATRIX: "0.1.0",
+    NULLABLE_INTEGER: "0.1.0",
+    NULLABLE_BOOLEAN: "0.1.0",
+    NULLABLE_STRING_ARRAY: "0.1.0",
+    NUMERIC_SCALAR: "0.2.0",
+    STRING: "0.2.0",
+    AWKWARD_ARRAY: "0.1.0",
+    NULL: "0.1.0",
+}
+
+SPARSE_TYPES = frozenset({CSR_MATRIX, CSC_MATRIX})
+NULLABLE_TYPES = frozenset(
+    {NULLABLE_INTEGER, NULLABLE_BOOLEAN, NULLABLE_STRING_ARRAY}
+)
+#: Encodings stored as a group holding `values` and `mask`.
+MASKED_TYPES = NULLABLE_TYPES
+#: Encodings that yield text when read as strings.
+STRING_TYPES = frozenset({STRING_ARRAY, NULLABLE_STRING_ARRAY, STRING, CATEGORICAL})
+#: Group encodings that are plain mappings of further elements.
+MAPPING_TYPES = frozenset({DICT, ANNDATA, RAW})
+
+#: Keys that are structural rather than dataframe columns.
+RESERVED_DATAFRAME_KEYS = frozenset({"_index", "__categories"})
+
+
+def decode_attr(value: Any) -> Any:
+    """Decode a raw attribute value to a plain Python value.
+
+    HDF5 hands back `bytes` for string attributes while Zarr hands back `str`;
+    numpy scalars appear on both. Normalising here is what lets the rest of the
+    codebase compare attributes with `==`.
+    """
+    if isinstance(value, bytes):
+        return value.decode("utf-8")
+    if hasattr(value, "item") and getattr(value, "shape", None) == ():
+        return value.item()
+    return value
+
+
+def encoding_of(obj: Any) -> Tuple[Optional[str], Optional[str]]:
+    """Return the declared ``(encoding-type, encoding-version)`` of an element.
+
+    Either component is None when the element is untagged, which is the normal
+    case for files written by anndata 0.7.x and earlier.
+    """
+    attrs = getattr(obj, "attrs", {})
+    enc = decode_attr(attrs.get("encoding-type", None))
+    ver = decode_attr(attrs.get("encoding-version", None))
+    return (enc or None, ver or None)
+
+
+def encoding_type(obj: Any) -> Optional[str]:
+    """Return just the declared ``encoding-type``, or None if untagged."""
+    return encoding_of(obj)[0]
+
+
+def set_attrs(obj: Any, values: Dict[str, Any]) -> None:
+    """Set several attributes in as few writes as the backend allows.
+
+    On Zarr each assignment persists the whole metadata document through its
+    sync-over-async bridge, so setting attributes one at a time costs a round
+    trip each. HDF5 has no such batch API and falls back to assignment.
+    """
+    if not values:
+        return
+
+    put = getattr(obj.attrs, "put", None)
+    if callable(put):
+        put({**dict(obj.attrs), **values})
+        return
+
+    for key, value in values.items():
+        obj.attrs[key] = value
+
+
+def set_encoding(obj: Any, enc_type: str, version: Optional[str] = None) -> None:
+    """Stamp ``encoding-type``/``encoding-version`` on a group or dataset.
+
+    The version defaults to the current spec version for that type, so callers
+    never hardcode one.
+    """
+    if version is None:
+        version = CURRENT_VERSION[enc_type]
+    set_attrs(obj, {"encoding-type": enc_type, "encoding-version": version})

adata/elements/strings.py

@@ -0,0 +1,90 @@
+"""Backend-neutral handling of string elements.
+
+The AnnData spec requires `string-array` to be variable-length UTF-8: a vlen
+str dtype in HDF5, `VariableLengthUTF8` in Zarr. Neither backend accepts the
+other's spelling, and a dtype round-tripped verbatim across backends fails
+outright -- an h5py vlen dataset reports `dtype == object`, which Zarr rejects
+as ambiguous. Everything that creates or copies a string element therefore
+resolves the dtype through this module rather than reusing the source's.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Optional
+
+import h5py
+import numpy as np
+
+from adata.storage import is_dataset, is_zarr_array
+
+
+def is_string_dtype(dtype: Any) -> bool:
+    """True if `dtype` holds text, in any spelling either backend produces.
+
+    Covers h5py vlen-UTF-8 (object), fixed-width bytes ('S'), numpy unicode
+    ('U'), numpy StringDType ('T', what zarr-python 3 reports for both
+    `VariableLengthUTF8` and v2 `VLenUTF8`), and bare object arrays.
+    """
+    if dtype is None:
+        return False
+    if dtype is str:
+        return True
+    try:
+        if h5py.check_string_dtype(dtype) is not None:
+            return True
+    except (TypeError, AttributeError):
+        pass
+    kind = getattr(dtype, "kind", None)
+    return kind in ("O", "S", "U", "T")
+
+
+def is_string_element(obj: Any) -> bool:
+    """True if `obj` is a dataset/array holding text."""
+    return is_dataset(obj) and is_string_dtype(getattr(obj, "dtype", None))
+
+
+def string_dtype_for(backend: str, zarr_format: Optional[int] = None) -> Any:
+    """Return the spec-compliant variable-length UTF-8 dtype for `backend`.
+
+    `zarr_format` is accepted for symmetry; zarr-python 3 maps `str` to
+    `VariableLengthUTF8` under both v2 and v3, so it does not currently affect
+    the result.
+    """
+    if backend == "zarr":
+        return str
+    return h5py.string_dtype(encoding="utf-8")
+
+
+def target_dtype(src_dtype: Any, target_backend: str, zarr_format: Optional[int] = None) -> Any:
+    """Map a source dtype onto one the target backend can actually create.
+
+    String dtypes are normalised to variable-length UTF-8; everything else is
+    passed through unchanged.
+    """
+    if is_string_dtype(src_dtype):
+        return string_dtype_for(target_backend, zarr_format)
+    return src_dtype
+
+
+def as_str_array(values: Any) -> np.ndarray:
+    """Coerce a sequence to a 1-D numpy array of Python str, decoding bytes."""
+    arr = np.asarray(values, dtype=object)
+    flat = [
+        v.decode("utf-8", errors="replace") if isinstance(v, (bytes, np.bytes_)) else str(v)
+        for v in arr.reshape(-1)
+    ]
+    return np.asarray(flat, dtype=object).reshape(arr.shape)
+
+
+def strip_string_filters(kwargs: dict, target_backend: str, zarr_format: Optional[int]) -> dict:
+    """Drop codecs that cannot survive the jump to `target_backend`.
+
+    A Zarr v2 string array carries `VLenUTF8` in `filters`; forwarding it to a
+    v3 array raises `TypeError: Expected an ArrayArrayCodec`. The v3 string
+    dtype encodes variable length itself, so the filter is simply dropped.
+    """
+    if target_backend != "zarr" or zarr_format != 3:
+        return kwargs
+    out = dict(kwargs)
+    out.pop("filters", None)
+    return out