diff --git a/CLAUDE.md b/CLAUDE.md index 0c03683..ef8c829 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -25,7 +25,7 @@ Or directly via `py src/main.py [variable] [options]` (backward-compat Where `` is a **path relative to `PSC_PLOT_DATA_DIR`, ending in the file prefix** — e.g. `pfd` (data root itself) or `run5/pfd` (subdirectory). The prefix part selects the data source: field prefixes (`pfd`, `pfd_moments`, `gauss`, `continuity`) or particle prefixes (`prt`). `Prepath` is just a `str` alias (`src/lib/file_util.py`); `split_prepath()` splits it into `(subdir, prefix)`. Examples live in `plots/check.sh` and `plots/check2.sh` and serve as the de-facto smoke tests / usage reference. Common flags: -- `-q` quiet (don't show interactively), `-s [dir]` save output (defaults to `.` if dir omitted) +- `-q` quiet (don't show interactively), `-s`/`--save` save output — each argument is either a path fragment `[dir/][stem][.ext]` or one of `dir=`/`name=`/`format=`; bare `-s` means cwd, a stem derived from `name_fragments`, and the default format. A directory fragment must end in `/` (or use `dir=`), otherwise it's taken as the filename stem. `--save-dpi` is deliberately separate, pending a `--dpi` that would apply to live figures too - `-w`/`--with [prepath::][var_key]` switch the active dataset and/or variable mid-pipeline (loads the prepath if not already in the world); `--copy [new=old|old]` duplicate a variable + its metadata; `-c`/`--compute` force full computation - Adaptor flags such as `--roll`, `--reduce`, `--bin`, `--scatter [name]`, `--mag`, `--nan0`, `--scale`, `--fourier`/`-f`, `--pos`, `--species`, `--transform-spherical`, `-v` (versus, sets axes), `-b` (bin) @@ -58,7 +58,7 @@ pytest --mpl The figure tests live in `tests/test_plots.py` and use `pytest-mpl` for image comparison against baseline PNGs in `tests/baseline/`. Each test runs the full CLI pipeline (parsing → loading → adaptors → plotting) against small test datasets in `tests/data/` (test-2d: x=1,y=8,z=16; test-3d: x=4,y=4,z=4). -Other suites (plain `pytest`, no `--mpl` needed): `test_save.py`/`test_save_filename.py` (save formats and the `name_fragments`-derived filename), `test_idx_efficient.py` (asserts `Idx` prunes dask partitions instead of filtering), `test_particle_bp_vs_h5.py` + `test_h5_species_discovery.py` (loader parity/species discovery), `test_memory.py`, `test_particle_bp_perf.py`, `test_dask_graph.py`, and `test_synthetic_particles.py` (with the `synthetic_particles.py` generator). +Other suites (plain `pytest`, no `--mpl` needed): `test_save.py` (everything about saving, in three sections: the `--save` argument grammar via `parse_save`, the `name_fragments`-derived filename, and end-to-end pipeline saves), `test_parse_util.py`/`test_parse_errors.py` (parsing helpers and argparse error rendering), `test_idx_efficient.py` (asserts `Idx` prunes dask partitions instead of filtering), `test_particle_bp_vs_h5.py` + `test_h5_species_discovery.py` (loader parity/species discovery), `test_memory.py`, `test_particle_bp_perf.py`, `test_dask_graph.py`, and `test_synthetic_particles.py` (with the `synthetic_particles.py` generator). To regenerate baselines after intentional visual changes: @@ -118,7 +118,7 @@ The code lives under `src/lib/` and is organized around **loaders** (sources), * 2. Everything runs through a lazy, memoized **node graph** (`src/lib/data/node.py`). `compile_action_nodes` chains: `RootNode(config)` → `AdaptorNode(With(args.prepath, args.variable))` → one `AdaptorNode` per user adaptor → `PlotNode(hooks)` → action node(s): `ShowPlotNode` (`--show`, default on), `SavePlotNode` (`--save`), and/or `DaskGraphNode` (`--dask-graph`, renders the dask graph as SVG instead of plotting). Note the positional args are just an implicit **`--with`** — loading is `With`'s job, not a dedicated loader node. Each `DataProcessingNode` has a `@cache`d `pull()` and accumulates `name_fragments` (used to derive save filenames). If the adaptor list contains no `Versus`, a default `Versus(["y","z"], time_dim_rule="guess", color_dim=None)` is appended (`_with_versus`) — this is what selects axes/color/time dims and appends the plot target. 3. Data flows as a **`DataWorld`** (`src/lib/data/data_world.py`): a frozen dataclass holding `datas: dict[Prepath, DataWithAttrs]`, an `active_prepath: Prepath | None`, `plot_targets: list[PlotTarget]`, and `config`. `RootNode.pull()` returns an empty world; each `AdaptorNode` applies a `WorldAdaptor` to it. `active_data` returns `None` when there's no active prepath, `require_active_data()` raises; mutate with `with_active(data=…, prepath=…)` (sets/replaces the active entry) or `with_data(prepath, data)` (inserts without activating — what `Loader.apply_world` does). A `__post_init__` assert enforces that `active_prepath` is a key of `datas`. Holding **multiple named datas + multiple plot targets** is the "split vars" capability this branch is named for. 4. `PlotNode.pull()` calls `get_plot(world)` (`src/lib/plotting/get_plot.py`), which builds one `Renderer` **per `PlotTarget`**: `Field1dRenderer` when the target has no `color_dim`; `PolarFieldRenderer` / `Field2dRenderer` chosen by `SpatialDimsRTheta` / `SpatialDimsXY`; `ScatterRenderer` for `List` data with exactly 2 spatial dims. It wraps the renderer **list** in `StaticPlot` or `AnimatedPlot` based on `n_frames = max(r.get_n_frames())` — **animated iff `n_frames > 1`**, not on `time_dim` presence. Hooks are then attached via `plot.add_hook()`. -5. `ShowPlotNode`/`SavePlotNode.pull()` call `plot.show()` / `plot.save_to_path()`, which lazily `_initialize()` the figure (see the PlotInfo/AxesManager layer below). +5. `ShowPlotNode`/`SavePlotNode.pull()` call `plot.show()` / `plot.save_to_path()`, which lazily `_initialize()` the figure (see the PlotInfo/AxesManager layer below). `SavePlotNode` carries a `SaveSpec(dir, name, format)` (`src/lib/parsing/parse_save.py`) — the parsed `--save` arguments, each `None` when the user didn't specify it — and resolves the defaults at pull time: `dir` → cwd, `name` → `get_save_file_stem()`, `format` → `plot.default_save_format()`, warning and falling back when the requested format isn't in `plot.allowed_save_formats()`. **Resolution lives in the node, not on `SaveSpec`**, because picking a default format needs the `Plot`; that's also why an unwritable extension isn't rejected at parse time, and why `--save` accepts any extension. An explicit `name=` is used verbatim — only the derived stem goes through `sanitize_stem`. ### PlotTarget diff --git a/src/lib/data/adaptors/bin.py b/src/lib/data/adaptors/bin.py index e6ef969..e8abf12 100644 --- a/src/lib/data/adaptors/bin.py +++ b/src/lib/data/adaptors/bin.py @@ -7,81 +7,82 @@ from lib import var_info_registry from lib.data.adaptor import MetadataAdaptor -from lib.data.data_with_attrs import Field, FieldMetadata, List +from lib.data.data_with_attrs import Field, FieldMetadata, LazyList, List +from lib.data.types import VarKey from lib.parsing import parse_util from lib.parsing.args_registry import arg_parser -def _guess_bin_edgess(data: List, varname_to_nbins: dict[str, int | None]) -> list: - varname_to_edges: dict[str, np.array] = {} +def _guess_bin_edgess(data: List, keys_to_nbins: dict[VarKey, int | None]) -> list: + keys_to_edges: dict[VarKey, np.ndarray] = {} - compute_varnames = [] + compute_keys: list[VarKey] = [] mins_to_compute = [] maxs_to_compute = [] nbins_so_far = 1 - varnames_with_missing_nbins = [] + keys_with_missing_nbins: list[VarKey] = [] df = data.data # Calculate edges using metadata when possible - for varname, nbins in varname_to_nbins.items(): - if varname in data.coordss(): - coords = data.coordss()[varname] + for key, nbins in keys_to_nbins.items(): + if key in data.coordss(): + coords = data.coordss()[key] if nbins is None: nbins = len(coords) # note: use inf as right edge for convenience; it gets sliced out later - varname_to_edges[varname] = np.concat((coords, [np.inf])) + keys_to_edges[key] = np.concatenate((coords, [np.inf])) else: - varname_to_edges[varname] = np.linspace(coords[0], coords[-1] + coords[1] - coords[0], nbins + 1, endpoint=True) - elif varname in data.metadata.var_infos and (data.metadata.var_infos[varname].geometry == "polar:theta" or data.metadata.var_infos[varname].geometry == "spherical:phi"): - varname_to_edges[varname] = np.linspace(-np.pi, np.pi, nbins + 1, endpoint=True) - elif varname in data.metadata.var_infos and data.metadata.var_infos[varname].geometry == "spherical:theta": - varname_to_edges[varname] = np.linspace(0.0, np.pi, nbins + 1, endpoint=True) + keys_to_edges[key] = np.linspace(coords[0], coords[-1] + coords[1] - coords[0], nbins + 1, endpoint=True) + elif key in data.metadata.var_infos and (data.metadata.var_infos[key].geometry == "polar:theta" or data.metadata.var_infos[key].geometry == "spherical:phi"): + keys_to_edges[key] = np.linspace(-np.pi, np.pi, nbins + 1, endpoint=True) + elif key in data.metadata.var_infos and data.metadata.var_infos[key].geometry == "spherical:theta": + keys_to_edges[key] = np.linspace(0.0, np.pi, nbins + 1, endpoint=True) else: - compute_varnames.append(varname) - mins_to_compute.append(df[varname].min()) - maxs_to_compute.append(df[varname].max()) + compute_keys.append(key) + mins_to_compute.append(df[key].min()) + maxs_to_compute.append(df[key].max()) if nbins: nbins_so_far *= nbins else: - varnames_with_missing_nbins.append(varname) + keys_with_missing_nbins.append(key) # If needed, batch-compute the missing edges - if compute_varnames: + if compute_keys: if isinstance(df, dd.DataFrame): computed_mins, computed_maxs = dask.array.compute(mins_to_compute, maxs_to_compute) else: computed_mins, computed_maxs = mins_to_compute, maxs_to_compute - if varnames_with_missing_nbins: + if keys_with_missing_nbins: # split bins evenly across remaining dimensions n_data = len(df) mean_n_data_per_bin_so_far = n_data / nbins_so_far target_mean_n_data_per_bin = 10 # arbitrary number - guessed_nbins = math.ceil((mean_n_data_per_bin_so_far / target_mean_n_data_per_bin) ** (1 / len(compute_varnames))) - for varname_with_missing_nbins in varnames_with_missing_nbins: - varname_to_nbins[varname_with_missing_nbins] = guessed_nbins + guessed_nbins = math.ceil((mean_n_data_per_bin_so_far / target_mean_n_data_per_bin) ** (1 / len(compute_keys))) + for varname_with_missing_nbins in keys_with_missing_nbins: + keys_to_nbins[varname_with_missing_nbins] = guessed_nbins - for varname, min, max in zip(compute_varnames, computed_mins, computed_maxs): - nbins = varname_to_nbins[varname] - varname_to_edges[varname] = np.linspace(min, max, nbins + 1, endpoint=True) + for key, min, max in zip(compute_keys, computed_mins, computed_maxs): + nbins = keys_to_nbins[key] + keys_to_edges[key] = np.linspace(min, max, nbins + 1, endpoint=True) # ensure edges are in same order as bin values - edgess = [varname_to_edges[varname] for varname in varname_to_nbins] + edgess = [keys_to_edges[key] for key in keys_to_nbins] return edgess class Bin(MetadataAdaptor): - def __init__(self, varname_to_nbins: dict[str, int | None]): - self.varname_to_nbins = varname_to_nbins + def __init__(self, key_to_nbins: dict[VarKey, int | None]): + self.keys_to_nbins = key_to_nbins def apply_field(self, data: Field) -> Field: dim_names_to_bin_size = {} - for dim_name, nbins in self.varname_to_nbins.items(): + for dim_name, nbins in self.keys_to_nbins.items(): if not nbins: continue @@ -96,48 +97,47 @@ def apply_field(self, data: Field) -> Field: return data.with_active(data=data.require_active_subdata().coarsen(dim_names_to_bin_size, boundary="pad").mean()) def apply_list(self, data: List) -> Field: - bin_edgess = _guess_bin_edgess(data, self.varname_to_nbins) + bin_edgess = _guess_bin_edgess(data, self.keys_to_nbins) - df = data.data - if isinstance(df, dd.DataFrame): + if isinstance(data, LazyList): binned_data, _ = dask.array.histogramdd( - [df[active_key].to_dask_array() for active_key in self.varname_to_nbins], + [data[key].to_dask_array() for key in self.keys_to_nbins], bin_edgess, density=False, - weights=df[data.metadata.weight_key].to_dask_array() if data.metadata.weight_key else None, + weights=data[data.metadata.weight_key].to_dask_array() if data.metadata.weight_key else None, ) else: binned_data, _ = np.histogramdd( - [df[active_key] for active_key in self.varname_to_nbins], + [data[key] for key in self.keys_to_nbins], bin_edgess, density=False, - weights=df[data.metadata.weight_key] if data.metadata.weight_key else None, + weights=data[data.metadata.weight_key] if data.metadata.weight_key else None, ) # note: the slice removes any infs - coords = dict(zip(self.varname_to_nbins.keys(), (edges[:-1] for edges in bin_edgess))) + coords = dict(zip(self.keys_to_nbins.keys(), (edges[:-1] for edges in bin_edgess))) da = xr.DataArray( binned_data, coords, - dims=self.varname_to_nbins.keys(), + dims=self.keys_to_nbins.keys(), ) - f_dim = var_info_registry.lookup("prt", "f") + f_info = var_info_registry.lookup("prt", "f") subject = data.metadata.subject if subject is not None and subject.latex == r"\text{Ions}": - f_dim = f_dim.assign(display=f_dim.display.latex + r"_\text{i}") + f_info = f_info.assign(display=f_info.display.latex + r"_\text{i}") elif subject is not None and subject.latex == r"\text{Electrons}": - f_dim = f_dim.assign(display=f_dim.display.latex + r"_\text{e}") + f_info = f_info.assign(display=f_info.display.latex + r"_\text{e}") new_var_infos = {key: data.metadata.var_infos[key] for key in da.coords if key in data.metadata.var_infos} # want: psc-plot prt.i --derive K="ux^2+uy^2+uz^2" --bin K=128 -v K --scale log - new_var_infos["f"] = f_dim + new_var_infos["f"] = f_info return Field({"f": da}, FieldMetadata.create_from(data.metadata, active_key="f", var_infos=new_var_infos)) def get_name_fragments(self) -> list[str]: - subfrags = "_".join(f"{varname}={nbins}" if nbins else varname for varname, nbins in self.varname_to_nbins.items()) + subfrags = "_".join(f"{varname}={nbins}" if nbins else varname for varname, nbins in self.keys_to_nbins.items()) return [f"bin_{subfrags}"] @@ -152,7 +152,7 @@ def get_name_fragments(self) -> list[str]: nargs="+", ) def parse_bin(args: list[str]) -> Bin: - varname_to_nbins = {} + keys_to_nbins = {} insert_bin_t = True for arg in args: @@ -160,22 +160,22 @@ def parse_bin(args: list[str]) -> Bin: if len(split_arg) == 2 and not split_arg[1]: # arg is "t=", i.e., disable implicit binning along t - parse_util.parse_value(split_arg[0], "active_key", ["t"]) + parse_util.parse_value(split_arg[0], "var_key", ["t"]) insert_bin_t = False continue elif len(split_arg) > 2: parse_util.fail_format(arg, _BIN_FORMAT) - [active_key, nbins_arg, *_] = split_arg + [""] + [key, nbins_arg, *_] = split_arg + [""] - parse_util.parse_identifier(active_key, "active_key") + parse_util.parse_identifier(key, "var_key") nbins = parse_util.parse_optional_number(nbins_arg, "nbins", int) - varname_to_nbins[active_key] = nbins - if active_key == "t": + keys_to_nbins[key] = nbins + if key == "t": insert_bin_t = False if insert_bin_t: - varname_to_nbins["t"] = None + keys_to_nbins["t"] = None - return Bin(varname_to_nbins) + return Bin(keys_to_nbins) diff --git a/src/lib/data/compile.py b/src/lib/data/compile.py index 6743e30..3c4ee7f 100644 --- a/src/lib/data/compile.py +++ b/src/lib/data/compile.py @@ -42,28 +42,17 @@ def compile_action_nodes(args: Args, config: PscPlotConfig) -> list[DataProcessi action_nodes = [] if args.dask_graph: - action_nodes.append(DaskGraphNode(plot_node.input_node, save_dir=args.save, show=args.show)) + action_nodes.append(DaskGraphNode(plot_node.input_node, save=args.save, show=args.show)) return action_nodes if args.show: action_nodes.append(ShowPlotNode(plot_node)) - if args.save is None and args.save_format: - print("error: --save-format requires --save", file=sys.stderr) - sys.exit(1) - - if args.save_format == "mp4" and not config.ffmpeg_bin: - print("error: --save-format mp4 requires ffmpeg", file=sys.stderr) - sys.exit(1) - if args.save is not None: - action_nodes.append( - SavePlotNode( - plot_node, - save_dir=args.save, - save_format=args.save_format, - save_dpi=args.save_dpi, - ) - ) + if args.save.format == "mp4" and not config.ffmpeg_bin: + print("error: format=mp4 requires ffmpeg", file=sys.stderr) + sys.exit(1) + + action_nodes.append(SavePlotNode(plot_node, save=args.save, save_dpi=args.save_dpi)) return action_nodes diff --git a/src/lib/data/node.py b/src/lib/data/node.py index f4a560e..aefcf82 100644 --- a/src/lib/data/node.py +++ b/src/lib/data/node.py @@ -8,9 +8,10 @@ from lib.config import PscPlotConfig from lib.data.adaptor import Adaptor from lib.data.data_world import DataWorld +from lib.parsing.parse_save import SaveSpec from lib.plotting.get_plot import get_plot from lib.plotting.hook import Hook -from lib.plotting.plot import Plot, SaveFormat +from lib.plotting.plot import Plot class DataProcessingNode[D](ABC): @@ -77,20 +78,18 @@ def __init__( self, input_node: DataProcessingNode[Plot], *, - save_dir: Path, - save_format: SaveFormat | None, + save: SaveSpec, save_dpi: float | None, ): super().__init__(input_node.name_fragments) self.input_node = input_node - self.save_dir = save_dir - self.save_format = save_format + self.save = save self.save_dpi = save_dpi def pull(self) -> None: plot = self.input_node.pull() - save_format = self.save_format + save_format = self.save.format if save_format not in plot.allowed_save_formats(): if save_format is not None: message = f"{save_format} is incompatible with the data; reverting to default ({plot.default_save_format()})" @@ -98,8 +97,9 @@ def pull(self) -> None: save_format = plot.default_save_format() - self.save_dir.mkdir(exist_ok=True, parents=True) - path = self.save_dir / f"{self.get_save_file_stem()}.{save_format}" + save_dir = self.save.dir or Path(".") + save_dir.mkdir(exist_ok=True, parents=True) + path = save_dir / f"{self.save.name or self.get_save_file_stem()}.{save_format}" plot.save_to_path(path, dpi=self.save_dpi) print(f"wrote to {path}") @@ -109,12 +109,12 @@ def __init__( self, input_node: DataProcessingNode[DataWorld], *, - save_dir: Path | None, + save: SaveSpec | None, show: bool, ): super().__init__(input_node.name_fragments) self.input_node = input_node - self.save_dir = save_dir or Path.cwd() + self.save = save or SaveSpec() self.show = show def pull(self) -> None: @@ -139,8 +139,10 @@ def pull(self) -> None: import dask - self.save_dir.mkdir(exist_ok=True, parents=True) - path = self.save_dir / f"{self.get_save_file_stem()}.daskgraph.svg" + # save.format is ignored: the extension here is always .daskgraph.svg + save_dir = self.save.dir or Path.cwd() + save_dir.mkdir(exist_ok=True, parents=True) + path = save_dir / f"{self.save.name or self.get_save_file_stem()}.daskgraph.svg" # dask.visualize's optimize_graph flag only lowers legacy HLG collections # (e.g. dask Arrays), not new-style Expr ones (dask DataFrames) — without # pre-optimizing the latter, un-lowered nodes (e.g. Concat from dd.concat) diff --git a/src/lib/parsing/args.py b/src/lib/parsing/args.py index 13a9541..4079536 100644 --- a/src/lib/parsing/args.py +++ b/src/lib/parsing/args.py @@ -1,7 +1,7 @@ import argparse -from pathlib import Path from lib.data.adaptor import Adaptor +from lib.parsing.parse_save import SaveSpec from lib.plotting.hook import Hook @@ -11,7 +11,6 @@ class Args(argparse.Namespace): adaptors: list[Adaptor] hooks: list[Hook] show: bool - save: Path | None - save_format: str | None + save: SaveSpec | None save_dpi: float | None dask_graph: bool diff --git a/src/lib/parsing/args_registry.py b/src/lib/parsing/args_registry.py index 7a62714..eb8f323 100644 --- a/src/lib/parsing/args_registry.py +++ b/src/lib/parsing/args_registry.py @@ -1,24 +1,35 @@ from __future__ import annotations import typing -from argparse import Action, ArgumentParser +from argparse import Action, ArgumentParser, ArgumentTypeError from dataclasses import dataclass from typing import Any, Callable CUSTOM_ARGS: list[ArgparseArgAdder] = [] +def _normalize_values(values: Any) -> list[Any]: + if values is None: + return [] + if isinstance(values, str): + return [values] + return list(values) + + +def _combine(parser: ArgumentParser, action: Action, combiner: typing.Callable[[list[Any]], Any], values: list[Any]) -> Any: + # argparse only turns ArgumentTypeError into a clean error message inside its own + # `type=` conversion (_get_value). One raised here, inside an Action, would escape + # as an uncaught traceback, so route it through parser.error ourselves. + try: + return combiner(values) + except ArgumentTypeError as e: + parser.error(f"argument {'/'.join(action.option_strings)}: {e}") + + def get_combine_args_action(combiner: typing.Callable[[list[Any]], Any]) -> Action: class CombineArgs(Action): def __call__(self, parser, namespace, values, option_string=None): - if values is None: - values = [] - elif isinstance(values, str): - values = [values] - else: - values = list(values) - - combined_value = combiner(values) + combined_value = _combine(parser, self, combiner, _normalize_values(values)) items = getattr(namespace, self.dest, []) items.append(combined_value) @@ -27,6 +38,18 @@ def __call__(self, parser, namespace, values, option_string=None): return CombineArgs +def get_store_combined_args_action(combiner: typing.Callable[[list[Any]], Any]) -> Action: + """Like get_combine_args_action, but stores the combined value instead of appending + it to a list — for single-instance args such as --save. Paired with nargs="*" and + default=None, this distinguishes "flag absent" (None) from "flag with no args".""" + + class StoreCombinedArgs(Action): + def __call__(self, parser, namespace, values, option_string=None): + setattr(namespace, self.dest, _combine(parser, self, combiner, _normalize_values(values))) + + return StoreCombinedArgs + + type ArgparseNArgs = int | typing.Literal["+", "*", "?"] | None type NArgs = ArgparseNArgs | typing.Literal["just one"] diff --git a/src/lib/parsing/parse.py b/src/lib/parsing/parse.py index e8b8e71..b6ae0e1 100644 --- a/src/lib/parsing/parse.py +++ b/src/lib/parsing/parse.py @@ -1,8 +1,8 @@ import argparse -from pathlib import Path from lib.parsing.args import Args -from lib.parsing.args_registry import CUSTOM_ARGS +from lib.parsing.args_registry import CUSTOM_ARGS, get_store_combined_args_action +from lib.parsing.parse_save import SAVE_METAVAR, parse_save def _get_parser() -> argparse.ArgumentParser: @@ -13,21 +13,14 @@ def _get_parser() -> argparse.ArgumentParser: parser.add_argument( "-s", "--save", - action="store", - metavar="dir", - nargs="?", + action=get_store_combined_args_action(parse_save), + dest="save", + metavar=SAVE_METAVAR, + nargs="*", default=None, - const=".", - help="save the figure (to the given dir, if present)", - type=Path, + help="save the figure. Each argument is either a path fragment '[dir/][stem][.ext]' or one of 'dir=', 'name=', 'format='. With no arguments, saves to the current directory using a filename derived from the pipeline and the default format for the data. A bare fragment naming a directory must end in '/' (or use dir=), otherwise it is taken as the filename stem.", ) parser.add_argument("-q", "--quiet", action="store_false", dest="show", help="don't show the figure") - parser.add_argument( - "--save-format", - choices=["mp4", "gif"], - default=None, - help="format for saved animations (default: mp4, falls back to gif if ffmpeg unavailable)", - ) parser.add_argument( "--save-dpi", type=float, diff --git a/src/lib/parsing/parse_save.py b/src/lib/parsing/parse_save.py new file mode 100644 index 0000000..ff585ca --- /dev/null +++ b/src/lib/parsing/parse_save.py @@ -0,0 +1,115 @@ +import argparse +from dataclasses import dataclass +from pathlib import Path + +from lib.parsing import parse_util + + +@dataclass(frozen=True) +class SaveSpec: + """The components of a save path, each None when the user did not specify it. + Resolution of the defaults happens in SavePlotNode, which needs the Plot to + decide a default format.""" + + dir: Path | None = None + name: str | None = None + format: str | None = None + + +_KEYS = ("dir", "name", "format") + +SAVE_METAVAR = "[dir/][stem][.ext] | dir= | name= | format=" + + +def _is_valid_stem(val: str) -> bool: + """A stem must be nonempty and hold at least one non-'.' character, so that + '.', '..' and friends fall through to being parsed as a dir.""" + return bool(val) and any(char != "." for char in val) + + +def _parse_dir(val: str) -> Path: + """Expand a leading '~'. bash leaves it alone after a '=' (dir=~/figs) and + inside quotes, so without this the save dir would be a literal '~'. A path + like './~foo' normalizes to '~foo', which expanduser() rejects outright when + no such user exists — keep it literal there, as os.path.expanduser does.""" + path = Path(val) + try: + return path.expanduser() + except RuntimeError: + return path + + +def _split_fragment(fragment: str) -> tuple[Path | None, str | None, str | None]: + """Parse a bare '[dir/][stem][.ext]' fragment right-to-left, claiming each + component only if it is valid. Returns (dir, stem, ext).""" + rest = fragment + + # ext: text after the last "." that follows the last "/"; valid iff nonempty + ext = None + dot_idx = rest.rfind(".") + if dot_idx > rest.rfind("/"): + candidate = rest[dot_idx + 1 :] + if candidate: + ext = candidate + rest = rest[:dot_idx] + + # stem: text after the last "/"; valid iff _is_valid_stem + name = None + slash_idx = rest.rfind("/") + candidate = rest[slash_idx + 1 :] + if _is_valid_stem(candidate): + name = candidate + rest = rest[: slash_idx + 1] + + # dir: whatever remains; Path() strips the trailing slash + dir = _parse_dir(rest) if rest else None + + return dir, name, ext + + +def parse_save(args: list[str]) -> SaveSpec: + components: dict[str, Path | str] = {} + seen_fragment = False + + def set_component(key: str, value: Path | str): + if key in components: + raise argparse.ArgumentTypeError(f"Expected {key} to be specified at most once; got '{components[key]}' and '{value}'") + components[key] = value + + for arg in args: + key, sep, value = arg.partition("=") + + # An arg is a key=value attempt iff an identifier precedes the first "=". + # Anything else (out/fig=x, a-b=c) is a bare fragment, which is how stems + # containing "=" — as derived stems do — get through without a name= key. + if sep and parse_util.is_identifier(key): + parse_util.parse_value(key, "save key", _KEYS) + + if key == "dir": + if not value: + raise argparse.ArgumentTypeError("Expected dir to be a nonempty path; got ''") + set_component("dir", _parse_dir(value)) + elif key == "name": + if "/" in value or not _is_valid_stem(value): + raise argparse.ArgumentTypeError(f"Expected name to be a stem with no '/' and at least one non-'.' character; got '{value}'") + set_component("name", value) + elif key == "format": + if not value or "." in value or "/" in value: + raise argparse.ArgumentTypeError(f"Expected format to be a nonempty extension with no '.' or '/'; got '{value}'") + set_component("format", value) + else: + raise AssertionError(key) + else: + if seen_fragment: + raise argparse.ArgumentTypeError(f"Expected at most one path fragment; got a second one: '{arg}'") + seen_fragment = True + + dir, name, ext = _split_fragment(arg) + if dir is not None: + set_component("dir", dir) + if name is not None: + set_component("name", name) + if ext is not None: + set_component("format", ext) + + return SaveSpec(**components) diff --git a/src/lib/parsing/parse_util.py b/src/lib/parsing/parse_util.py index b0a7d5f..a16d322 100644 --- a/src/lib/parsing/parse_util.py +++ b/src/lib/parsing/parse_util.py @@ -3,7 +3,7 @@ import typing -def _is_identifier(val: str) -> bool: +def is_identifier(val: str) -> bool: return all(re.match(r"^\w[\d\w]*$", v) for v in val.split(".")) @@ -18,15 +18,17 @@ def parse_value[T](val: T, val_name: str, valid_options: typing.Container[T]) -> def parse_identifier(val: str, val_name: str) -> str: - if not _is_identifier(val): + if not is_identifier(val): raise argparse.ArgumentTypeError(f"Expected {val_name} to be an identifier; got '{val}'") return val def parse_optional_identifier(val: str | None, val_name: str) -> str | None: - if val and not _is_identifier(val): - raise argparse.ArgumentTypeError(f"Expected {val_name} to be an identifier or ''; got '{val}'") - return val + if not val: + return None + if is_identifier(val): + return val + raise argparse.ArgumentTypeError(f"Expected {val_name} to be an identifier or ''; got '{val}'") def check_order[T](lower: T | None, upper: T | None, lower_name: str, upper_name: str): diff --git a/tests/test_parse_errors.py b/tests/test_parse_errors.py new file mode 100644 index 0000000..61f8a59 --- /dev/null +++ b/tests/test_parse_errors.py @@ -0,0 +1,62 @@ +import argparse + +import pytest + +from lib.parsing.args_registry import get_store_combined_args_action +from lib.parsing.parse import parse_args + + +def test_bad_versus_arg_exits_cleanly(): + """A validation failure inside a multi-arg adaptor must render as an argparse + error, not escape as an uncaught ArgumentTypeError traceback.""" + with pytest.raises(SystemExit) as excinfo: + parse_args(["pfd", "hx_fc", "-v", "loc=oops"]) + + assert excinfo.value.code == 2 + + +def test_bad_versus_arg_message_names_the_flag(capsys): + with pytest.raises(SystemExit): + parse_args(["pfd", "hx_fc", "-v", "loc=oops"]) + + err = capsys.readouterr().err + error_line = err.rsplit("error: ", 1)[-1] + assert "-v" in error_line + assert "--versus" in error_line + assert "i,j" in err + + +def _combining_parser(): + parser = argparse.ArgumentParser(prog="t") + parser.add_argument("-s", nargs="*", default=None, dest="save", action=get_store_combined_args_action(tuple)) + return parser + + +def test_store_combined_args_action_distinguishes_absent_bare_and_valued(): + parser = _combining_parser() + assert parser.parse_args([]).save is None + assert parser.parse_args(["-s"]).save == () + assert parser.parse_args(["-s", "a", "b"]).save == ("a", "b") + + +def test_store_combined_args_action_stores_rather_than_appends(): + """Unlike get_combine_args_action, a second occurrence replaces rather than appends.""" + assert _combining_parser().parse_args(["-s", "a", "-s", "b"]).save == ("b",) + + +def test_store_combined_args_action_routes_argument_type_error(capsys): + def combiner(values): + raise argparse.ArgumentTypeError("bad value") + + parser = argparse.ArgumentParser(prog="t") + parser.add_argument("-s", "--save", nargs="*", default=None, dest="save", action=get_store_combined_args_action(combiner)) + + with pytest.raises(SystemExit) as excinfo: + parser.parse_args(["-s", "x"]) + + assert excinfo.value.code == 2 + err = capsys.readouterr().err + error_line = err.rsplit("error: ", 1)[-1] + assert "-s" in error_line + assert "--save" in error_line + assert "bad value" in error_line diff --git a/tests/test_parse_util.py b/tests/test_parse_util.py new file mode 100644 index 0000000..4887d87 --- /dev/null +++ b/tests/test_parse_util.py @@ -0,0 +1,24 @@ +import pytest + +from lib.parsing import parse_util + + +@pytest.mark.parametrize( + "val, expected", + [ + ("dir", True), + ("name", True), + ("a1", True), + ("_x", True), + ("a.b", True), + ("", False), + ("a.", False), + (".a", False), + ("out/fig", False), + ("a-b", False), + ("(a)", False), + ("1a", True), + ], +) +def test_is_identifier(val, expected): + assert parse_util.is_identifier(val) is expected diff --git a/tests/test_save.py b/tests/test_save.py index 0c080e4..0d0d51a 100644 --- a/tests/test_save.py +++ b/tests/test_save.py @@ -1,5 +1,112 @@ +"""Everything about saving a figure, from the innermost layer outward: + +1. `parse_save` — the pure `--save` argument grammar, no data and no pipeline. +2. `get_save_file_stem` — the filename derived from the node graph's `name_fragments`. +3. The whole pipeline — `compile_action_nodes` through to a file on disk. +""" + +import argparse +from pathlib import Path + import pytest -from conftest import make_save +from conftest import CONFIG_2D, make_save + +from lib.data.compile import compile_action_nodes, compile_plot_node +from lib.parsing.parse import parse_args +from lib.parsing.parse_save import SaveSpec, parse_save + +# --- 1. The --save argument grammar ------------------------------------------------- + + +@pytest.mark.parametrize( + "args, expected", + [ + # bare -s: every component defaulted + ([], SaveSpec()), + ([""], SaveSpec()), + # fragment: each component alone + (["out/"], SaveSpec(dir=Path("out"))), + (["fig"], SaveSpec(name="fig")), + ([".gif"], SaveSpec(format="gif")), + # fragment: combinations + (["out/fig.gif"], SaveSpec(dir=Path("out"), name="fig", format="gif")), + (["out/.gif"], SaveSpec(dir=Path("out"), format="gif")), + (["/out/fig.png"], SaveSpec(dir=Path("/out"), name="fig", format="png")), + (["../out/"], SaveSpec(dir=Path("../out"))), + # rule 2 fallthrough: all-dot segments are dirs, not stems + (["."], SaveSpec(dir=Path("."))), + ([".."], SaveSpec(dir=Path(".."))), + (["./.."], SaveSpec(dir=Path("./.."))), + # last dot wins; trailing dot is absorbed into the stem + (["fig.tar.gz"], SaveSpec(name="fig.tar", format="gz")), + (["fig."], SaveSpec(name="fig.")), + # keys + (["dir=out", "name=fig", "format=gif"], SaveSpec(dir=Path("out"), name="fig", format="gif")), + (["out/", "name=fig.i"], SaveSpec(dir=Path("out"), name="fig.i")), + (["name=abc=efg"], SaveSpec(name="abc=efg")), + # key detection: a non-identifier before "=" means it is a fragment + (["out/fig=x"], SaveSpec(dir=Path("out"), name="fig=x")), + (["a-b=c"], SaveSpec(name="a-b=c")), + # a leading "~" expands; bash does not do it after "=" or inside quotes + (["dir=~/figs"], SaveSpec(dir=Path.home() / "figs")), + (["dir=~"], SaveSpec(dir=Path.home())), + (["~/figs/"], SaveSpec(dir=Path.home() / "figs")), + # "~" naming no real user stays literal rather than raising + (["dir=~nosuchuser42/x"], SaveSpec(dir=Path("~nosuchuser42/x"))), + ], +) +def test_parse_save(args, expected): + assert parse_save(args) == expected + + +@pytest.mark.parametrize( + "args, message_fragment", + [ + # identifier before "=" that is not a known key + (["stem=foo"], "save key"), + (["ext=gif"], "save key"), + (["a.b=c"], "save key"), + # component set twice + (["out/", "dir=other"], "at most once"), + (["out/fig.gif", "format=png"], "at most once"), + (["name=a", "fig"], "at most once"), + # more than one fragment + (["fig", "other"], "at most one path fragment"), + # per-key validation + (["name=a/b"], "name"), + (["name=.."], "name"), + (["name="], "name"), + (["format=a.b"], "format"), + (["format=a/b"], "format"), + (["format="], "format"), + (["dir="], "dir"), + ], +) +def test_parse_save_errors(args, message_fragment): + with pytest.raises(argparse.ArgumentTypeError, match=message_fragment): + parse_save(args) + + +# --- 2. The derived filename stem --------------------------------------------------- + + +@pytest.mark.parametrize( + "args_list, expected_stem", + [ + (["pfd", "hx_fc"], "pfd__hx_fc-v_y,z"), + (["pfd", "hx_fc", "--nan0"], "pfd__hx_fc-nan0-v_y,z"), + (["pfd", "hx_fc", "--scale", "log"], "pfd__hx_fc-scale_log-v_y,z"), + (["pfd", "hx_fc", "-v", "y", "z", "time="], "pfd__hx_fc-v_y,z;time="), + ], +) +def test_save_file_stem(args_list, expected_stem): + actual_stem = compile_plot_node(parse_args(args_list), CONFIG_2D).get_save_file_stem() + assert actual_stem == expected_stem + + +# --- 3. End to end, through the pipeline to a file ---------------------------------- + +_BASE = ["pfd", "hx_fc", "-i", "t=-1", "-v", "y", "time=", "-q"] def test_save_static_png(tmp_path): @@ -19,3 +126,45 @@ def test_save_animated_gif(tmp_path): with Image.open(path) as img: assert img.n_frames == 11 + + +def test_no_save_flag_produces_no_action_nodes(): + assert compile_action_nodes(parse_args(_BASE), CONFIG_2D) == [] + + +def test_save_uses_derived_stem_by_default(tmp_path): + [node] = compile_action_nodes(parse_args([*_BASE, "-s", f"{tmp_path}/"]), CONFIG_2D) + node.pull() + assert (tmp_path / f"{node.get_save_file_stem()}.png").exists() + + +def test_save_name_and_format_override_the_output_path(tmp_path): + [node] = compile_action_nodes(parse_args([*_BASE, "-s", f"{tmp_path}/", "name=myfig", "format=png"]), CONFIG_2D) + node.pull() + assert (tmp_path / "myfig.png").exists() + + +def test_save_fragment_sets_dir_name_and_ext(tmp_path): + [node] = compile_action_nodes(parse_args([*_BASE, "-s", f"{tmp_path}/myfig.png"]), CONFIG_2D) + node.pull() + assert (tmp_path / "myfig.png").exists() + + +def test_save_name_is_not_sanitized(tmp_path): + """A ':' in a derived stem is rewritten by sanitize_stem; an explicit name is not.""" + [node] = compile_action_nodes(parse_args([*_BASE, "-s", f"{tmp_path}/", "name=a:b"]), CONFIG_2D) + node.pull() + assert (tmp_path / "a:b.png").exists() + + +def test_save_format_flag_is_gone(): + with pytest.raises(SystemExit): + parse_args([*_BASE, "--save-format", "gif"]) + + +def test_save_incompatible_format_falls_back_to_default(tmp_path): + """A static plot only allows 'png'; requesting 'jpg' should warn and fall back.""" + [node] = compile_action_nodes(parse_args([*_BASE, "-s", f"{tmp_path}/", "name=myfig", "format=jpg"]), CONFIG_2D) + with pytest.warns(UserWarning, match="jpg is incompatible with the data; reverting to default"): + node.pull() + assert (tmp_path / "myfig.png").exists() diff --git a/tests/test_save_filename.py b/tests/test_save_filename.py deleted file mode 100644 index a35c03e..0000000 --- a/tests/test_save_filename.py +++ /dev/null @@ -1,19 +0,0 @@ -import pytest -from conftest import CONFIG_2D - -from lib.data.compile import compile_plot_node -from lib.parsing.parse import parse_args - - -@pytest.mark.parametrize( - "args_list, expected_stem", - [ - (["pfd", "hx_fc"], "pfd__hx_fc-v_y,z"), - (["pfd", "hx_fc", "--nan0"], "pfd__hx_fc-nan0-v_y,z"), - (["pfd", "hx_fc", "--scale", "log"], "pfd__hx_fc-scale_log-v_y,z"), - (["pfd", "hx_fc", "-v", "y", "z", "time="], "pfd__hx_fc-v_y,z;time="), - ], -) -def test_save_file_stem(args_list, expected_stem): - actual_stem = compile_plot_node(parse_args(args_list), CONFIG_2D).get_save_file_stem() - assert actual_stem == expected_stem