diff --git a/cellpy/batch/__init__.py b/cellpy/batch/__init__.py index ad026e35..a1b24ad3 100644 --- a/cellpy/batch/__init__.py +++ b/cellpy/batch/__init__.py @@ -30,6 +30,14 @@ parse_argument, resolve_specs, ) +from cellpy.batch.result import ( + BatchLoadError, + BatchResult, + CellOutcome, + CellResult, +) +from cellpy.batch.runner import load_cell, run +from cellpy.batch.store import CellStore __all__ = [ "Journal", @@ -45,4 +53,11 @@ "CellSpec", "resolve_specs", "parse_argument", + "CellResult", + "BatchResult", + "CellOutcome", + "BatchLoadError", + "load_cell", + "run", + "CellStore", ] diff --git a/cellpy/batch/result.py b/cellpy/batch/result.py new file mode 100644 index 00000000..a1a5543f --- /dev/null +++ b/cellpy/batch/result.py @@ -0,0 +1,97 @@ +"""Batch run results (batch v3, #700). + +"Errors are data": a batch run returns a :class:`BatchResult` with a per-cell +outcome, timing and captured exception, instead of the legacy mix of printing, +partial ``errors`` lists and aborting. Printing/raising is the caller's policy +(``result.raise_if_failed()`` for strict mode). +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from enum import Enum +from typing import Any, Iterator + +import polars as pl + +from cellpy.exceptions import CellpyError + + +class CellOutcome(str, Enum): + LOADED = "loaded" + FAILED = "failed" + SKIPPED = "skipped" + + +class BatchLoadError(CellpyError): + """Raised by :meth:`BatchResult.raise_if_failed` when cells failed.""" + + +@dataclass +class CellResult: + """The outcome of loading one cell.""" + + label: str + outcome: CellOutcome + cell: Any | None = None + source: str | None = None # "cellpy" | "raw" | None + seconds: float = 0.0 + error: BaseException | None = None + + @property + def ok(self) -> bool: + return self.outcome == CellOutcome.LOADED + + +@dataclass +class BatchResult: + """The outcome of a batch run: one :class:`CellResult` per cell.""" + + results: list[CellResult] = field(default_factory=list) + + def __len__(self) -> int: + return len(self.results) + + def __iter__(self) -> Iterator[CellResult]: + return iter(self.results) + + def __getitem__(self, label: str) -> CellResult: + for result in self.results: + if result.label == label: + return result + raise KeyError(label) + + @property + def loaded(self) -> list[CellResult]: + return [r for r in self.results if r.outcome == CellOutcome.LOADED] + + @property + def failed(self) -> list[CellResult]: + return [r for r in self.results if r.outcome == CellOutcome.FAILED] + + @property + def skipped(self) -> list[CellResult]: + return [r for r in self.results if r.outcome == CellOutcome.SKIPPED] + + def cells(self) -> dict[str, Any]: + """Mapping of label -> loaded cell, successful cells only.""" + return {r.label: r.cell for r in self.loaded} + + def raise_if_failed(self) -> "BatchResult": + """Strict mode: raise if any cell failed; otherwise return self.""" + if self.failed: + labels = ", ".join(r.label for r in self.failed) + raise BatchLoadError(f"{len(self.failed)} cell(s) failed to load: {labels}") + return self + + def report(self) -> pl.DataFrame: + """A tidy per-cell outcome frame (the dataframe ``errors`` only hinted at).""" + return pl.DataFrame( + { + "cell": [r.label for r in self.results], + "outcome": [r.outcome.value for r in self.results], + "source": [r.source for r in self.results], + "seconds": [r.seconds for r in self.results], + "error": [None if r.error is None else str(r.error) for r in self.results], + } + ) diff --git a/cellpy/batch/runner.py b/cellpy/batch/runner.py new file mode 100644 index 00000000..0df33b22 --- /dev/null +++ b/cellpy/batch/runner.py @@ -0,0 +1,110 @@ +"""Batch runner (batch v3, #700). + +Per-cell work is a pure function -- one cell in, one result out, no shared +mutable state. Serial vs parallel execution is then a choice of executor, not a +second 300-line method (the legacy ``update`` / ``parallel_update`` clone). +This arc ships the serial executor; the process pool is A8 (#704). +""" + +from __future__ import annotations + +import time +from typing import Any, Callable + +from cellpy import get as _cellpy_get +from cellpy.batch.journal import Journal +from cellpy.batch.policy import CellSpec, LoadPolicy, SourcePreference, resolve_specs +from cellpy.batch.result import BatchResult, CellOutcome, CellResult + +ProgressHook = Callable[[int, int, CellResult], None] + + +def _get_kwargs(spec: CellSpec, policy: LoadPolicy) -> tuple[dict, str | None]: + """Map a resolved :class:`CellSpec` + policy onto ``cellpy.get`` kwargs. + + Returns the kwargs and the source label ("cellpy"/"raw"/None) we expect. + """ + kwargs: dict[str, Any] = { + "mass": spec.mass, + "nominal_capacity": spec.nom_cap, + "area": spec.area, + "cycle_mode": spec.cycle_mode, + "instrument": spec.instrument, + "model": spec.model, + "selector": policy.selector, + } + + raw = spec.raw_files or None + if policy.source is SourcePreference.RAW_ONLY: + kwargs["filename"] = raw + source = "raw" if raw else None + elif policy.source is SourcePreference.CELLPY_ONLY: + kwargs["cellpy_file"] = spec.cellpy_file + source = "cellpy" if spec.cellpy_file else None + else: # AUTO + kwargs["cellpy_file"] = spec.cellpy_file + kwargs["filename"] = raw + source = "cellpy" if spec.cellpy_file else ("raw" if raw else None) + + kwargs = {key: val for key, val in kwargs.items() if val is not None} + kwargs.update(policy.loader_kwargs) + return kwargs, source + + +def load_cell(spec: CellSpec, policy: LoadPolicy | None = None) -> CellResult: + """Load one cell from its resolved spec. Pure-ish: no prints, no mutation. + + Returns a :class:`CellResult` carrying the cell or the exception; only + re-raises when ``policy.accept_errors`` is False. + """ + policy = policy or LoadPolicy() + kwargs, source = _get_kwargs(spec, policy) + + started = time.perf_counter() + try: + cell = _cellpy_get(**kwargs) + except Exception as error: # noqa: BLE001 - errors are data (accept_errors) + if not policy.accept_errors: + raise + return CellResult( + label=spec.label, + outcome=CellOutcome.FAILED, + source=source, + seconds=time.perf_counter() - started, + error=error, + ) + return CellResult( + label=spec.label, + outcome=CellOutcome.LOADED, + cell=cell, + source=source, + seconds=time.perf_counter() - started, + ) + + +def run( + journal: Journal, + policy: LoadPolicy | None = None, + per_cell: dict | None = None, + on_progress: ProgressHook | None = None, +) -> BatchResult: + """Load every cell in ``journal`` (serial), returning a :class:`BatchResult`. + + Progress is reported via the ``on_progress`` callback; the runner never + imports tqdm or prints. + """ + policy = policy or LoadPolicy() + specs = resolve_specs(journal, policy, per_cell) + bad = set(journal.session.get("bad_cells") or []) if policy.skip_bad_cells else set() + + results: list[CellResult] = [] + total = len(specs) + for index, spec in enumerate(specs, start=1): + if spec.label in bad: + result = CellResult(spec.label, CellOutcome.SKIPPED, source=None) + else: + result = load_cell(spec, policy) + results.append(result) + if on_progress is not None: + on_progress(index, total, result) + return BatchResult(results) diff --git a/cellpy/batch/store.py b/cellpy/batch/store.py new file mode 100644 index 00000000..7ff96045 --- /dev/null +++ b/cellpy/batch/store.py @@ -0,0 +1,76 @@ +"""Lazy cell store (batch v3, #700). + +Replaces ``batch_core.Data`` + ``experiment.cell_data_frames`` + the ``x_`` +prefixed accessor dict. A standard ``Mapping`` with lazy loading; tab completion +comes from ``_ipython_key_completions_`` (the supported mechanism for +``store[""]``) instead of prefix-mangled attribute names -- which also +removes the ``str.lstrip`` label-mangling bug (batch_core.py:180, where a cell +named ``xenon_cell`` round-tripped to ``enon_cell``). +""" + +from __future__ import annotations + +from collections.abc import Callable, Mapping +from typing import Any, Iterator + + +class CellStore(Mapping): + """A lazy ``Mapping[str, CellpyCell]``. + + Construct with per-label zero-argument loaders (called on first access) or + with already-loaded cells (:meth:`from_cells`). Loaded cells are cached. + """ + + def __init__( + self, + loaders: Mapping[str, Callable[[], Any]] | None = None, + cells: Mapping[str, Any] | None = None, + ) -> None: + self._loaders: dict[str, Callable[[], Any]] = dict(loaders or {}) + self._cache: dict[str, Any] = dict(cells or {}) + # preserve insertion order, loaders first then any cache-only labels + self._labels: list[str] = list(self._loaders) + for label in self._cache: + if label not in self._loaders: + self._labels.append(label) + + @classmethod + def from_cells(cls, cells: Mapping[str, Any]) -> "CellStore": + """Build a store over already-loaded cells (e.g. from a BatchResult).""" + return cls(cells=cells) + + def __getitem__(self, label: str) -> Any: + if label in self._cache: + return self._cache[label] + if label in self._loaders: + cell = self._loaders[label]() + self._cache[label] = cell + return cell + raise KeyError(label) + + def __iter__(self) -> Iterator[str]: + return iter(self._labels) + + def __len__(self) -> int: + return len(self._labels) + + def first(self) -> Any: + """Load and return the first cell.""" + if not self._labels: + raise KeyError("no cells in store") + return self[self._labels[0]] + + def sample(self) -> Any: + """Alias for :meth:`first` (a representative cell).""" + return self.first() + + def is_loaded(self, label: str) -> bool: + return label in self._cache + + def unload(self, label: str) -> None: + """Drop a loaded cell from the cache (explicit memory management).""" + self._cache.pop(label, None) + + def _ipython_key_completions_(self) -> list[str]: + """Tab completion for ``store[""]`` -- no prefix mangling.""" + return list(self._labels) diff --git a/tests/test_batch_v3_runner.py b/tests/test_batch_v3_runner.py new file mode 100644 index 00000000..d48716e4 --- /dev/null +++ b/tests/test_batch_v3_runner.py @@ -0,0 +1,141 @@ +"""Tests for batch v3 runner/result/store (#700).""" + +import polars as pl +import pytest + +from cellpy.batch import ( + BatchLoadError, + BatchResult, + CellOutcome, + CellResult, + CellSpec, + CellStore, + Journal, + LoadPolicy, + SourcePreference, + load_cell, + run, +) +from cellpy.batch.journal import FILENAME + + +# ---- result.py ---------------------------------------------------------- + + +def _mk_results(): + return BatchResult( + [ + CellResult("a", CellOutcome.LOADED, cell=object(), source="cellpy", seconds=0.1), + CellResult("b", CellOutcome.FAILED, error=ValueError("boom"), seconds=0.2), + CellResult("c", CellOutcome.SKIPPED), + ] + ) + + +def test_batchresult_partitions_and_report(): + br = _mk_results() + assert [r.label for r in br.loaded] == ["a"] + assert [r.label for r in br.failed] == ["b"] + assert [r.label for r in br.skipped] == ["c"] + assert set(br.cells()) == {"a"} + assert br["b"].error.args[0] == "boom" + + rep = br.report() + assert isinstance(rep, pl.DataFrame) + assert rep.height == 3 + assert set(rep.columns) == {"cell", "outcome", "source", "seconds", "error"} + assert rep.filter(pl.col("cell") == "b")["error"].item() == "boom" + + +def test_raise_if_failed(): + with pytest.raises(BatchLoadError, match="b"): + _mk_results().raise_if_failed() + ok = BatchResult([CellResult("a", CellOutcome.LOADED)]) + assert ok.raise_if_failed() is ok + + +# ---- store.py (incl. the lstrip bug fix) -------------------------------- + + +def test_cellstore_is_lazy(): + calls = [] + + def make(label): + return lambda: (calls.append(label), f"cell::{label}")[1] + + store = CellStore({"a": make("a"), "b": make("b")}) + assert list(store) == ["a", "b"] + assert len(store) == 2 + assert calls == [] # nothing loaded yet + assert store["a"] == "cell::a" + assert calls == ["a"] + assert store["a"] == "cell::a" # cached, not reloaded + assert calls == ["a"] + assert store.is_loaded("a") and not store.is_loaded("b") + store.unload("a") + assert not store.is_loaded("a") + + +def test_cellstore_no_label_mangling(): + """The legacy x_/lstrip accessor turned 'xenon_cell' into 'enon_cell'.""" + store = CellStore.from_cells({"xenon_cell": "X", "x_ray": "R"}) + assert set(store) == {"xenon_cell", "x_ray"} + assert store["xenon_cell"] == "X" + assert store["x_ray"] == "R" + assert set(store._ipython_key_completions_()) == {"xenon_cell", "x_ray"} + + +# ---- runner.py (integration with cellpy.get) ---------------------------- + + +def _one_cell_journal(label, cellpy_file): + return Journal( + name="t", + project="p", + pages=pl.DataFrame({FILENAME: [label], "cellpy_file_name": [str(cellpy_file)]}), + ) + + +def test_load_cell_from_cellpy_file(parameters): + spec = CellSpec(label="c45", cellpy_file=parameters.cellpy_file_path) + result = load_cell(spec, LoadPolicy(source=SourcePreference.CELLPY_ONLY)) + assert result.ok + assert result.source == "cellpy" + assert result.cell is not None + assert result.seconds >= 0 + + +def test_load_cell_error_is_captured(tmp_path): + spec = CellSpec(label="bad", cellpy_file=str(tmp_path / "nope.h5")) + result = load_cell(spec, LoadPolicy(source=SourcePreference.CELLPY_ONLY)) + assert result.outcome == CellOutcome.FAILED + assert result.error is not None + + +def test_load_cell_reraises_when_not_accepting(tmp_path): + spec = CellSpec(label="bad", cellpy_file=str(tmp_path / "nope.h5")) + with pytest.raises(Exception): + load_cell( + spec, + LoadPolicy(source=SourcePreference.CELLPY_ONLY, accept_errors=False), + ) + + +def test_run_over_journal(parameters): + j = _one_cell_journal("c45", parameters.cellpy_file_path) + seen = [] + br = run( + j, + LoadPolicy(source=SourcePreference.CELLPY_ONLY), + on_progress=lambda i, n, r: seen.append((i, n, r.label)), + ) + assert len(br) == 1 + assert br["c45"].ok + assert seen == [(1, 1, "c45")] + + +def test_run_skips_bad_cells(parameters): + j = _one_cell_journal("c45", parameters.cellpy_file_path) + j.session["bad_cells"] = ["c45"] + br = run(j, LoadPolicy(source=SourcePreference.CELLPY_ONLY, skip_bad_cells=True)) + assert br["c45"].outcome == CellOutcome.SKIPPED