diff --git a/docs/guides/inspect-export.md b/docs/guides/inspect-export.md index faae969..183df85 100644 --- a/docs/guides/inspect-export.md +++ b/docs/guides/inspect-export.md @@ -64,7 +64,10 @@ missing, and text as fixed-width unicode. ## Metadata and images `export dict` writes a group or scalar as JSON, and `export image` writes a -2-D or 3-D array as PNG, JPEG or TIFF, such as a Visium tissue image: +2-D or 3-D array as PNG, JPEG or TIFF, such as a Visium tissue image. It +keeps the values, a 16-bit image as a 16-bit PNG or TIFF, and refuses ones +an image cannot hold unless `--scale` stretches them onto 0-255 (see +[`export`](../reference/commands.md#export)): ```bash adata export dict data.h5ad uns/neighbors -o neighbors.json diff --git a/docs/reference/commands.md b/docs/reference/commands.md index f8af9a9..770aee7 100644 --- a/docs/reference/commands.md +++ b/docs/reference/commands.md @@ -460,6 +460,20 @@ with `export dataframe`, which writes them as empty fields. `export dict` writes NaN and infinite values as `null`, since JSON has no number for them, and says how many it wrote. +`export image` keeps the values as they are. Whole numbers from 0 to 255 +are an 8-bit image, whatever their dtype. One channel of `uint16`, or of +whole numbers up to 65,535, is a 16-bit image as `.png` or `.tiff`, the +formats that keep 16 bits; a colour image keeps 8 bits a channel. Floats +from 0 to 1 are a normalised image, scaled up to 0-255; booleans become 0 +and 255. Values an image cannot keep (negative, too large for the format, +not finite) are refused, naming their range, rather than clipped. +`--scale` stretches them onto 0-255 instead, the smallest to 0 and the +largest to 255, with anything not finite at 0: + +```bash +adata export image data.h5ad uns/spatial/morphology -o morphology.png --scale +``` + A file output is written beside its destination and renamed into place once complete, so an export that fails or is refused leaves an existing file as it was. An output that names the input store, lies inside a Zarr input, or names diff --git a/docs/reference/formats/cistopic.md b/docs/reference/formats/cistopic.md index bd230db..57d8c3c 100644 --- a/docs/reference/formats/cistopic.md +++ b/docs/reference/formats/cistopic.md @@ -29,20 +29,51 @@ adata convert imputed_acc_obj.pkl -o imputed.h5ad --dense | `projections["cell"]` and `["region"]`, such as UMAP | `obsm["X_"]`, `varm["X_"]` | | `project`, `path_to_fragments`, the topic names, the model's metrics, coherence, parameters and other tables | `uns["pycistopic"]` | -A `CistopicImputedFeatures` becomes X, cells x features, CSR (`--dense` -keeps it dense), with `feature_names` and `cell_names` as var and obs names. - -The pickle is read without running anything from it: an unpickler that -resolves only the names a pycisTopic object holds (numpy arrays, scipy -sparse matrices, pandas frames, a few builtins) and refuses any other by -name, so a pickle naming `os.system` is refused rather than run. -pycisTopic's own classes are stood in for by plain holders of their -attributes. A pickle written by pycisTopic's pandas 1.5 reads with pandas 2 -or 3, as CI checks. Columns are written as anndata writes them: Categorical -as categorical, nullable integers and booleans as nullable, numbers held as +A `CistopicImputedFeatures` becomes X, cells x features, CSR, or dense with +`--dense`, with `feature_names` and `cell_names` as var and obs names. Its +matrix is dense for imputed accessibility and sparse for gene activity; +`--dense` makes X dense either way. A matrix that is not numbers is refused. + +## Pickles are trusted input + +Read a pickle only from a source you would trust with the data itself. It +is read without running anything from it: an unpickler that resolves only +the names a pycisTopic object holds (numpy arrays, scipy sparse matrices, +pandas frames, a few builtins) and refuses any other by name, so a pickle +naming `os.system` is refused rather than run. pycisTopic's own classes are +stood in for by plain holders of their attributes. + +What those names are given is checked too, before numpy, scipy or pandas +act on it, and a pickle that fails a check is refused before anything is +written: + +- every numpy dtype and array is one numpy itself would write: a plain + dtype (no structured ones), and data that fills the array's shape + exactly; +- pandas' and scipy's classes are created empty and given their state, as + every pickle of them does, never constructed from the file's arguments; +- every sparse matrix has its indices within its shape and an `indptr` + that rises from 0 to its number of entries, every categorical's codes + lie within its categories, and every table's columns hold as many rows + as its index. + +Dates and times are not read; pycisTopic's objects hold none. + +A pickle written by pycisTopic's pandas 1.5 reads with pandas 2 or 3, as CI +checks. Columns are written as anndata writes them: Categorical as +categorical, nullable integers and booleans as nullable, numbers held as `object` as numbers, repeated text as categorical. Anything else on the object is named in a warning and left out. +## Memory + A pickle cannot be read in part, so it is held in memory whole, as it was -when pycisTopic saved it; the fragment matrix is turned into CSR on disk by -the streaming transpose rather than copied in memory. +when pycisTopic saved it. Reading it checks each value, which takes about +ten times as long as Python's own `pickle.load`: under a second for a +200 MB pickle. The +fragment matrix is turned into CSR on disk by the streaming transpose +rather than copied in memory. An imputed matrix is written a block of about +a million entries at a time (a block of cells, or of one cell's features +when a cell has more), whatever the number of features, and a dense X is +chunked the same way. A sparse matrix written with `--dense` goes through +CSR on disk first, then a block of cells at a time. diff --git a/docs/reference/formats/xenium.md b/docs/reference/formats/xenium.md index d31248d..98f8ffa 100644 --- a/docs/reference/formats/xenium.md +++ b/docs/reference/formats/xenium.md @@ -30,7 +30,9 @@ adata convert xenium_bundle -o region1.h5ad --feature-type "Gene Expression" Xenium ships `analysis/` packed as `analysis.tar.gz`. When it has not been unpacked, the three tables above are copied out of it, by name, into a -temporary folder beside the output, removed when the conversion ends. A cell an +temporary folder beside the output (in the nearest of its folders that +exists, when the output goes into a new one), removed when the conversion +ends. A cell an analysis table lacks is left empty (no cluster, NaN), with a warning giving the count. diff --git a/src/adata/cli.py b/src/adata/cli.py index 93281ef..b97532c 100644 --- a/src/adata/cli.py +++ b/src/adata/cli.py @@ -1445,21 +1445,30 @@ def export_image( "-o", help="Output image file (.png, .jpg, .tiff); defaults to PNG on stdout", ), + scale: bool = typer.Option( + False, + "--scale", + help="Stretch the values onto 0-255 (smallest to 0, largest to 255), " + "for values an image cannot keep as they are", + ), ) -> None: """ Export an image-like array to PNG/JPG/TIFF format. The array should be 2D (H,W) or 3D (H,W,C) with C in {1,3,4}. The format follows the output's extension; without --output the image goes - to stdout as PNG. + to stdout as PNG. Values are kept as they are: one channel of 16 bits + as a 16-bit PNG or TIFF. Values an image cannot keep are refused, not + clipped; --scale stretches them onto 0-255. Examples: adata export image data.h5ad uns/spatial/image -o tissue.png adata export image data.h5ad uns/spatial/image > tissue.png + adata export image data.h5ad uns/spatial/morphology -o morphology.png --scale """ try: - export_image_cmd(file=file, obj=entry, out=output, console=console) + export_image_cmd(file=file, obj=entry, out=output, console=console, scale=scale) except Exception as e: console.print(f"[bold red]Error:[/] {escape(str(e))}") raise typer.Exit(code=1) diff --git a/src/adata/commands/export.py b/src/adata/commands/export.py index e3badac..67d1bea 100644 --- a/src/adata/commands/export.py +++ b/src/adata/commands/export.py @@ -128,9 +128,11 @@ def export_json( ) -def export_image(file: Path, obj: str, out: Optional[Path], console: Console) -> None: +def export_image( + file: Path, obj: str, out: Optional[Path], console: Console, scale: bool = False +) -> None: with _output(file, out, console) as target, open_store(file, "r") as store: - export_image_format(store.root, obj=obj, out=target, console=console) + export_image_format(store.root, obj=obj, out=target, console=console, scale=scale) __all__ = [ diff --git a/src/adata/conversion/cistopic.py b/src/adata/conversion/cistopic.py index 6a15229..d303476 100644 --- a/src/adata/conversion/cistopic.py +++ b/src/adata/conversion/cistopic.py @@ -24,14 +24,17 @@ A pickle cannot be read in part, so the object is held in memory, as it was when pycisTopic saved it. What is written from it is not copied again: the fragment matrix, regions x cells CSR, is the same three arrays as a CSC cells -x regions, and is turned into CSR on disk by the streaming transpose. +x regions, and is turned into CSR on disk by the streaming transpose. A +dense imputed matrix is written `BLOCK_ENTRIES` at a time, which is also +the dense output's chunk, so neither grows with the number of features; a +sparse one asked for dense goes through that CSR on disk first. """ from __future__ import annotations from dataclasses import dataclass, field from pathlib import Path -from typing import Any, Dict, List, Optional, Sequence, Tuple +from typing import Any, Dict, Iterator, List, Optional, Sequence, Tuple import numpy as np from rich.console import Console @@ -69,8 +72,11 @@ "topic_region", }) -#: Cells of a dense imputed matrix written at once. -DENSE_CELLS = 4096 +#: Entries of an imputed matrix written at once: a block of cells, every +#: feature of each, or a run of one cell's features when a cell has more. +#: Also the chunk of a dense output. Sparse output from a dense matrix takes +#: about 40 bytes an entry while a block is sorted out, so 40 MiB at most. +BLOCK_ENTRIES = 1 << 20 def require_extra() -> None: @@ -130,8 +136,10 @@ def plan( name = Path(source).name try: obj = pickled.load(source) - except pickled.RefusedGlobal as exc: + except pickled.RefusedPickle as exc: raise ValueError(f"{name}: {exc}") from None + except Exception as exc: # numpy's, scipy's or pandas' own refusal, or a file cut short + raise ValueError(f"{name} could not be read as a pycisTopic pickle: {exc}") from None kind = getattr(obj, "kind", type(obj).__name__) warnings: List[str] = [] if kind == "CistopicObject": @@ -294,6 +302,11 @@ def _plan_imputed(obj: Any, name: str) -> _Import: f"{name}: its matrix is {shape}, but it names {len(features):,} features " f"and {len(cells):,} cells." ) + if not _is_sparse(matrix): + if not isinstance(matrix, np.ndarray) or matrix.ndim != 2: + raise ValueError(f"{name}: its matrix is a {type(matrix).__name__}, not an array.") + if matrix.dtype.kind not in "biuf": + raise ValueError(f"{name}: its matrix holds {matrix.dtype} values, not numbers.") return _Import("CistopicImputedFeatures", obj, len(cells), len(features), cells, features) @@ -350,35 +363,90 @@ def _write_csr(parent: Any, name: str, matrix: Any, shape: Tuple[int, int], deta create_dataset(group, key, data=values, **kw) +def _tiles(n_cells: int, n_features: int) -> Tuple[int, int]: + """Cells and features a block holds, `BLOCK_ENTRIES` at most: whole + cells when one fits, else a run of one cell's features.""" + budget = max(1, BLOCK_ENTRIES) + if n_features <= budget: + return max(1, budget // max(1, n_features)), max(1, n_features) + return 1, budget + + +def _blocks(n_cells: int, n_features: int) -> Iterator[Tuple[slice, slice]]: + """(cells, features) of each block, in row order: a cell is split only + when it is a block of its own.""" + rows, cols = _tiles(n_cells, n_features) + for lo in range(0, n_cells, rows): + for start in range(0, n_features, cols): + yield slice(lo, min(lo + rows, n_cells)), slice(start, min(start + cols, n_features)) + + +def _dense_target(root: Any, detail: _Import, dtype: np.dtype) -> Any: + """X, dense, chunked as the blocks it is written in.""" + shape = (detail.n_obs, detail.n_var) + kw: Dict[str, Any] = {} + if all(shape): # no chunk is larger than the array, nor compressed without one + rows, cols = _tiles(*shape) + kw = {"chunks": (min(rows, shape[0]), min(cols, shape[1]))} + if detail.compression is not None: + kw = with_compression( + kw, detail.compression, shape=shape, dtype=dtype, + backend="zarr" if is_zarr_group(root) else "hdf5", zarr_format=detail.zarr_format, + ) + target = create_dataset(root, "X", shape=shape, dtype=dtype, **kw) + spec.set_encoding(target, spec.ARRAY) + return target + + def _write_imputed(root: Any, detail: _Import) -> None: - """The features x cells matrix as cells x features, a block of cells at a time.""" + """The features x cells matrix as cells x features, a block at a time.""" matrix = detail.obj.mtx if _is_sparse(matrix): - _write_fragments(root, "X", matrix, detail) + if detail.dense: + _write_dense_from_sparse(root, matrix, detail) + else: + _write_fragments(root, "X", matrix, detail) return - matrix = np.asarray(matrix) shape = (detail.n_obs, detail.n_var) dtype = np.dtype(matrix.dtype.str) if detail.dense: - kw: Dict[str, Any] = {"chunks": (min(max(shape[0], 1), DENSE_CELLS), max(shape[1], 1))} - if detail.compression is not None: - kw = with_compression( - kw, detail.compression, shape=shape, dtype=dtype, - backend="zarr" if is_zarr_group(root) else "hdf5", zarr_format=detail.zarr_format, - ) - target = create_dataset(root, "X", shape=shape, dtype=dtype, **kw) - spec.set_encoding(target, spec.ARRAY) - for lo in range(0, shape[0], DENSE_CELLS): - target[lo : lo + DENSE_CELLS] = matrix[:, lo : lo + DENSE_CELLS].T + target = _dense_target(root, detail, dtype) + for cells, features in _blocks(*shape): + target[cells, features] = matrix[features, cells].T return appender = CsrAppender(root, "X", shape=shape, data_dtype=dtype, compression=detail.compression) - for lo in range(0, shape[0], DENSE_CELLS): - block = matrix[:, lo : lo + DENSE_CELLS].T + for cells, features in _blocks(*shape): + block = matrix[features, cells].T rows, cols = np.nonzero(block) - appender.append(rows + lo, cols, block[rows, cols]) + appender.append(rows + cells.start, cols + features.start, block[rows, cols]) appender.finish() +def _write_dense_from_sparse(root: Any, matrix: Any, detail: _Import) -> None: + """A sparse features x cells matrix as dense cells x features: as CSR + first, on disk through the streaming transpose, then a block of cells at + a time from that, entries repeated at a place summed as scipy sums + them.""" + dtype = np.dtype(matrix.dtype.str) + with scratch_group(root, "__cistopic_dense__") as scratch: + _write_fragments(scratch, "csr", matrix, detail) + csr = scratch["csr"] + indptr = np.asarray(csr["indptr"][...], dtype=np.int64) + target = _dense_target(root, detail, dtype) + read = None # the cells whose entries are held: one read for all their runs of features + for cells, features in _blocks(detail.n_obs, detail.n_var): + if read != cells: + read = cells + start, end = int(indptr[cells.start]), int(indptr[cells.stop]) + cols = np.asarray(csr["indices"][start:end], dtype=np.int64) + values = np.asarray(csr["data"][start:end]) + rows = np.repeat(np.arange(cells.stop - cells.start), np.diff(indptr[cells.start : cells.stop + 1])) + block = np.zeros((cells.stop - cells.start, features.stop - features.start), dtype=dtype) + keep = (cols >= features.start) & (cols < features.stop) + np.add.at(block, (rows[keep], cols[keep] - features.start), values[keep]) + target[cells, features] = block + + def _write_uns(uns: Any, detail: _Import, warnings: List[str]) -> None: from adata.formats.pandas_frame import write_frame diff --git a/src/adata/conversion/xenium.py b/src/adata/conversion/xenium.py index f4a2786..e8bbaff 100644 --- a/src/adata/conversion/xenium.py +++ b/src/adata/conversion/xenium.py @@ -109,15 +109,23 @@ def _bundle(source: Path) -> Path: return bundle -def _unpack_analysis(archive: Path, beside: Path) -> tempfile.TemporaryDirectory: +def _unpack_analysis(archive: Path, output: Path) -> tempfile.TemporaryDirectory: """The analysis tables this reads, out of Xenium's analysis.tar.gz. Only those members, each copied to a path built here rather than the one the archive names, so the archive cannot write anywhere else. Into a folder beside the output rather than the system's temporary one, - which on a cluster node is often too small for a large run's PCA. + which on a cluster node is often too small for a large run's PCA: the + nearest of the output's folders that exists, since this runs while + planning and the output's own is made only once every check has passed. """ - holder = tempfile.TemporaryDirectory(prefix=".adata-xenium-", dir=beside) + beside = Path(output).resolve().parent + while not beside.exists() and beside != beside.parent: + beside = beside.parent + try: + holder = tempfile.TemporaryDirectory(prefix=".adata-xenium-", dir=beside) + except OSError as exc: + raise ValueError(f"Cannot unpack {archive.name} in {beside}: {exc.strerror or exc}.") from None wanted = {str(Path("analysis") / p): p for p in analysis_paths()} with tarfile.open(archive, "r:*") as tar: for member in tar: @@ -156,7 +164,7 @@ def plan( unpacked = None packed_but_empty = False if not analysis.is_dir() and (bundle / "analysis.tar.gz").is_file(): - unpacked = _unpack_analysis(bundle / "analysis.tar.gz", Path(output).resolve().parent) + unpacked = _unpack_analysis(bundle / "analysis.tar.gz", Path(output)) analysis = Path(unpacked.name) packed_but_empty = not any(analysis.rglob("*.csv")) warnings: List[str] = [] diff --git a/src/adata/formats/image.py b/src/adata/formats/image.py index 75b4e97..6150d41 100644 --- a/src/adata/formats/image.py +++ b/src/adata/formats/image.py @@ -15,10 +15,19 @@ from adata.util.path import norm_path -def export_image(root: Any, obj: str, out: Optional[Path], console: Console) -> None: +#: Formats Pillow writes a channel of 16 bits to (mode "I;16"), for an +#: image of one channel; a colour image it writes with 8 bits a channel. +SIXTEEN_BIT = frozenset({"PNG", "TIFF"}) + + +def export_image( + root: Any, obj: str, out: Optional[Path], console: Console, *, scale: bool = False +) -> None: """Write a 2-D or 3-D array as an image, or as PNG to stdout when `out` is None. The format follows the output's extension; stdout has none, so it is PNG. + Values are kept as they are (`pixels`), or with `scale` stretched onto + 0-255. """ h5obj = _resolve(root, obj) if not is_dataset(h5obj): @@ -31,26 +40,13 @@ def export_image(root: Any, obj: str, out: Optional[Path], console: Console) -> raise ValueError( f"Expected last dimension (channels) to be 1, 3, or 4; got {arr.shape}." ) - - if np.issubdtype(arr.dtype, np.floating): - amax = float(np.nanmax(arr)) if arr.size else 0.0 - if amax <= 1.0: - arr = np.clip(arr, 0.0, 1.0) * 255.0 - else: - arr = np.clip(arr, 0.0, 255.0) - arr = arr.astype(np.uint8) - elif np.issubdtype(arr.dtype, np.integer): - arr = np.clip(arr, 0, 255).astype(np.uint8) - elif arr.dtype == np.bool_: - arr = arr.astype(np.uint8) * 255 - else: - raise ValueError(f"Unsupported image dtype: {arr.dtype}") - if arr.ndim == 3 and arr.shape[2] == 1: arr = arr[:, :, 0] - img = Image.fromarray(arr) - if out is None or str(out) == "-": + to_stdout = out is None or str(out) == "-" + fmt = "PNG" if to_stdout else Image.registered_extensions().get(out.suffix.lower()) + img = Image.fromarray(pixels(arr, obj, fmt, scale=scale)) + if to_stdout: img.save(sys.stdout.buffer, format="PNG") sys.stdout.buffer.flush() return @@ -58,6 +54,69 @@ def export_image(root: Any, obj: str, out: Optional[Path], console: Console) -> img.save(out) +def pixels(arr: np.ndarray, obj: str, fmt: Optional[str], *, scale: bool = False) -> np.ndarray: + """`arr` as an image holds it, its values unchanged: whole numbers from + 0 to 255 as 8 bits, and one channel of uint16 (or of whole numbers to + 65535) as 16 bits where `fmt` keeps them (PNG, TIFF). Floats from 0 to 1 + are a normalised image and scale up to 0-255; booleans are 0 and 255. + Other values are refused rather than clipped, unless `scale` stretches + them onto 0-255: the smallest to 0, the largest to 255.""" + if arr.dtype == np.bool_: + return arr.astype(np.uint8) * 255 + if arr.dtype.kind not in "iuf": + raise ValueError(f"Unsupported image dtype: {arr.dtype}") + if scale: + return _stretch(arr) + if arr.dtype == np.uint8: + return arr + deep = arr.ndim == 2 and fmt in SIXTEEN_BIT + if arr.dtype == np.uint16 and deep: + return arr + finite = bool(np.isfinite(arr).all()) if arr.dtype.kind == "f" else True + values = arr[np.isfinite(arr)] if not finite else arr + lo, hi = (values.min(), values.max()) if values.size else (0, 0) + if arr.dtype.kind == "f" and finite and lo >= 0: + if hi <= 1: + return (arr * 255.0).astype(np.uint8) + if hi <= 255: + return arr.astype(np.uint8) + elif arr.dtype.kind in "iu" and lo >= 0: + if hi <= 255: + return arr.astype(np.uint8) + if hi <= 65535 and deep: + return arr.astype(np.uint16) + + if not finite: + why = "some are not finite" + elif arr.dtype.kind in "iu" and lo >= 0 and hi <= 65535 and arr.ndim == 2: + why = f"{fmt or 'that format'} keeps 8 bits a pixel; a .png or .tiff keeps 16" + elif arr.dtype.kind in "iu" and lo >= 0 and hi <= 65535: + why = "a colour image keeps 8 bits a channel (one channel keeps 16, as .png or .tiff)" + else: + top = 65535 if deep else 255 + why = f"an image keeps whole numbers from 0 to {top}" + ( + ", or floats from 0 to 1" if arr.dtype.kind == "f" else "" + ) + raise ValueError( + f"'{obj}' holds {arr.dtype} values from {lo:g} to {hi:g}, and {why}. Pass " + "--scale to stretch them onto 0-255, or export the values as they are with " + "`adata export array`." + ) + + +def _stretch(arr: np.ndarray) -> np.ndarray: + """Values stretched onto 0-255: the smallest finite one to 0, the + largest to 255; any that is not finite to 0.""" + values = arr.astype(np.float64) + finite = np.isfinite(values) + out = np.zeros(arr.shape, dtype=np.float64) + if finite.any(): + lo, hi = values[finite].min(), values[finite].max() + if hi > lo: + out[finite] = (values[finite] - lo) / (hi - lo) * 255.0 + return np.rint(out).astype(np.uint8) + + def import_image(root: Any, obj: str, input_file: Path, console: Console) -> None: """Read an image file into the store as a dense array. diff --git a/src/adata/formats/pickled.py b/src/adata/formats/pickled.py index 1517989..79ae511 100644 --- a/src/adata/formats/pickled.py +++ b/src/adata/formats/pickled.py @@ -7,6 +7,34 @@ by an unpickler that resolves an allowlist of those names and refuses any other: an `os.system` in the file is refused by name, never called. +The allowed names are not harmless whatever they are given, so what they +are given is checked too, before numpy, scipy or pandas act on it: + +* numpy's dtypes and arrays must hold the state numpy itself writes for + them (numpy trusts a pickled dtype's flags and item size, which would let + raw bytes from the file stand for object pointers), and each is given its + state once, before anything else holds it; +* classes are created empty and given their state, as every pickle of them + does; they are never constructed from a file's arguments, so pandas never + parses a dtype or a timezone from it (`dateutil/` opens that path); +* containers are filled only through `dict`, `list` and `set`, never by + writing into an array that has been checked; +* once the whole file is read, every sparse matrix's indices must lie within + its shape and its `indptr` must rise from 0 to its entries, every + Categorical's and MultiIndex's codes within its categories, and every + nullable array's mask must match its values, and every frame's blocks + must fill its rows and hold each column once: scipy's and pandas' + compiled code indexes with them unchecked. + +That keeps a crafted file from corrupting memory or reading files through +what it names, but it is still parsed by numpy, scipy and pandas: a pickle +is trusted input, to be read only from a source one would trust with the +data itself. + +The pure-Python unpickler is used, so that each opcode can be checked; it +is a few times slower than the C one, which is small beside writing the +data. + pycisTopic itself is not needed. Its classes become inert stand-ins that hold the restored attributes (`__dict__`); none of its code exists here to run. numpy, scipy and pandas are resolved to whatever versions are @@ -21,13 +49,17 @@ from __future__ import annotations import _compat_pickle +import copyreg import gzip import importlib import pickle import pickletools +from collections import OrderedDict from contextlib import contextmanager from pathlib import Path -from typing import Any, BinaryIO, Dict, FrozenSet, Iterator, Optional, Tuple +from typing import Any, BinaryIO, Callable, Dict, FrozenSet, Iterable, Iterator, Optional, Tuple + +import numpy as np #: Module prefix of pycisTopic's classes, as its pickles name them. PYCISTOPIC = "pycisTopic." @@ -47,7 +79,9 @@ "_reconstruct", "scalar", "_frombuffer", "ndarray", "dtype", } -#: The globals a pycisTopic pickle may name: module -> names. +#: The globals a pycisTopic pickle may name: module -> names. Dates and +#: times are not among them: pandas reads a timezone given as text from a +#: file named in it, and pycisTopic's objects hold none. ALLOWED: Dict[str, FrozenSet[str]] = { **{module: frozenset(_NUMPY) for module in ( "numpy", "numpy.core.multiarray", "numpy._core.multiarray", @@ -61,44 +95,75 @@ "pandas.core.internals.blocks": frozenset({"new_block", "Block", "ObjectBlock"}), "pandas._libs.internals": frozenset({"_unpickle_block", "BlockPlacement"}), "pandas._libs.arrays": frozenset({"__pyx_unpickle_NDArrayBacked"}), - "pandas.core.arrays.categorical": frozenset({"Categorical", "_unpickle"}), - "pandas.core.dtypes.dtypes": frozenset({ - "CategoricalDtype", "DatetimeTZDtype", "PeriodDtype", "IntervalDtype", - }), + "pandas._libs.missing": frozenset({"NA"}), + "pandas.core.arrays.categorical": frozenset({"Categorical"}), + "pandas.core.dtypes.dtypes": frozenset({"CategoricalDtype"}), "pandas.core.indexes.base": frozenset({"_new_Index", "Index"}), "pandas.core.indexes.range": frozenset({"RangeIndex"}), "pandas.core.indexes.numeric": frozenset({"Int64Index", "Float64Index", "UInt64Index"}), "pandas.core.indexes.category": frozenset({"CategoricalIndex"}), "pandas.core.indexes.multi": frozenset({"MultiIndex"}), - "pandas.core.indexes.datetimes": frozenset({"DatetimeIndex", "_new_DatetimeIndex"}), "pandas.core.arrays.integer": frozenset({"IntegerArray", *( f"{sign}Int{bits}Dtype" for sign in ("", "U") for bits in (8, 16, 32, 64) )}), "pandas.core.arrays.floating": frozenset({"FloatingArray", "Float32Dtype", "Float64Dtype"}), "pandas.core.arrays.boolean": frozenset({"BooleanArray", "BooleanDtype"}), "pandas.core.arrays.string_": frozenset({"StringArray", "StringDtype"}), - "pandas.core.arrays.datetimes": frozenset({"DatetimeArray"}), - "pandas._libs.tslibs.timestamps": frozenset({"_unpickle_timestamp", "Timestamp"}), - "pandas._libs.tslibs.nattype": frozenset({"__nat_unpickle"}), - "pandas._libs.tslibs.timezones": frozenset(), # pandas 3 pickles its classes under their public names. "pandas": frozenset({ "DataFrame", "Series", "Index", "RangeIndex", "MultiIndex", "CategoricalIndex", - "DatetimeIndex", "Categorical", "CategoricalDtype", "StringDtype", "BooleanDtype", - "Float32Dtype", "Float64Dtype", "Timestamp", "DatetimeTZDtype", *( + "Categorical", "CategoricalDtype", "StringDtype", "BooleanDtype", + "Float32Dtype", "Float64Dtype", "NA", *( f"{sign}Int{bits}Dtype" for sign in ("", "U") for bits in (8, 16, 32, 64) ), }), "pandas.arrays": frozenset({ "BooleanArray", "IntegerArray", "FloatingArray", "StringArray", "Categorical", - "DatetimeArray", }), - "builtins": frozenset({"slice", "set", "frozenset", "complex", "bytearray", "range", "object"}), + "builtins": frozenset({ + "slice", "set", "frozenset", "complex", "bytes", "bytearray", "range", "object", + }), "collections": frozenset({"OrderedDict"}), "copyreg": frozenset({"_reconstructor"}), "_codecs": frozenset({"encode"}), } +#: numpy dtype kinds an array may have: no structured or void dtypes, +#: which nothing pycisTopic holds uses. +_KINDS = frozenset("biufcOUSMm") + +#: Classes a pickle calls with arguments, as the allowlist names them, and +#: the most arguments it passes: none of them parses a dtype from text, and +#: pandas' managers get no third, which would turn off their own check of +#: their blocks. Any other class is only ever created empty and given its +#: state. +_CONSTRUCTED = { + **{("builtins", name): 3 for name in ("slice", "range", "bytes", "bytearray")}, + **{("builtins", name): 1 for name in ("set", "frozenset")}, + ("builtins", "complex"): 2, + ("collections", "OrderedDict"): 1, + ("pandas.core.internals.managers", "BlockManager"): 2, + ("pandas.core.internals.managers", "SingleBlockManager"): 2, + ("pandas._libs.internals", "BlockPlacement"): 1, + ("pandas.core.arrays.string_", "StringDtype"): 2, # (storage, na_value), from pandas 3 +} + +#: What pandas pickles a DataFrame's or Series' state with (pandas 1.5 to 3): +#: the state is set as attributes as it stands, so nothing else, such as a +#: cache of columns, gets in. +_FRAME_STATE = frozenset({"_mgr", "_data", "_typ", "_metadata", "attrs", "_attrs", "_flags", "name", "_name"}) + +#: What a scipy matrix pickles: its arrays, its shape, and how it prints. +#: The flags scipy caches about its indices are dropped, so that scipy works +#: them out again from the checked arrays. +_SPARSE_STATE = frozenset({"_shape", "data", "indices", "indptr", "coords", "maxprint", "_maxprint"}) +_SPARSE_CACHED = frozenset({"has_canonical_format", "_has_canonical_format", + "has_sorted_indices", "_has_sorted_indices"}) + +#: The keys pandas' `_new_Index` is pickled with, across versions. +_INDEX_KEYS = frozenset({"data", "name", "names", "levels", "codes", "labels", "sortorder", + "start", "stop", "step"}) + class PycistopicObject: """A pycisTopic object, its attributes restored and nothing more. @@ -121,11 +186,268 @@ def _stand_in(name: str) -> type: return _STAND_INS[name] -class RefusedGlobal(pickle.UnpicklingError): +class RefusedPickle(pickle.UnpicklingError): + """A pickle this does not read: it names, calls or holds something a + pycisTopic object's pickle never does.""" + + +class RefusedGlobal(RefusedPickle): pass -class _Unpickler(pickle.Unpickler): +class _ArrayClass: + """`numpy.ndarray` as a pickle names it: the class `_reconstruct` makes, + and nothing a pickle may call.""" + + def __repr__(self) -> str: + return "numpy.ndarray" + + +_NDARRAY = _ArrayClass() + +#: Not found: unlike None, never an object a pickle holds. +_MISSING = object() + + +def _name(value: Any) -> str: + if value is _NDARRAY: + return "numpy.ndarray" + module = getattr(value, "__module__", None) or type(value).__module__ + name = getattr(value, "__qualname__", None) or type(value).__qualname__ + return f"{module}.{name}" + + +def _refuse(message: str) -> RefusedPickle: + return RefusedPickle(f"the pickle {message}, which a pycisTopic object's never does, so it is not read.") + + +def _is_count(value: Any) -> bool: + return isinstance(value, (int, np.integer)) and not isinstance(value, (bool, np.bool_)) and value >= 0 + + +def _shape(value: Any, what: str) -> Tuple[int, ...]: + if type(value) is not tuple or not all(_is_count(n) for n in value): + raise _refuse(f"gives {what} the shape {value!r}") + return tuple(int(n) for n in value) + + +def _size(shape: Tuple[int, ...]) -> int: + size = 1 + for n in shape: + size *= n + return size + + +def check_dtype(dtype: Any, what: str = "an array") -> np.dtype: + """`dtype`, if it is one numpy builds for itself: a kind pycisTopic's + arrays use, and the state numpy writes for it (flags, item size, + alignment), which numpy trusts when it reads a pickle.""" + if not isinstance(dtype, np.dtype): + raise _refuse(f"gives {what} the dtype {type(dtype).__name__}") + if dtype.kind not in _KINDS or dtype.names is not None or dtype.subdtype is not None: + raise _refuse(f"gives {what} the numpy dtype {dtype}") + if dtype.__reduce__() != np.dtype(dtype.str).__reduce__(): + raise _refuse(f"gives {what} a numpy dtype {dtype.str} whose state numpy never writes for it") + return dtype + + +def _vector(state: Dict[str, Any], key: str, kinds: str, what: str) -> np.ndarray: + value = state.get(key) + if type(value) is not np.ndarray or value.ndim != 1 or value.dtype.kind not in kinds: + described = ( + f"a {value.ndim}-d {value.dtype} array" if type(value) is np.ndarray else type(value).__name__ + ) + raise _refuse(f"gives {what} {key} as {described}") + return value + + +def _index(state: Dict[str, Any], key: str, what: str) -> np.ndarray: + """An index array as scipy makes them, and its compiled routines take: + int32 or int64, in the machine's byte order.""" + value = _vector(state, key, "i", what) + if value.dtype not in (np.dtype(np.int32), np.dtype(np.int64)): + raise _refuse(f"gives {what} {key} as {value.dtype}") + return value + + +def _within(values: np.ndarray, lo: int, hi: int, what: str) -> None: + """Every value in [lo, hi).""" + if values.size and (values.min() < lo or values.max() >= hi): + raise _refuse(f"holds {what} out of range: from {values.min()} to {values.max()}, outside [{lo}, {hi})") + + +def check_sparse(matrix: Any) -> None: + """A scipy matrix's state as scipy would have built it: an array of + each kind its format needs, a 2-D shape, indices within the shape, and + an `indptr` rising from 0 to the number of entries. scipy's C++ code + trusts all of it.""" + kind = type(matrix).__name__.split("_")[0] # by its class: `format` reads its state + what = f"a scipy {type(matrix).__name__}" + state = vars(matrix) + unknown = sorted(set(state) - _SPARSE_STATE) + if unknown: + raise _refuse(f"gives {what} {', '.join(unknown)}") + shape = state.get("_shape") + if type(shape) is not tuple or len(shape) != 2 or not all(_is_count(n) for n in shape): + raise _refuse(f"gives {what} the shape {shape!r}") + # Plain integers, as scipy makes them: it sizes its arrays from the + # shape, and `np.uint8(255) + 1` is 0. + n_rows, n_cols = (int(n) for n in shape) + matrix._shape = (n_rows, n_cols) + data = _vector(state, "data", "biufc", what) + if not data.dtype.isnative: + raise _refuse(f"gives {what} data as {data.dtype}") + if kind == "coo": + coords = state.get("coords") + if type(coords) is not tuple or len(coords) != 2: + raise _refuse(f"gives {what} coords that are not a row and a column") + rows = _index({"row": coords[0]}, "row", what) + cols = _index({"col": coords[1]}, "col", what) + if not len(rows) == len(cols) == len(data): + raise _refuse(f"gives {what} rows, columns and values of different length") + _within(rows, 0, n_rows, f"{what}'s rows") + _within(cols, 0, n_cols, f"{what}'s columns") + matrix.has_canonical_format = False # scipy sums duplicates when it needs to + return + indices = _index(state, "indices", what) + indptr = _index(state, "indptr", what) + n_major, n_minor = (n_rows, n_cols) if kind == "csr" else (n_cols, n_rows) + if len(indptr) != n_major + 1: + raise _refuse(f"gives {what} an indptr of {len(indptr):,} for {n_major:,} rows") + bounds = indptr.astype(np.int64) + if bounds[0] != 0 or (np.diff(bounds) < 0).any(): + raise _refuse(f"gives {what} an indptr that does not rise from 0") + nnz = int(bounds[-1]) + if nnz > len(indices) or nnz > len(data): + raise _refuse(f"gives {what} an indptr past its {min(len(indices), len(data)):,} entries") + if len(indices) != nnz or len(data) != nnz: + matrix.indices, matrix.data = indices[:nnz], data[:nnz] + indices = matrix.indices + _within(indices, 0, n_minor, f"{what}'s indices") + + +def _is_frame_or_series(value: Any) -> bool: + import pandas as pd + + return isinstance(value, (pd.DataFrame, pd.Series)) + + +def _check_manager(manager: Any, what: str) -> None: + """A DataFrame's or Series' blocks fill its axes: each block as many + rows as the index, and between them every column once. pandas checks + neither when it unpickles a frame, and indexes the blocks by the axes.""" + import pandas as pd + from pandas.api.extensions import ExtensionArray + from pandas.core.internals.managers import SingleBlockManager + + axes, blocks = getattr(manager, "axes", None), getattr(manager, "blocks", None) + if type(axes) is not list or not all(isinstance(axis, pd.Index) for axis in axes) or type(blocks) is not tuple: + raise _refuse(f"gives {what} a manager without its axes and blocks") + single = isinstance(manager, SingleBlockManager) + if len(axes) != (1 if single else 2) or (single and len(blocks) != 1): + raise _refuse(f"gives {what} {len(axes)} axes and {len(blocks)} blocks") + n_rows = len(axes[-1]) + covered = [] + for block in blocks: + values = getattr(block, "values", None) + if not (type(values) is np.ndarray or isinstance(values, ExtensionArray)): + raise _refuse(f"gives {what} a block of {type(values).__name__}") + if single: + if block.ndim != 1 or values.shape != (n_rows,): + raise _refuse(f"gives {what} {values.shape} values for {n_rows:,} rows") + continue + locs = block.mgr_locs.as_array + fits = values.shape == (len(locs), n_rows) or ( + values.ndim == 1 and len(locs) == 1 and values.shape[0] == n_rows + ) + if block.ndim != 2 or not fits: + raise _refuse( + f"gives {what} a block of shape {values.shape} for {len(locs):,} columns " + f"of {n_rows:,} rows" + ) + covered.append(np.asarray(locs, dtype=np.int64)) + if not single: + columns = np.sort(np.concatenate(covered)) if covered else np.zeros(0, dtype=np.int64) + if not np.array_equal(columns, np.arange(len(axes[0]))): + raise _refuse(f"gives {what} blocks that do not hold each of its {len(axes[0]):,} columns once") + + +def _check_pandas(value: Any) -> None: + """pandas objects whose integer codes and blocks its compiled code + indexes with unchecked: within what they index.""" + import pandas as pd + from pandas.core.arrays.masked import BaseMaskedArray + from pandas.core.internals.managers import BlockManager, SingleBlockManager + + if _is_frame_or_series(value): + manager = vars(value).get("_mgr") + kind = SingleBlockManager if isinstance(value, pd.Series) else BlockManager + if type(manager) is not kind: + raise _refuse(f"gives a {type(value).__name__} a {type(manager).__name__} for its data") + _check_manager(manager, f"a {type(value).__name__}") + elif type(value) in (BlockManager, SingleBlockManager): + _check_manager(value, f"a {type(value).__name__}") + elif isinstance(value, pd.Categorical): + what = "a Categorical" + codes = value._ndarray + categories = getattr(value.dtype, "categories", None) + if not isinstance(value.dtype, pd.CategoricalDtype) or not isinstance(categories, pd.Index): + raise _refuse(f"gives {what} no categories") + if type(codes) is not np.ndarray or codes.ndim != 1 or codes.dtype.kind != "i": + raise _refuse(f"gives {what} its codes as {type(codes).__name__}") + _within(codes, -1, len(categories), f"{what}'s codes") + elif isinstance(value, pd.MultiIndex): + levels, codes = list(value.levels), list(value.codes) + if len(levels) != len(codes) or len({len(c) for c in codes}) > 1: + raise _refuse("gives a MultiIndex levels and codes that do not match") + for level, level_codes in zip(levels, codes): + if type(level_codes) is not np.ndarray or level_codes.dtype.kind != "i": + raise _refuse("gives a MultiIndex codes that are not integers") + _within(level_codes, -1, len(level), "a MultiIndex's codes") + elif isinstance(value, BaseMaskedArray): + data, mask = getattr(value, "_data", None), getattr(value, "_mask", None) + if (type(data) is not np.ndarray or type(mask) is not np.ndarray + or mask.dtype != np.bool_ or data.shape != mask.shape): + raise _refuse(f"gives a {type(value).__name__} a mask that does not match its values") + elif isinstance(value, pd.arrays.StringArray): + values = value._ndarray + if type(values) is not np.ndarray or values.ndim != 1 or values.dtype != object: + raise _refuse("gives a StringArray values that are not text") + + +class _Unpickler(pickle._Unpickler): + """The pure-Python unpickler, every opcode that calls a name, sets an + object's state or fills a container checked first. + + `_fresh`: objects this pickle created and has not yet given their state; + each is given it once. `_unbuilt`: those of them that are numpy arrays + and dtypes, which nothing may hold until they have their state. + """ + + dispatch = dict(pickle._Unpickler.dispatch) + + def __init__(self, file: BinaryIO) -> None: + super().__init__(file) + self._fresh: Dict[int, Any] = {} + self._unbuilt: Dict[int, Any] = {} + self._dtype_args: Dict[int, Tuple[str, bool]] = {} + self._loaded: Dict[int, Any] = {} + self._own: Dict[Tuple[str, int], Callable[..., Any]] = {} + self._classes: Optional[Dict[int, Tuple[type, int]]] = None + + def load(self) -> Any: + value = super().load() + for item in self._loaded.values(): + if isinstance(item, np.ndarray): + continue + if type(item).__module__.startswith("scipy"): + check_sparse(item) + else: + _check_pandas(item) + return value + + # -- names --------------------------------------------------------- + def find_class(self, module: str, name: str) -> Any: if module.startswith(PYCISTOPIC): return _stand_in(name) @@ -141,7 +463,348 @@ def find_class(self, module: str, name: str) -> Any: f"the pickle asks for {module}.{name}, which a pycisTopic object " "never holds, so it is not read." ) - return _resolve(module, name) + value = _resolve(module, name) + if value is np.ndarray: + return _NDARRAY + own = self._wrappers().get(_key(value)) + return own if own is not None else value + + def _constructed(self) -> Dict[int, Tuple[type, int]]: + """The classes called with arguments, by identity: each and the most + arguments it is given.""" + if self._classes is None: + self._classes = {} + for (module, name), most in _CONSTRUCTED.items(): + try: + found = _resolve(module, name) + except RefusedGlobal: # pragma: no cover - pandas comes with the extra + continue + self._classes[id(found)] = (found, most) + return self._classes + + def _wrappers(self) -> Dict[Tuple[str, int], Callable[..., Any]]: + """The allowlisted callables given checked arguments instead, by the + real one's identity.""" + if not self._own: + real = { + np.dtype: self._dtype, + copyreg._reconstructor: self._reconstructor, + _resolve("_codecs", "encode"): self._encode, + _resolve("numpy._core.multiarray", "_reconstruct"): self._reconstruct, + _resolve("numpy._core.multiarray", "scalar"): self._scalar, + _resolve("numpy._core.numeric", "_frombuffer"): self._frombuffer, + } + try: + real[_resolve("pandas.core.indexes.base", "_new_Index")] = self._new_index + real[_ndarray_backed()] = self._ndarray_backed + except RefusedGlobal: # pragma: no cover - pandas comes with the extra + pass + self._own = {_key(value): own for value, own in real.items()} + return self._own + + # -- what may be held, and by what --------------------------------- + + def _created(self, obj: Any, fresh: bool = True) -> Any: + if fresh: + self._fresh[id(obj)] = obj + if not isinstance(obj, (PycistopicObject, np.dtype)) and ( + type(obj).__module__.startswith(("scipy", "pandas")) + ): + self._loaded[id(obj)] = obj + return obj + + def _held(self, items: Iterable[Any], by: str) -> None: + """Refuse a numpy array or dtype that is handed on before it has + its state: what it is handed to sees it empty.""" + unbuilt = self._unbuilt + for item in items: + if unbuilt.get(id(item), _MISSING) is item: + raise _refuse(f"hands a numpy {type(item).__name__} to {by} before giving it its state") + + def _operands(self, count: Optional[int] = None) -> list: + """The items an opcode is about to consume: the last `count`, or all + since the mark.""" + return self.stack if count is None else self.stack[len(self.stack) - count:] + + # -- calls ----------------------------------------------------------- + + def _call(self, func: Any, args: Any) -> Any: + if type(args) is not tuple: + raise _refuse(f"calls {_name(func)} with {type(args).__name__} as its arguments") + self._held(args, _name(func)) + if getattr(func, "__self__", None) is self: + return func(*args) + if isinstance(func, type): + if not args: + return self._created(func()) + found, most = self._constructed().get(id(func), (None, 0)) + if found is func: + if len(args) > most: + raise _refuse(f"constructs {_name(func)} from {len(args)} arguments") + if func in (bytes, bytearray) and any(not isinstance(a, (str, bytes)) for a in args): + raise _refuse(f"calls {func.__name__} with {', '.join(type(a).__name__ for a in args)}") + return self._created(func(*args), fresh=False) + raise _refuse(f"constructs {_name(func)} from arguments") + module = getattr(func, "__module__", "") or "" + if module.startswith("pandas") and getattr(func, "__name__", "") in ("_unpickle_block", "new_block"): + return self._created(func(*args), fresh=False) + raise _refuse(f"calls {_name(func)}") + + def load_reduce(self) -> None: + stack = self.stack + args = stack.pop() + stack[-1] = self._call(stack[-1], args) + + dispatch[pickle.REDUCE[0]] = load_reduce + + def _new(self, cls: Any, args: Any, kwargs: Any) -> Any: + if not isinstance(cls, type): + raise _refuse(f"creates an instance of {cls!r}") + if args or kwargs: + raise _refuse(f"constructs {_name(cls)} from arguments") + return self._created(cls.__new__(cls)) + + def load_newobj(self) -> None: + args = self.stack.pop() + self.stack[-1] = self._new(self.stack[-1], args, None) + + dispatch[pickle.NEWOBJ[0]] = load_newobj + + def load_newobj_ex(self) -> None: + kwargs = self.stack.pop() + args = self.stack.pop() + self.stack[-1] = self._new(self.stack[-1], args, kwargs) + + dispatch[pickle.NEWOBJ_EX[0]] = load_newobj_ex + + def _instantiate(self, klass: Any, args: Any) -> None: + self.append(self._new(klass, args, None)) + + # -- state ----------------------------------------------------------- + + def load_build(self) -> None: + stack = self.stack + state = stack.pop() + inst = stack[-1] + if self._fresh.pop(id(inst), _MISSING) is not inst: + raise _refuse(f"sets the state of a {_name(type(inst))} it did not just create, or more than once") + self._held((state,), "a state") + if isinstance(inst, np.dtype): + self._build_dtype(inst, state) + elif type(inst) is np.ndarray: + self._build_array(inst, state) + elif isinstance(inst, PycistopicObject): + if type(state) is not dict or not all(type(k) is str for k in state): + raise _refuse(f"gives a {inst.kind} the state {type(state).__name__}") + inst.__dict__.update(state) + elif type(inst).__module__.startswith("scipy"): + if type(state) is not dict or not all(type(k) is str for k in state): + raise _refuse(f"gives a {type(inst).__name__} the state {type(state).__name__}") + state = {k: v for k, v in state.items() if k not in _SPARSE_CACHED} + setstate = getattr(inst, "__setstate__", None) + if setstate is not None: + setstate(state) + else: + inst.__dict__.update(state) + elif type(inst).__module__.startswith("pandas"): + if _is_frame_or_series(inst) and type(state) is dict: + unknown = sorted(map(str, set(state) - _FRAME_STATE)) + if unknown: + raise _refuse(f"gives a {type(inst).__name__} {', '.join(unknown)}") + stack.append(state) + super().load_build() + else: + raise _refuse(f"sets the state of {_name(type(inst))}") + self._unbuilt.pop(id(inst), None) + + dispatch[pickle.BUILD[0]] = load_build + + # -- containers: only dict, list and set are filled -------------------- + + def _fill(self, target: Any, kinds: Tuple[type, ...], items: list, by: str) -> None: + if type(target) not in kinds: + raise _refuse(f"{by} into a {_name(type(target))}") + self._held(items, f"a {type(target).__name__}") + + def load_append(self) -> None: + self._fill(self.stack[-2], (list,), self.stack[-1:], "appends") + super().load_append() + + dispatch[pickle.APPEND[0]] = load_append + + def load_appends(self) -> None: + self._fill(self.metastack[-1][-1], (list,), self.stack, "appends") + super().load_appends() + + dispatch[pickle.APPENDS[0]] = load_appends + + def load_setitem(self) -> None: + self._fill(self.stack[-3], (dict, OrderedDict), self.stack[-2:], "writes") + super().load_setitem() + + dispatch[pickle.SETITEM[0]] = load_setitem + + def load_setitems(self) -> None: + self._fill(self.metastack[-1][-1], (dict, OrderedDict), self.stack, "writes") + super().load_setitems() + + dispatch[pickle.SETITEMS[0]] = load_setitems + + def load_additems(self) -> None: + self._fill(self.metastack[-1][-1], (set,), self.stack, "adds") + super().load_additems() + + dispatch[pickle.ADDITEMS[0]] = load_additems + + def _packed(load: Callable[[Any], None], count: Optional[int]) -> Callable[[Any], None]: # type: ignore[misc] + def packed(self: "_Unpickler") -> None: + if self._unbuilt: + self._held(self._operands(count), "a tuple, list, dict or set") + load(self) + + return packed + + for _opcode, _count in ( + (pickle.TUPLE, None), (pickle.TUPLE1, 1), (pickle.TUPLE2, 2), (pickle.TUPLE3, 3), + (pickle.LIST, None), (pickle.DICT, None), (pickle.FROZENSET, None), + ): + dispatch[_opcode[0]] = _packed(dispatch[_opcode[0]], _count) + del _opcode, _count, _packed + + # -- numpy ----------------------------------------------------------- + + def _dtype(self, spec: Any, align: Any = False, copy: Any = True) -> np.dtype: + """`numpy.dtype(spec, align, copy)`, from a type's name only; always + a copy, so the state set on it next is set on nothing else.""" + if type(spec) is not str: + raise _refuse(f"builds a numpy dtype from {type(spec).__name__}") + try: + dtype = np.dtype(spec, bool(align), True) + except (TypeError, ValueError) as exc: + raise _refuse(f"builds the numpy dtype {spec!r} ({exc})") from None + check_dtype(dtype, "a dtype") + self._dtype_args[id(dtype)] = (spec, bool(align)) + self._unbuilt[id(dtype)] = dtype + return self._created(dtype) + + def _build_dtype(self, dtype: np.dtype, state: Any) -> None: + """The state set first on a copy, and only on `dtype` when the copy + comes out as numpy builds that dtype: so `dtype`, which an array may + already use, is never wrong, even until the refusal.""" + spec, align = self._dtype_args.pop(id(dtype)) + if type(state) is not tuple: + raise _refuse(f"gives a numpy dtype the state {type(state).__name__}") + probe = np.dtype(spec, align, True) + try: + probe.__setstate__(state) + except Exception as exc: # numpy's own refusal of the tuple + raise _refuse(f"gives the numpy dtype {spec!r} the state {state!r} ({exc})") from None + check_dtype(probe, "a dtype") + dtype.__setstate__(state) + + def _reconstruct(self, subtype: Any, shape: Any, typecode: Any = b"b") -> np.ndarray: + """numpy's `_reconstruct`, as numpy pickles an array: an empty + ndarray, given its state next.""" + if subtype is not _NDARRAY: + raise _refuse(f"reconstructs {subtype!r} through numpy's _reconstruct") + if shape != (0,) or type(shape) is not tuple: + raise _refuse(f"reconstructs an array of shape {shape!r} through numpy's _reconstruct") + array = np.ndarray.__new__(np.ndarray, (0,), np.dtype("b")) + self._unbuilt[id(array)] = array + return self._created(array) + + def _build_array(self, array: np.ndarray, state: Any) -> None: + """An array's state: its shape, a checked dtype, and data that fills + the shape exactly, as numpy writes it.""" + if type(state) is not tuple or len(state) not in (4, 5): + raise _refuse(f"gives a numpy array the state {type(state).__name__}") + shape, dtype, fortran, raw = state[-4:] + if len(state) == 5 and not _is_count(state[0]): + raise _refuse(f"gives a numpy array the version {state[0]!r}") + shape = _shape(shape, "a numpy array") + dtype = check_dtype(dtype) + if not isinstance(fortran, (bool, int)): + raise _refuse(f"gives a numpy array the order {fortran!r}") + size = _size(shape) + if dtype.hasobject: + if type(raw) is not list or len(raw) != size: + raise _refuse(f"gives a numpy object array of {size:,} items its data as {type(raw).__name__}") + elif type(raw) not in (bytes, bytearray) or len(raw) != size * dtype.itemsize: + length = len(raw) if isinstance(raw, (bytes, bytearray)) else type(raw).__name__ + raise _refuse( + f"gives a numpy {dtype} array of shape {shape} {length} bytes of data, " + f"not {size * dtype.itemsize:,}" + ) + array.__setstate__(state) + + def _frombuffer(self, buffer: Any, dtype: Any, shape: Any, order: Any) -> np.ndarray: + """numpy's `_frombuffer` (protocol 5): a checked dtype, and a buffer + the shape fills exactly.""" + from numpy import frombuffer + + dtype = check_dtype(dtype) + shape = _shape(shape, "a numpy array") + if dtype.hasobject or type(buffer) not in (bytes, bytearray) or order not in ("C", "F"): + raise _refuse(f"builds a numpy {dtype} array from {type(buffer).__name__}") + if len(buffer) != _size(shape) * dtype.itemsize: + raise _refuse(f"gives a numpy {dtype} array of shape {shape} {len(buffer):,} bytes of data") + return frombuffer(buffer, dtype=dtype).reshape(shape, order=order) + + def _scalar(self, dtype: Any, raw: Any = None) -> Any: + """numpy's `scalar`: an object as itself, or bytes as wide as the + dtype's items.""" + dtype = check_dtype(dtype, "a scalar") + if dtype.hasobject: + return raw + if type(raw) is not bytes or len(raw) != dtype.itemsize: + raise _refuse(f"gives a numpy {dtype} scalar {type(raw).__name__} of the wrong size") + return np.frombuffer(raw, dtype=dtype)[0] + + # -- builtins and pandas ------------------------------------------- + + def _encode(self, text: Any, encoding: Any = "utf-8") -> bytes: + """`_codecs.encode`, which protocols 0-2 write bytes with.""" + if type(text) is not str or encoding not in ("latin1", "latin-1"): + raise _refuse(f"encodes {type(text).__name__} as {encoding!r}") + return text.encode("latin-1") + + def _reconstructor(self, cls: Any, base: Any, state: Any) -> Any: + """`copyreg._reconstructor` as protocols 0 and 1 write it: an empty + instance, given its state next.""" + if base is not object or state is not None or not isinstance(cls, type): + raise _refuse(f"calls copyreg._reconstructor for {cls!r} from {base!r}") + return self._created(object.__new__(cls)) + + def _new_index(self, cls: Any, d: Any) -> Any: + """pandas' `_new_Index(cls, d)`, which passes `d` to the class: only + the keys pandas pickles an index with, so never a dtype or a + timezone to parse.""" + import pandas as pd + from pandas.core.indexes.base import _new_Index + + if not (isinstance(cls, type) and issubclass(cls, pd.Index)) or type(d) is not dict: + raise _refuse(f"builds an index with pandas' _new_Index from {cls!r}") + unknown = sorted(set(d) - _INDEX_KEYS) + if unknown: + raise _refuse(f"builds a {cls.__name__} with {', '.join(map(str, unknown))}") + return self._created(_new_Index(cls, d), fresh=False) + + def _ndarray_backed(self, cls: Any, checksum: Any, state: Any = None) -> Any: + """pandas' unpickling of its Categorical and StringArray.""" + import pandas as pd + + if cls not in (pd.Categorical, pd.arrays.StringArray): + raise _refuse(f"builds {cls!r} through pandas' __pyx_unpickle_NDArrayBacked") + return self._created(_ndarray_backed()(cls, checksum, state), fresh=state is None) + + +def _ndarray_backed() -> Callable[..., Any]: + # Looked up by name: in a class body, `__pyx_...` would be mangled. + return _resolve("pandas._libs.arrays", "__pyx_unpickle_NDArrayBacked") + + +def _key(value: Any) -> Tuple[str, int]: + return ("callable", id(value)) def _renamed(module: str, name: str) -> Tuple[str, str]: @@ -192,7 +855,7 @@ def open_pickle(path: Path) -> Iterator[BinaryIO]: def load(path: Path) -> Any: - """Unpickle `path` through the allowlist.""" + """Unpickle `path` through the allowlist and the checks.""" with open_pickle(path) as fh: return _Unpickler(fh).load() diff --git a/tests/test_cistopic.py b/tests/test_cistopic.py index 3365be7..f731847 100644 --- a/tests/test_cistopic.py +++ b/tests/test_cistopic.py @@ -109,15 +109,20 @@ def test_harmony_topics_are_kept_too(tmp_path): assert np.array_equal(got.obsm["X_topic_harmony"], harmony[truth.obj.cell_names].to_numpy().T) +@pytest.mark.parametrize("budget", [20, 5], ids=["cells", "features"]) @pytest.mark.parametrize("dense", [False, True], ids=["csr", "dense"]) -def test_an_imputed_features_pickle_is_cells_by_features(tmp_path, dense, monkeypatch): +@pytest.mark.parametrize("suffix", [".h5ad", ".zarr2", ".zarr3"]) +def test_an_imputed_features_pickle_is_cells_by_features(tmp_path, dense, suffix, budget, monkeypatch): + """Blocks of a few cells (8 features to a cell fit `cells`), or runs of + one cell's features (they do not fit `features`), straddling the edges.""" from adata.conversion import cistopic - monkeypatch.setattr(cistopic, "DENSE_CELLS", 4) # blocks straddle the cells + monkeypatch.setattr(cistopic, "BLOCK_ENTRIES", budget, raising=False) truth = write_imputed_pickle(tmp_path / "imputed_acc_obj.pkl") - out = tmp_path / "out.h5ad" - result = _ok(truth.path, "-o", out, *(["--dense"] if dense else [])) - assert "pycisTopic CistopicImputedFeatures → h5ad, 15 cells x 8 features" in _out(result) + out = tmp_path / f"out{suffix[:5]}" + zarr = ["--zarr-format", suffix[-1]] if suffix != ".h5ad" else [] + result = _ok(truth.path, "-o", out, *zarr, *(["--dense"] if dense else [])) + assert f"pycisTopic CistopicImputedFeatures → {suffix[1:5]}, 15 cells x 8 features" in _out(result) got = _read(out) assert isinstance(got.X, np.ndarray) == dense X = got.X if dense else got.X.toarray() @@ -126,6 +131,116 @@ def test_an_imputed_features_pickle_is_cells_by_features(tmp_path, dense, monkey assert list(got.obs_names) == truth.obj.cell_names +@pytest.mark.parametrize("suffix", [".h5ad", ".zarr"]) +def test_dense_output_chunks_are_bounded_in_bytes_however_wide_the_matrix(tmp_path, suffix, monkeypatch): + """F007: `--dense` chunked X as (4096 cells, every feature), so a chunk + grew with the features: 3.2 GB at 200,000 float32 ones, and past + 262,144 HDF5 refused it (a chunk must be under 4 GB). A chunk now holds + a fixed number of entries, a run of one cell's features when even one + cell is wider than that.""" + import h5py + import zarr + + from adata.conversion import cistopic + + monkeypatch.setattr(cistopic, "BLOCK_ENTRIES", 64, raising=False) + for n_features, expected in ((16, (4, 16)), (1000, (1, 64))): + truth = write_imputed_pickle(tmp_path / f"wide{n_features}.pkl", n_cells=10, n_features=n_features) + out = tmp_path / f"wide{n_features}{suffix}" + _ok(truth.path, "-o", out, "--dense") + if suffix == ".h5ad": + with h5py.File(out, "r") as f: + chunks = f["X"].chunks + else: + chunks = zarr.open_group(out, mode="r")["X"].chunks + assert chunks == expected, chunks + assert np.array_equal(_read(out).X, truth.obj.mtx.T) + + +def _sparse_imputed(tmp_path, fmt): + truth = write_imputed_pickle(tmp_path / f"gene_act_{fmt}.pkl", n_cells=15, n_features=8) + truth.obj.mtx = getattr(sparse, f"{fmt}_matrix")(truth.obj.mtx) + truth.path.write_bytes(pickle.dumps(truth.obj, protocol=4)) + return truth + + +@pytest.mark.parametrize("fmt", ["csr", "csc", "coo"]) +@pytest.mark.parametrize("suffix", [".h5ad", ".zarr2", ".zarr3"]) +def test_dense_is_honoured_for_a_sparse_imputed_matrix(tmp_path, fmt, suffix, monkeypatch): + """F100: gene activity is pickled with a sparse `mtx`, and `--dense` was + then ignored: the summary said dense and X came out CSR. It is written + dense, a block of cells at a time from the CSR the transpose makes.""" + from adata.conversion import cistopic + + monkeypatch.setattr(cistopic, "BLOCK_ENTRIES", 20, raising=False) + truth = _sparse_imputed(tmp_path, fmt) + out = tmp_path / f"out{suffix[:5]}" + zarr = ["--zarr-format", suffix[-1]] if suffix != ".h5ad" else [] + result = _ok(truth.path, "-o", out, "--dense", *zarr) + assert ", dense" in _out(result) + got = _read(out) + assert isinstance(got.X, np.ndarray), type(got.X) + assert got.X.dtype == np.float32 and np.array_equal(got.X, truth.obj.mtx.T.toarray()) + assert sorted(p.name for p in tmp_path.iterdir()) == sorted([truth.path.name, out.name]) + if suffix == ".h5ad": + import h5py + + with h5py.File(out, "r") as f: + assert not [k for k in f if k.startswith("__cistopic")], "scratch left behind" + else: + import zarr + + assert not [k for k in zarr.open_group(out, mode="r") if k.startswith("__cistopic")] + + +@pytest.mark.parametrize("fmt", ["csr", "csc", "coo"]) +def test_a_sparse_imputed_matrix_is_csr_without_dense(tmp_path, fmt): + truth = _sparse_imputed(tmp_path, fmt) + out = tmp_path / "out.h5ad" + _ok(truth.path, "-o", out) + got = _read(out) + assert got.X.format == "csr" and np.array_equal(got.X.toarray(), truth.obj.mtx.T.toarray()) + + +def test_an_imputed_matrix_that_is_not_numbers_is_refused(tmp_path): + truth = write_imputed_pickle(tmp_path / "imputed.pkl", n_cells=3, n_features=2) + truth.obj.mtx = truth.obj.mtx.astype(object) + truth.path.write_bytes(pickle.dumps(truth.obj, protocol=4)) + result = _convert(truth.path, "-o", tmp_path / "out.h5ad") + assert result.exit_code == 1 + assert "its matrix holds object values, not numbers" in _out(result) + + +@pytest.mark.perf +@pytest.mark.parametrize("dense", [False, True], ids=["csr", "dense"]) +def test_writing_an_imputed_matrix_holds_a_block_not_every_feature_of_4096_cells(tmp_path, dense, monkeypatch): + """F025 (and F007's copy): each block was 4096 cells x every feature, + and CSR took about 9x its bytes in `np.nonzero`'s int64 coordinates: 3 + GB for 4096 x 20,000, about 50 GB at 500,000 regions, where loading the + pickle took 360 MB. A block now holds a fixed number of entries, so + what writing allocates does not grow with the features.""" + from adata.conversion import cistopic + from adata.conversion.base import Options + from adata.storage import open_store + from tests.perf_counters import FLAT_SIZES, assert_independent_of, count_allocations + + # 256 cells of 64 features are one block; of 4096, 64 blocks of 4 cells. + monkeypatch.setattr(cistopic, "BLOCK_ENTRIES", 1 << 14, raising=False) + + def measure(n_features: int) -> int: + truth = write_imputed_pickle(tmp_path / f"imputed{n_features}.pkl", n_cells=256, n_features=n_features) + out = tmp_path / f"out{n_features}{int(dense)}.h5ad" + plan = cistopic.CistopicImport().plan( + truth.path, out, "cistopic", "h5ad", Options(dense=dense, given=frozenset()) + ) + with open_store(out, "w") as store: + with count_allocations() as peak: + cistopic._write_imputed(store.root, plan.detail) + return peak[0] + + assert_independent_of(measure, what="bytes allocated writing X", axis="n_features", sizes=FLAT_SIZES) + + def test_a_pickle_that_asks_for_code_is_refused_unrun(tmp_path): """The point of the allowlist: a global that a pycisTopic object never holds is refused by name, and never called.""" @@ -148,6 +263,19 @@ def __reduce__(self): assert not out.exists() +def test_a_pickle_numpy_or_pandas_cannot_read_is_named_as_such(tmp_path): + """A pickle cut short, or one whose contents pandas itself refuses, + failed with the library's bare error; the file is named.""" + truth = write_cistopic_pickle(tmp_path / "obj.pkl", model=False, projections=False) + data = truth.path.read_bytes() + truth.path.write_bytes(data[: len(data) // 2]) + out = tmp_path / "out.h5ad" + result = _convert(truth.path, "-o", out) + assert result.exit_code == 1 + assert "obj.pkl could not be read as a pycisTopic pickle" in _out(result) + assert not out.exists() + + def test_a_pickle_that_is_not_pycistopics_needs_from(tmp_path): path = tmp_path / "other.pkl" path.write_bytes(pickle.dumps({"a": [1, 2, 3]})) diff --git a/tests/test_formats.py b/tests/test_formats.py index 4962420..728947f 100644 --- a/tests/test_formats.py +++ b/tests/test_formats.py @@ -85,18 +85,117 @@ def test_image_export_scales_unit_range_floats(store, temp_dir): assert 120 <= got[0, 1] <= 135 -def test_image_export_clips_floats_already_in_byte_range(store, temp_dir): +def test_image_export_keeps_floats_already_in_byte_range(store, temp_dir): path, opener = store with opener("a") as root: - ew.write_dense(root["uns"], "img", np.array([[-5.0, 300.0], [10.0, 20.0]])) + ew.write_dense(root["uns"], "img", np.array([[0.0, 255.0], [10.0, 20.0]])) out = temp_dir / "f.png" export_image(root, "uns/img", out, console) from PIL import Image + assert np.array_equal(np.asarray(Image.open(out)), [[0, 255], [10, 20]]) + + +@pytest.mark.parametrize("suffix", [".png", ".tiff", ".tif"]) +def test_a_16_bit_image_exports_as_16_bit(store, png, temp_dir, suffix): + """F066: values from 255 up were clipped to 255 with no warning, so a + 16-bit fluorescence or morphology image (stored as uint16 by `import + image`) came back as [0, 255, 255, ...]. One channel of 16 bits is + written as a 16-bit PNG or TIFF, values as they are.""" + from PIL import Image + + array = (np.arange(24, dtype=np.uint16) * 2500).reshape(4, 6) + source = temp_dir / "deep.png" + Image.fromarray(array).save(source) + path, opener = store + with opener("a") as root: + import_image(root, "uns/img", source, console) + with opener("r") as root: + assert root["uns"]["img"].dtype == np.uint16 + out = temp_dir / f"out{suffix}" + export_image(root, "uns/img", out, console) + + got = np.asarray(Image.open(out)) + assert got.dtype == np.uint16 and np.array_equal(got, array) + + +@pytest.mark.parametrize( + "array, phrase", + [ + ((np.arange(12, dtype=np.uint16) * 1000).reshape(2, 2, 3), "a colour image keeps 8 bits a channel"), + (np.array([[0.0, 4095.0], [10.0, 20.0]]), "from 0 to 4095"), + (np.array([[-5.0, 3.0], [10.0, 20.0]]), "from -5 to 20"), + (np.array([[0, 70000], [1, 2]], dtype=np.int32), "whole numbers from 0 to 65535"), + (np.array([[0.5, np.nan], [0.1, 0.2]]), "not finite"), + ], + ids=["rgb16", "float-4095", "negative", "past-16-bit", "nan"], +) +def test_image_values_an_image_cannot_hold_are_refused_unless_scaled(store, temp_dir, array, phrase): + """F066: such values were clipped into 0-255 silently. They are refused, + naming the values, unless --scale asks for them to be stretched onto + 0-255, which maps the smallest to 0 and the largest to 255.""" + from PIL import Image + + path, opener = store + with opener("a") as root: + ew.write_dense(root["uns"], "img", array) + out = temp_dir / "x.png" + with opener("r") as root: + with pytest.raises(ValueError, match="--scale") as caught: + export_image(root, "uns/img", out, console) + assert phrase in str(caught.value) + assert not out.exists() + export_image(root, "uns/img", out, console, scale=True) + got = np.asarray(Image.open(out)) - assert got[0, 0] == 0, "negatives clip to 0" - assert got[0, 1] == 255, "values above 255 clip to 255" + values = array.astype(np.float64) + finite = np.isfinite(values) + lo, hi = values[finite].min(), values[finite].max() + expected = np.where(finite, np.rint((np.where(finite, values, lo) - lo) / (hi - lo) * 255), 0) + assert got.dtype == np.uint8 and np.array_equal(got, expected.astype(np.uint8)) + + +def test_sixteen_bits_into_a_format_that_keeps_eight_is_refused(store, temp_dir): + """F066: JPEG holds 8 bits a pixel, so a 16-bit image is refused there + rather than clipped, pointing at the formats that keep 16.""" + path, opener = store + with opener("a") as root: + ew.write_dense(root["uns"], "img", (np.arange(4, dtype=np.uint16) * 1000).reshape(2, 2)) + with pytest.raises(ValueError, match=r"JPEG keeps 8 bits a pixel; a \.png or \.tiff keeps 16"): + export_image(root, "uns/img", temp_dir / "x.jpg", console) + + +def test_integers_that_fit_a_byte_export_as_they_are(store, temp_dir): + """Whatever their dtype, whole numbers from 0 to 255 are an 8-bit image.""" + from PIL import Image + + path, opener = store + array = np.array([[0, 255], [7, 128]], dtype=np.int64) + with opener("a") as root: + ew.write_dense(root["uns"], "img", np.stack([array] * 3, axis=-1)) + export_image(root, "uns/img", temp_dir / "rgb.png", console) + assert np.array_equal(np.asarray(Image.open(temp_dir / "rgb.png")), np.stack([array] * 3, axis=-1)) + + +def test_export_image_takes_scale_on_the_command_line(store, temp_dir): + from typer.testing import CliRunner + + from adata.cli import app + + path, opener = store + with opener("a") as root: + ew.write_dense(root["uns"], "img", np.array([[0.0, 4095.0], [1.0, 2.0]])) + runner = CliRunner() + out = temp_dir / "x.png" + refused = runner.invoke(app, ["export", "image", str(path), "uns/img", "-o", str(out)]) + assert refused.exit_code == 1 and "--scale" in " ".join((refused.stdout + refused.stderr).split()) + assert not out.exists() + done = runner.invoke(app, ["export", "image", str(path), "uns/img", "-o", str(out), "--scale"]) + assert done.exit_code == 0, done.stdout + done.stderr + from PIL import Image + + assert np.array_equal(np.asarray(Image.open(out)), [[0, 255], [0, 0]]) def test_image_export_handles_booleans(store, temp_dir): diff --git a/tests/test_pickled.py b/tests/test_pickled.py new file mode 100644 index 0000000..9ea92ed --- /dev/null +++ b/tests/test_pickled.py @@ -0,0 +1,536 @@ +"""The restricted unpickler: what a crafted pickle cannot get past. + +`formats/pickled.py` reads pycisTopic's pickles through an allowlist of the +names a pycisTopic object holds, and checks what numpy, scipy and pandas are +given before any of them acts on it. Every crafted pickle here names only +allowlisted globals: the checks are what refuse it. + +The crafted files are loaded in a child process, because before the checks +some of them crashed the interpreter, which would have taken the test run +with it. The child loads the file and then uses what it got, as a +conversion would, so a pickle that slipped through shows as a crash or a +wrong answer rather than as nothing. +""" + +from __future__ import annotations + +import copyreg +import io +import pickle +import subprocess +import sys +from pathlib import Path + +import numpy as np +import pytest + +from adata.formats import pickled +from tests.conftest import _pycistopic_classes, write_cistopic_pickle, write_imputed_pickle + +pd = pytest.importorskip("pandas") +sp = pytest.importorskip("scipy.sparse") + +_reconstruct = np.ndarray.__reduce__(np.zeros(1))[0] + + +class _Crafted: + """Pickles as `func(*args)`, then `state` set on what that returned: + the REDUCE and BUILD a crafted file holds.""" + + def __init__(self, func, args, state=None): + self.func, self.args, self.state = func, args, state + + def __reduce__(self): + if self.state is None: + return self.func, self.args + return self.func, self.args, self.state + + +class _New: + """Pickles as `cls.__new__(cls)` (NEWOBJ), then `state` set on it: how + protocol 2 and later write a class's instance. `__class__` is what the + pickler checks NEWOBJ's class against.""" + + def __init__(self, cls, state): + self.cls, self.state = cls, state + + @property + def __class__(self): + return self.cls + + def __reduce_ex__(self, protocol): + return copyreg.__newobj__, (self.cls,), self.state + + +def _array(dtype, shape, raw): + """An ndarray as numpy pickles one, with whatever state is given.""" + return _Crafted(_reconstruct, (np.ndarray, (0,), b"b"), (1, shape, dtype, False, raw)) + + +def _dtype(spec, state): + """A numpy dtype as numpy pickles one, with whatever state is given.""" + return _Crafted(np.dtype, (spec, False, True), state) + + +# The child: load, then use what was loaded as a conversion would. Exit 1 +# with REFUSED when the unpickler refuses; OPENED lists the watched path if +# anything opened it. +_CHILD = r""" +import sys +watched = sys.argv[2] if len(sys.argv) > 2 else None +opened = [] +if watched: + sys.addaudithook(lambda event, args: opened.append(args[0]) + if event == "open" and str(args[0]) == watched else None) +from adata.formats import pickled +try: + try: + obj = pickled.load(sys.argv[1]) + except pickled.RefusedPickle as exc: + print("REFUSED", exc) + sys.exit(1) + import numpy as np + import scipy.sparse as sp + values = list(vars(obj).values()) if isinstance(obj, pickled.PycistopicObject) else [obj] + for value in values: + if sp.issparse(value): + print(value.tocsr().toarray().sum()) + else: + print(repr(value)[:200]) + print("LOADED") +finally: + if opened: + print("OPENED", opened) +""" + + +def _load_isolated(path: Path, watched: Path = None) -> subprocess.CompletedProcess: + args = [sys.executable, "-c", _CHILD, str(path)] + ([str(watched)] if watched else []) + return subprocess.run(args, capture_output=True, text=True, timeout=120) + + +def _write(path: Path, obj, protocol: int = 4) -> Path: + path.write_bytes(pickle.dumps(obj, protocol=protocol)) + return path + + +def _refused(result: subprocess.CompletedProcess, *phrases: str) -> str: + output = result.stdout + result.stderr + assert result.returncode == 1 and "REFUSED" in result.stdout, (result.returncode, output[-2000:]) + for phrase in phrases: + assert phrase in result.stdout, result.stdout + assert "OPENED" not in result.stdout, result.stdout + return result.stdout + + +# --------------------------------------------------------------------------- +# numpy: dtypes and arrays + + +@pytest.mark.parametrize( + "dtype, raw", + [ + # A structured dtype with an object field whose pickled flags were + # cleared: numpy then treats the field's pointer as plain bytes. + (_dtype("V8", (3, "|", None, ("f0",), {"f0": (np.dtype("O"), 0)}, 8, 1, 0)), b"\x41" * 8), + # A plain object dtype with its flags cleared: the raw bytes become + # the array's object pointers. + (_dtype("O8", (3, "|", None, None, None, -1, -1, 0)), b"\x41" * 8), + # A text dtype claiming 7 bytes an item, which UCS-4 text never is. + (_dtype("U1", (3, "<", None, None, None, 7, 3, 8)), b"\x41" * 7), + ], + ids=["structured", "object", "itemsize"], +) +def test_a_dtype_whose_pickled_state_numpy_never_writes_is_refused(tmp_path, dtype, raw): + """F040: numpy trusts the flags and item size a pickled dtype carries, + so a dtype could claim no objects while holding them, and raw bytes + from the file were read as object pointers: the interpreter crashed + (or worse), with only allowlisted names in the file.""" + path = _write(tmp_path / "evil.pkl", _array(dtype, (1,), raw)) + _refused(_load_isolated(path), "dtype") + + +def test_the_crafted_dtype_is_refused_inside_a_cistopic_object_too(tmp_path): + """F040, end to end: the crafted array as a CistopicObject's cell + names, which the converter iterates while planning.""" + import scipy.sparse as sp + + classes = _pycistopic_classes() + obj = classes["CistopicObject"]() + obj.fragment_matrix = sp.csr_matrix(np.ones((1, 1), dtype=np.int32)) + obj.region_names = ["r0"] + obj.cell_names = _array(_dtype("O8", (3, "|", None, None, None, -1, -1, 0)), (1,), b"\x41" * 8) + path = _write(tmp_path / "evil.pkl", obj) + out = tmp_path / "out.h5ad" + result = subprocess.run( + [sys.executable, "-c", "from adata.cli import app; app()", "convert", str(path), "-o", str(out)], + capture_output=True, text=True, timeout=120, + ) + assert result.returncode == 1, (result.returncode, result.stdout[-2000:], result.stderr[-2000:]) + assert "whose state numpy never writes" in " ".join((result.stdout + result.stderr).split()) + assert not out.exists() + + +def test_array_data_that_does_not_fill_the_shape_is_refused(tmp_path): + path = _write(tmp_path / "short.pkl", _array(np.dtype(" bytes: + return b"\x80\x02" + b"".join(ops) + b"." + + +def _global(module: str, name: str) -> bytes: + return b"c" + module.encode() + b"\n" + name.encode() + b"\n" + + +# dtype(' bytes: + return ( + _global(_reconstruct.__module__, "_reconstruct") + _global("numpy", "ndarray") + + b"K\x00\x85" + b"C\x01b" + b"\x87R" # _reconstruct(ndarray, (0,), b'b') + + b"(K\x01K\x01\x85" + _INT64 + + b"\x89C" + bytes([len(raw)]) + raw + b"tb" # BUILD (1, (1,), dtype('i8'), False, raw) + ) + + +def test_an_array_cannot_be_rebuilt_after_it_was_checked(tmp_path): + """Checks made when an array is built mean nothing if the pickle can + build it again later with other contents (or write into it): neither + happens in a pickle numpy writes.""" + rebuilt = _assemble( + _int_array_ops(b"\x00" * 8), b"q\x00", + b"(K\x01K\x01\x85" + _INT64 + + b"\x89C\x08" + b"\xff" * 8 + b"tb", # a second BUILD on the same array + ) + _refused(_load_isolated(_bytes(tmp_path / "rebuilt.pkl", rebuilt)), "once") + + written = _assemble(_int_array_ops(b"\x00" * 8), b"K\x00K\x07s") # array[0] = 7 + _refused(_load_isolated(_bytes(tmp_path / "written.pkl", written)), "ndarray") + + +def _bytes(path: Path, data: bytes) -> Path: + path.write_bytes(data) + return path + + +def test_an_array_cannot_be_used_before_it_is_built(tmp_path): + """An array is empty until its state is set: one handed to something + else first, then filled, would bypass whatever that something checked.""" + used_early = _assemble( + _global(_reconstruct.__module__, "_reconstruct") + _global("numpy", "ndarray") + + b"K\x00\x85C\x01b\x87R" + b"q\x00" + + b"]q\x01h\x00a" # [array] before the array's BUILD + + b"0h\x00" + + b"(K\x01K\x01\x85" + _INT64 + + b"\x89C\x08" + b"\x00" * 8 + b"tb0h\x01", + ) + _refused(_load_isolated(_bytes(tmp_path / "early.pkl", used_early)), "before") + + +def test_numpys_own_reconstruction_is_not_callable_with_other_arguments(tmp_path): + """`numpy.ndarray` is named by a pickle only as the class + `_reconstruct` makes; called itself, it can view any buffer.""" + called = _assemble(_global("numpy", "ndarray") + b"K\x01\x85\x85R") # ndarray((1,)) + _refused(_load_isolated(_bytes(tmp_path / "called.pkl", called)), "numpy.ndarray") + + big = _assemble( + _global(_reconstruct.__module__, "_reconstruct") + _global("numpy", "ndarray") + + b"J\x00\x00\x00\x40\x85C\x01b\x87R" # _reconstruct(ndarray, (2**30,), b'b') + ) + _refused(_load_isolated(_bytes(tmp_path / "big.pkl", big)), "_reconstruct") + + +def test_a_scalar_whose_bytes_do_not_fit_its_dtype_is_refused(tmp_path): + scalar = np.int64(1).__reduce__()[0] # numpy's `scalar`, as a pickle names it + path = _write(tmp_path / "scalar.pkl", _Crafted(scalar, (np.dtype("` from that + path, so a pickle naming one made the process open any local file.""" + path = _write(tmp_path / "tz.pkl", _tz_routes(f"dateutil/{tz_file}")[route]) + _refused(_load_isolated(path, watched=tz_file)) + + +def test_a_class_is_never_constructed_from_a_pickles_arguments(tmp_path): + """pandas' and scipy's classes are created empty and given their state; + their constructors, which parse dtypes and build from anything, are not + run on what a file holds.""" + path = _write(tmp_path / "frame.pkl", _Crafted(pd.DataFrame, ({"a": [1, 2]},))) + _refused(_load_isolated(path), "pandas.DataFrame") + path = _write(tmp_path / "csr.pkl", _Crafted(sp.csr_matrix, ((np.array([1]), np.array([5]), np.array([0, 1])),))) + _refused(_load_isolated(path), "csr_matrix") + path = _write(tmp_path / "recon.pkl", _Crafted(copyreg._reconstructor, (pd.DataFrame, pd.DataFrame, {"a": [1]}))) + _refused(_load_isolated(path), "_reconstructor") + + +# --------------------------------------------------------------------------- +# what pycisTopic writes still reads + + +@pytest.mark.parametrize("protocol", [0, 1, 2, 3, 4, 5]) +def test_real_cistopic_pickles_read_under_every_protocol(tmp_path, protocol): + truth = write_cistopic_pickle(tmp_path / "obj.pkl", harmony=True) + _write(truth.path, truth.obj, protocol=protocol) + obj = pickled.load(truth.path) + assert np.array_equal(obj.fragment_matrix.toarray(), truth.obj.fragment_matrix.toarray()) + pd.testing.assert_frame_equal(obj.cell_data, truth.obj.cell_data) + pd.testing.assert_frame_equal(obj.selected_model.cell_topic, truth.obj.selected_model.cell_topic) + + imputed = write_imputed_pickle(tmp_path / "imputed.pkl") + _write(imputed.path, imputed.obj, protocol=protocol) + assert np.array_equal(pickled.load(imputed.path).mtx, imputed.obj.mtx) + + +def test_what_numpy_and_pandas_pickle_besides_reads_back_equal(tmp_path): + """Values a pycisTopic object may hold beyond the fixtures: numpy + scalars, big-endian and Fortran-ordered arrays, text and bytes, empty + arrays, nullable columns with missing values (pd.NA), sets.""" + with pd.option_context("future.infer_string", False): + frame = pd.DataFrame({ + "n": pd.array([1, None], dtype="Int64"), + "f": pd.array([0.5, None], dtype="Float64"), + "b": pd.array([True, None], dtype="boolean"), + "s": pd.array(["a", None], dtype=pd.StringDtype("python")), + "c": pd.Categorical(["x", None]), + }, index=pd.Index(["c0", "c1"], name="cell")) + multi = pd.Series([1, 2], index=pd.MultiIndex.from_tuples([("a", 1), ("b", 2)])) + value = { + "scalars": [np.int64(3), np.float32(0.5), np.bool_(True), np.str_("x")], + "big": np.arange(4, dtype=">i4"), + "fortran": np.asfortranarray(np.arange(6.0).reshape(2, 3)), + "text": np.array(["a", "bb"]), "bytes": np.array([b"a", b"bb"]), + "objects": np.array(["a", 1, None], dtype=object), + "empty": np.zeros((0, 3)), + "frame": frame, + "multi": multi, + "sets": ({1, 2}, frozenset({3})), + } + for protocol in (0, 2, 4, 5): + got = pickled.load(_write(tmp_path / f"p{protocol}.pkl", value, protocol=protocol)) + expected = pickle.loads(pickle.dumps(value, protocol=protocol)) # numpy's own reading + assert got["scalars"] == value["scalars"] + for key in ("big", "fortran", "text", "bytes", "objects", "empty"): + assert got[key].dtype == expected[key].dtype and np.array_equal(got[key], value[key]), key + assert got["fortran"].flags.f_contiguous + pd.testing.assert_frame_equal(got["frame"], frame) + pd.testing.assert_series_equal(got["multi"], value["multi"]) + assert got["sets"] == value["sets"] + + +def test_a_pandas_15_pickle_reads_with_pd_na_in_its_text(tmp_path): + """pandas 1.5 (pycisTopic's) pickles a nullable text column's missing + values as `pandas._libs.missing.NA`, which the allowlist lacked.""" + data = _assemble(_global("pandas._libs.missing", "NA")) + assert pickled.load(_bytes(tmp_path / "na.pkl", data)) is pd.NA diff --git a/tests/test_xenium.py b/tests/test_xenium.py index 5f80847..68ad421 100644 --- a/tests/test_xenium.py +++ b/tests/test_xenium.py @@ -141,6 +141,48 @@ def test_the_analysis_is_read_from_the_archive_xenium_ships_it_in(tmp_path): np.testing.assert_allclose(adata.obsm["X_umap"], truth.umap) +@pytest.mark.parametrize("suffix", [".h5ad", ".zarr"]) +def test_a_packed_analysis_converts_into_a_folder_that_does_not_exist_yet(tmp_path, suffix): + """F112: the archive was unpacked beside the output while planning, + before anything had made the output's folder, so the conversion failed + with '[Errno 2] No such file or directory: .../.adata-xenium-...', a + path the user never chose. Other conversions make the folder, and this + one now does too: the archive is unpacked in the nearest folder that + exists, and nothing is left there.""" + truth = write_xenium_bundle(tmp_path / "xenium") + shutil.make_archive(str(truth.bundle / "analysis"), "gztar", truth.bundle, "analysis") + shutil.rmtree(truth.bundle / "analysis") + out = tmp_path / "new" / "deeper" / f"out{suffix}" + + result = _ok(truth.bundle, "-o", out) + + assert "analysis: clusters, X_umap" in _out(result) + adata = _read(out) + assert adata.obs["graphclust"].astype(object).fillna("").tolist() == list(truth.clusters) + np.testing.assert_allclose(adata.obsm["X_umap"], truth.umap) + assert not list(tmp_path.rglob(".adata-xenium-*")), "the unpacked tables are left behind" + + +def test_an_archive_that_cannot_be_unpacked_beside_the_output_is_refused_by_name(tmp_path, monkeypatch): + """F112: a folder that cannot take the unpacked tables is named, with + why, rather than as a bare Errno about a temporary path.""" + import tempfile + + truth = write_xenium_bundle(tmp_path / "xenium") + shutil.make_archive(str(truth.bundle / "analysis"), "gztar", truth.bundle, "analysis") + shutil.rmtree(truth.bundle / "analysis") + + def unwritable(*args, **kwargs): + raise PermissionError(13, "Permission denied", str(kwargs.get("dir"))) + + monkeypatch.setattr(tempfile, "TemporaryDirectory", unwritable) + out = tmp_path / "new" / "out.h5ad" + message = _refused(truth.bundle, "-o", out) + assert "Cannot unpack analysis.tar.gz in" in message and "Permission denied" in message + assert "Errno" not in message + assert not out.parent.exists() + + def test_an_archive_cannot_write_outside_its_temporary_folder(tmp_path): """Members are copied to paths built from the names this wants, never extracted where the archive says."""