From 1501121fe2dd7554a92eb21c369e526226e28067 Mon Sep 17 00:00:00 2001 From: wpbonelli Date: Mon, 12 Oct 2026 18:33:44 -0700 Subject: [PATCH] make a list of unions' keyword arms words, and round-trip union rows through dataframes PRP's ALL, FIRST and LAST were classes (Prp.All, ...) while OC's ocsetting takes the same words as strings. A list of unions' keyword arms are now words too: PRP's period element is Literal['ALL', 'FIRST', 'LAST'] | Frequency | Steps | Fraction. Rows read, write and coerce as a union field's value does, and a row no arm takes raises rather than being dropped. to_dataframe kept a union field's arm (MAW's mawsetting) as a dict, which from_dataframe couldn't read back; it keeps the instance now. A list of unions' frame names each row's arm, or is its word, in a keyword column, so from_dataframe takes it too. Co-Authored-By: Claude Opus 5.5 --- flopy4/mf6/converter/egress/unstructure.py | 13 +-- flopy4/mf6/converter/ingress/structure.py | 10 +- flopy4/mf6/item.py | 42 +++++-- flopy4/mf6/package.py | 110 +++++++++++------- flopy4/mf6/prt/prp.py | 19 +-- flopy4/mf6/utils/codegen/make.py | 27 ++++- .../utils/codegen/templates/package.py.jinja | 4 +- test/mf6/test_mf6_codec.py | 60 +++++++--- test/mf6/test_mf6_codegen.py | 10 ++ test/test_dataframe_api.py | 35 +++++- 10 files changed, 224 insertions(+), 106 deletions(-) diff --git a/flopy4/mf6/converter/egress/unstructure.py b/flopy4/mf6/converter/egress/unstructure.py index f686f0aba..f9c7d67f9 100644 --- a/flopy4/mf6/converter/egress/unstructure.py +++ b/flopy4/mf6/converter/egress/unstructure.py @@ -133,10 +133,11 @@ def _period_dataarray( return _mark_netcdf(da, meta) -def _rows_to_tuples(row_list: list) -> list[tuple]: +def _rows_to_tuples(row_list: list) -> list: """Convert a list of Item instances to MF6 record tuples via each - Item's own to_tokens().""" - return [row.to_tokens() for row in row_list] + Item's own to_tokens(), and a list of unions' words (PRP's ALL) to + lines of their own.""" + return [[row] if isinstance(row, str) else row.to_tokens() for row in row_list] def _wrap_array(value: Any) -> xr.DataArray: @@ -247,11 +248,7 @@ def _unstructure_package(value: Package) -> dict[str, Any]: kper_int = _normalize_kper(kper) if kper_int is None: continue - rows = ( - _rows_to_tuples(row_list) - if isinstance(row_list, list) and row_list and isinstance(row_list[0], Item) - else [] - ) + rows = _rows_to_tuples(row_list) if isinstance(row_list, list) else [] spd_period.setdefault(kper_int, []).extend(rows) continue diff --git a/flopy4/mf6/converter/ingress/structure.py b/flopy4/mf6/converter/ingress/structure.py index f55802525..339c03d80 100644 --- a/flopy4/mf6/converter/ingress/structure.py +++ b/flopy4/mf6/converter/ingress/structure.py @@ -19,6 +19,7 @@ Item, dim_lookup, item_list_type, + item_list_words, parse_union_items, sized_by, ) @@ -54,15 +55,18 @@ def _parse_rows( union_arm: str = "cellid", union_forms: Callable[[Mapping[str, Any]], Any] | None = None, load_child: Callable[[Any, str], Any] | None = None, + words: tuple[str, ...] = (), ) -> list | None: - """Parse lists of tokens (rows) into a list of Items, their cellids' - widths per the context (see `LoadContext.ncelldim`).""" + """Parse lists of tokens (rows) into a list of Items (or, in a list of + unions, its ``words``), their cellids' widths per the context (see + `LoadContext.ncelldim`).""" if not rows: return None if isinstance(item_cls, tuple): return parse_union_items( rows, item_cls, + words=words, sizes=sizes, boundnames=boundnames, dims=dict(context.dims), @@ -1093,6 +1097,7 @@ def load_child(f: Any, token: str) -> Any: union_arm=union_arm, union_forms=union_forms, load_child=load_child, + words=item_list_words(f.type), ) if row_list is not None: init_key = f.alias if (f.alias and not f.alias.startswith("_")) else f.name @@ -1153,6 +1158,7 @@ def load_child(f: Any, token: str) -> Any: union_arm=union_arm, union_forms=union_forms, load_child=load_child, + words=item_list_words(period_field.type), ) if row_list is not None: spd[kper] = row_list diff --git a/flopy4/mf6/item.py b/flopy4/mf6/item.py index 8dc5a06b1..642fc448f 100644 --- a/flopy4/mf6/item.py +++ b/flopy4/mf6/item.py @@ -19,7 +19,7 @@ from collections.abc import Callable, Mapping from functools import lru_cache from pathlib import Path -from typing import Any, ClassVar, Union, cast, get_args, get_origin +from typing import Any, ClassVar, Literal, Union, cast, get_args, get_origin import attrs @@ -177,12 +177,12 @@ def to_union_arm(value: Any, self: Any, f: attrs.Attribute) -> Any: if value is None or (arms := _union_field_arms(type(self), f)) is None: return value values = list(value) if isinstance(value, (list, tuple)) else [value] - if (converted := _union_value(values, arms)) is None: + if (converted := union_value(values, arms)) is None: raise ValueError(f"{type(self).__name__}.{f.name}: no arm of {f.type} takes {value!r}") return converted -def _union_value(values: list, arms: "tuple[tuple[str, ...], tuple[type[Item], ...]]") -> Any: +def union_value(values: list, arms: "tuple[tuple[str, ...], tuple[type[Item], ...]]") -> Any: """A union field's value from the user-supplied values it spans: one of its words, an arm instance given as is, or an arm built from them by its keyword (see construct_union_item).""" @@ -258,7 +258,7 @@ class also has boundname (always declared last) -- a string there if arms is None and len(trailing) == 1 and isinstance(trailing[0], (list, tuple)): # The array given as one value, e.g. (0, 0.5, 0.5, 4, (0, 1, 4, 3)). trailing = list(trailing[0]) - tuple_vals = _union_value(trailing, arms) if arms is not None else tuple(trailing) + tuple_vals = union_value(trailing, arms) if arms is not None else tuple(trailing) if boundname_val is not None: return cast("Item", item_cls(*before, tuple_vals, boundname=boundname_val)) return cast("Item", item_cls(*before, tuple_vals)) @@ -843,23 +843,36 @@ def _unwrap_item(item) -> "type[Item] | tuple[type[Item], ...] | None": return None -def item_list_type(field_type) -> "type[Item] | tuple[type[Item], ...] | None": - """For Optional[list[C]] or Optional[dict[int, list[C]]], return C (or - the tuple of arm classes for a Union item type).""" +def _list_elem(field_type) -> Any: + """For Optional[list[C]] or Optional[dict[int, list[C]]], C.""" args = get_args(field_type) inner = next((a for a in args if a is not type(None)), None) if inner is None: return None origin = get_origin(inner) if origin is list: - return _unwrap_item(get_args(inner)[0]) + return get_args(inner)[0] if origin is dict: _, val = get_args(inner) if get_origin(val) is list: - return _unwrap_item(get_args(val)[0]) + return get_args(val)[0] return None +def item_list_type(field_type) -> "type[Item] | tuple[type[Item], ...] | None": + """For Optional[list[C]] or Optional[dict[int, list[C]]], return C (or + the tuple of arm classes for a Union item type).""" + elem = _list_elem(field_type) + return None if elem is None else _unwrap_item(elem) + + +def item_list_words(field_type) -> tuple[str, ...]: + """The words a list of unions' rows may be, besides its arm classes + (PRP's ALL, FIRST and LAST).""" + elem = _list_elem(field_type) + return tuple(w for a in get_args(elem) if get_origin(a) is Literal for w in get_args(a)) + + def dispatch_union_item(item: list, arm_classes: "tuple[type[Item], ...]") -> "type[Item] | None": """Find which arm class a raw token item belongs to, by its leading _keyword token.""" @@ -897,20 +910,25 @@ def parse_union_items( items: list, arm_classes: "tuple[type[Item], ...]", *, + words: tuple[str, ...] = (), sizes: Mapping[str, int] | None = None, boundnames: bool = False, dims: "dict | None" = None, load_child: Callable[[attrs.Attribute, str], Any] | None = None, ) -> list | None: """Parse raw token items into Item instances, dispatching each by - keyword (see dispatch_union_item); unmatched items are skipped. - ``load_child`` loads a child column's component (see from_tokens).""" + keyword (see dispatch_union_item), or into ``words`` (PRP's ``ALL``); + unmatched items are skipped. ``load_child`` loads a child column's + component (see from_tokens).""" if not items: return None - result = [] + result: list = [] for item in items: if not item: continue + if len(item) == 1 and (word := str(item[0]).upper()) in words: + result.append(word) + continue arm_cls = dispatch_union_item(item, arm_classes) if arm_cls is None: continue diff --git a/flopy4/mf6/package.py b/flopy4/mf6/package.py index 05b3885c2..c4d14b8bb 100644 --- a/flopy4/mf6/package.py +++ b/flopy4/mf6/package.py @@ -18,10 +18,11 @@ from flopy4.mf6.item import ( Item, construct_item, - construct_union_item, dim_lookup, item_list_type, + item_list_words, package_sized_fields, + union_value, ) from flopy4.mf6.load_context import LoadContext from flopy4.mf6.period_arrays import ( @@ -153,9 +154,11 @@ def _init_item_lists(self, fields) -> None: if raw is None: continue + words = item_list_words(f.type) if f.metadata.get("fill_forward"): coerced = { - kper: self._coerce_item_list(rows, item_cls) for kper, rows in raw.items() + kper: self._coerce_item_list(rows, item_cls, words) + for kper, rows in raw.items() } object.__setattr__(self, f.name, coerced) rows = [r for kper_rows in coerced.values() for r in kper_rows] @@ -163,7 +166,7 @@ def _init_item_lists(self, fields) -> None: nrows = max((len(r) for r in coerced.values()), default=0) self._set_dim_from_rows(dim, nrows, bound) else: - rows = self._coerce_item_list(raw, item_cls) + rows = self._coerce_item_list(raw, item_cls, words) object.__setattr__(self, f.name, rows) if dim: self._set_dim_from_rows(dim, len(rows), bound) @@ -353,7 +356,9 @@ def lookup(name: str) -> Any: raise ValueError(f"{shape}={declared} but {f.name} has {n} values") @staticmethod - def _coerce_item_list(data, item_cls: "type[Item] | tuple[type[Item], ...]") -> list: + def _coerce_item_list( + data, item_cls: "type[Item] | tuple[type[Item], ...]", words: tuple[str, ...] = () + ) -> list: """Convert user-supplied list/dict data to a list of Item instances. For a plain (non-union) item_cls, accepts: @@ -364,30 +369,24 @@ def _coerce_item_list(data, item_cls: "type[Item] | tuple[type[Item], ...]") -> - dict of lists → column-oriented {col_name: [values]} For a union item_cls (a tuple of arm classes, e.g. OC's (OcSave, - OcPrint)): existing arm instances pass through; + OcPrint)): existing arm instances pass through, as do its + ``words`` (PRP's ALL), in any case; tuples/lists are dispatched to the right arm by their keyword token and built positionally (see flopy4.mf6.item.construct_union_item -- NOT from_tokens, since these values are already Python-side, not raw 1-based/string file tokens); dicts are dispatched by a - "keyword" key. Ambiguous columnar dict-of-lists input isn't - supported (no single arm to build columns from). + "keyword" key. Raises on a row no arm takes. Ambiguous columnar + dict-of-lists input isn't supported (no single arm to build + columns from). """ if isinstance(item_cls, tuple): items = [] for row in data: - if isinstance(row, item_cls): - items.append(row) - elif isinstance(row, dict): - kw = str(row.get("keyword", "")).upper() - arm = next( - (c for c in item_cls if c.__dict__.get("_keyword", "").upper() == kw), None - ) - if arm is not None: - items.append(arm(**{k: v for k, v in row.items() if k != "keyword"})) - else: - item = construct_union_item(row, item_cls) - if item is not None: - items.append(item) + item = _union_row(row, item_cls, words) + if item is None: + names = [*words, *(c.__name__ for c in item_cls)] + raise ValueError(f"no arm of {names} takes row {row!r}") + items.append(item) return items if isinstance(data, dict): n = len(next(iter(data.values()))) if data else 0 @@ -639,9 +638,13 @@ def to_dataframe(self) -> pd.DataFrame: _spd = self.__dict__.get("_stress_period_data") if not _spd: return pd.DataFrame() + union = isinstance(item_list_type(self._period_item_field().type), tuple) frames = [] for kper in sorted(_spd): - df = pd.DataFrame([_row_dict(row) for row in _spd[kper]]) + # object columns in a list of unions, whose arms leave each + # other's empty: else an int column with gaps would turn float + dtype = object if union else None + df = pd.DataFrame([_row_dict(row, keyword=union) for row in _spd[kper]], dtype=dtype) df.insert(0, "kper", kper) frames.append(df) return pd.concat(frames, ignore_index=True) if frames else pd.DataFrame() @@ -650,37 +653,35 @@ def from_dataframe(self, df: pd.DataFrame) -> None: """Set stress_period_data from a tidy DataFrame. The DataFrame must have a ``kper`` column and data columns matching - the period Item class's own fields (as produced by ``to_dataframe()``). + the period Item class's own fields (as produced by ``to_dataframe()``), + and, for a list of unions, a ``keyword`` column naming each row's arm. """ if df.empty: self.__dict__["_stress_period_data"] = {} return if "kper" not in df.columns: raise ValueError("DataFrame must have a 'kper' column") - item_cls = self._period_item_cls() - if isinstance(item_cls, tuple): - raise ValueError( - f"{type(self).__name__}.from_dataframe() doesn't support a union period " - "field (multiple possible row shapes) -- construct arm instances directly." - ) + f = self._period_item_field() + item_cls = item_list_type(f.type) + assert item_cls is not None spd: dict[int, list] = {} for kper, group in df.groupby("kper"): group = group.drop(columns=["kper"]) - # A missing optional column (e.g. boundname) round-trips through - # pandas as NaN, not absent. Drop it so the Item class's own - # default applies rather than storing a float in a str field. - spd[int(kper)] = [ - item_cls(**{k: v for k, v in row.items() if not (is_scalar(v) and pd.isna(v))}) + # A missing optional column (e.g. boundname), or another arm's + # in a list of unions, round-trips through pandas as NaN, not + # absent. Drop it so the Item class's own default applies rather + # than storing a float in a str field. + records = [ + {k: v for k, v in row.items() if not (is_scalar(v) and pd.isna(v))} for row in group.to_dict("records") ] + spd[int(kper)] = self._coerce_item_list(records, item_cls, item_list_words(f.type)) self.__dict__["_stress_period_data"] = spd - def _period_item_cls(self) -> "type[Item] | tuple[type[Item], ...]": + def _period_item_field(self) -> attrs.Attribute: for f in attrs.fields(type(self)): # type: ignore[arg-type] - if f.metadata.get("fill_forward"): - item_cls = item_list_type(f.type) - if item_cls is not None: - return item_cls + if f.metadata.get("fill_forward") and item_list_type(f.type) is not None: + return f raise ValueError(f"{type(self).__name__} has no period Item-list field") @property @@ -802,15 +803,36 @@ def to_xarray(self) -> "xr.Dataset": # type: ignore[override] return super().to_xarray() # type: ignore[return-value] -def _row_dict(row: Any) -> dict[str, Any]: +def _union_row(row: Any, arms: "tuple[type[Item], ...]", words: tuple[str, ...]) -> Any: + """A list of unions' row from user input (see Package._coerce_item_list), + as a union field's value is (see union_value), or None if no arm + takes it.""" + if isinstance(row, dict): + kw = str(row.get("keyword", "")).upper() + if kw in words and len(row) == 1: + return kw + arm = next((c for c in arms if c.keyword().upper() == kw), None) + return arm(**{k: v for k, v in row.items() if k != "keyword"}) if arm else None + values = list(row) if isinstance(row, (list, tuple)) else [row] + return union_value(values, (words, arms)) + + +def _row_dict(row: Any, keyword: bool = False) -> dict[str, Any]: """A table row as a DataFrame record: `attrs.asdict`, except that - time-series references stay as they are (asdict would make each a - dict), so the frame round-trips through `from_dataframe`.""" - record = attrs.asdict(row) + time-series references and union arms stay as they are (asdict would + make each a dict), so the frame round-trips through `from_dataframe`. + With `keyword`, a list of unions' row names its arm, or is its word + (PRP's ALL), in a "keyword" column.""" + if isinstance(row, str): + return {"keyword": row} + record = {"keyword": type(row).keyword().upper()} if keyword else {} + record |= attrs.asdict(row) for f in attrs.fields(type(row)): + v = getattr(row, f.name) if f.metadata.get("time_series"): - v = getattr(row, f.name) record[f.name] = list(v) if isinstance(v, tuple) else v + elif isinstance(v, Item): + record[f.name] = v return record @@ -819,6 +841,8 @@ def _ts_names(row: Any) -> list[str]: time-series columns, including in its nested union arms (not its child files, which check their own).""" names: list[str] = [] + if isinstance(row, str): + return names for f in attrs.fields(type(row)): value = getattr(row, f.name) if f.metadata.get("time_series"): diff --git a/flopy4/mf6/prt/prp.py b/flopy4/mf6/prt/prp.py index 526245425..f2b6aabda 100644 --- a/flopy4/mf6/prt/prp.py +++ b/flopy4/mf6/prt/prp.py @@ -1,6 +1,6 @@ # autogenerated file, do not modify from pathlib import Path -from typing import ClassVar, Optional +from typing import ClassVar, Literal, Optional import attrs @@ -28,18 +28,6 @@ class Packagedata(Item): class Releasetimes(Item): time: float = field() - @attrs.define - class All(Item): - _keyword: ClassVar[str] = "all" - - @attrs.define - class First(Item): - _keyword: ClassVar[str] = "first" - - @attrs.define - class Last(Item): - _keyword: ClassVar[str] = "last" - @attrs.define class Frequency(Item): _keyword: ClassVar[str] = "frequency" @@ -55,7 +43,7 @@ class Fraction(Item): _keyword: ClassVar[str] = "fraction" fraction: tuple[float, ...] = field(default=(), array=True) - _StressPeriodDataItem = All | First | Last | Frequency | Steps | Fraction + _StressPeriodDataItem = Literal["ALL", "FIRST", "LAST"] | Frequency | Steps | Fraction boundnames: bool = field( default=False, @@ -187,9 +175,6 @@ class Fraction(Item): PrpPackagedata = Prp.Packagedata PrpReleasetimes = Prp.Releasetimes -PrpAll = Prp.All -PrpFirst = Prp.First -PrpLast = Prp.Last PrpFrequency = Prp.Frequency PrpSteps = Prp.Steps PrpFraction = Prp.Fraction diff --git a/flopy4/mf6/utils/codegen/make.py b/flopy4/mf6/utils/codegen/make.py index 1b0d24c60..647e321f5 100644 --- a/flopy4/mf6/utils/codegen/make.py +++ b/flopy4/mf6/utils/codegen/make.py @@ -119,10 +119,19 @@ class ItemClassSpec: @dataclass class ItemUnionSpec: """A union list's element type alias (e.g. ``_StressPeriodDataItem = - Stage | Rate | ...``), keeping the field annotation short.""" + Stage | Rate | ...``), keeping the field annotation short. Its keyword + arms are words, like a union field's (PRP's ``Literal['ALL', 'FIRST', + 'LAST'] | Frequency | ...``).""" alias: str members: list[str] + words: list[str] = dc_field(default_factory=list) + + @property + def arms(self) -> list[str]: + """The alias's ``|``-joined arms: its words, then its classes.""" + literal = f"Literal[{', '.join(repr(w.upper()) for w in self.words)}]" + return ([literal] if self.words else []) + self.members @dataclass @@ -492,7 +501,7 @@ def _build_list_item_specs( - a union element (OC's output, LAK's period settings): one class per arm, dispatched by keyword, and a ``_Item`` alias for the - union + union, its keyword arms words (PRP's ALL, FIRST and LAST) - a tagged (keyword-led) record element: one class whose leading keyword is its ``_keyword``, built like a union arm - an untagged record element (packagedata, CHD's stress_period_data, @@ -509,7 +518,8 @@ def _build_list_item_specs( ) alias = f"_{class_name}Item" members = [s.class_name for s in specs if s.top_level] - return specs, alias, ItemUnionSpec(alias=alias, members=members) + words = [a.name for a in item.arms.values() if isinstance(a, KeywordField)] + return specs, alias, ItemUnionSpec(alias=alias, members=members, words=words) if list_field.tagged: specs = _build_arm_spec( list_field.name, @@ -589,10 +599,13 @@ def _build_arm_specs_from_union( children: "Mapping[str, list[str]] | None" = None, ) -> list[ItemClassSpec]: """Build one ItemClassSpec per arm of a list's union element (see - _build_arm_spec): OC's output, PRP's period settings, or LAK's, each - arm of which leads with its own lake or outlet number.""" + _build_arm_spec), but for its keyword arms, which are words: OC's + output, PRP's period settings, or LAK's, each arm of which leads with + its own lake or outlet number.""" specs: list[ItemClassSpec] = [] for arm_name, arm in union.arms.items(): + if isinstance(arm, KeywordField): + continue specs.extend( _build_arm_spec( arm_name, @@ -905,6 +918,7 @@ def _generated_imports( has_derived_dims: bool = False, has_optional_child: bool = False, has_union_child: bool = False, + has_literal: bool = False, extra_imports: list[str] | None = None, ) -> dict[str, list[str]]: """Compute import lines for generated packages.""" @@ -946,7 +960,7 @@ def _generated_imports( typing_parts.append("Optional") if has_union: typing_parts.append("Union") - if any(col.uses_literal for col in _all_schema_cols): + if has_literal or any(col.uses_literal for col in _all_schema_cols): typing_parts.append("Literal") if typing_parts: stdlib.append(f"from typing import {', '.join(sorted(typing_parts))}") @@ -1848,6 +1862,7 @@ def _block_index(fs: FieldSpec) -> int: fs.spec_call.startswith("child(") and "Union[" in fs.type_annotation for fs in field_specs ), + has_literal=any(u.words for u in item_unions), extra_imports=[ *extra_imports, *(["from flopy4.mf6.block import Block"] if block_classes else []), diff --git a/flopy4/mf6/utils/codegen/templates/package.py.jinja b/flopy4/mf6/utils/codegen/templates/package.py.jinja index 98c5ee9a1..0713cd007 100644 --- a/flopy4/mf6/utils/codegen/templates/package.py.jinja +++ b/flopy4/mf6/utils/codegen/templates/package.py.jinja @@ -87,8 +87,8 @@ class {{ spec.class_name }}({% for m in spec.mixins %}{{ m }}, {% endfor %}{{ sp {% endfor %} {% for union in spec.item_unions %} {{ union.alias }} = ( -{% for member in union.members %} - {{ member }}{{ " |" if not loop.last }} +{% for arm in union.arms %} + {{ arm }}{{ " |" if not loop.last }} {% endfor %} ) diff --git a/test/mf6/test_mf6_codec.py b/test/mf6/test_mf6_codec.py index f6e80bb24..8d9f92ece 100644 --- a/test/mf6/test_mf6_codec.py +++ b/test/mf6/test_mf6_codec.py @@ -715,6 +715,30 @@ def test_union_field_takes_arm_values(): Maw.StressPeriodData(ifno=0, mawsetting=("bogus", 1)) +def test_union_field_rows_from_dicts(): + """A row with a union field (MAW's, LKT's) builds from a dict of its + columns, the union field given as its arm's values.""" + from flopy4.mf6.gwf import Maw + from flopy4.mf6.gwt import Lkt + + maw = Maw( + stress_period_data={ + 0: [ + {"ifno": 0, "mawsetting": ("rate", -1.0)}, + {"ifno": 1, "mawsetting": Maw.Status(status="INACTIVE")}, + ] + } + ) + assert maw.stress_period_data[0] == [ + Maw.StressPeriodData(ifno=0, mawsetting=Maw.Rate(rate=-1.0)), + Maw.StressPeriodData(ifno=1, mawsetting=Maw.Status(status="INACTIVE")), + ] + lkt = Lkt(stress_period_data={0: [{"ifno": 0, "laksetting": ("concentration", 1.5)}]}) + assert lkt.stress_period_data[0][0].laksetting == Lkt.Concentration(concentration=1.5) + with pytest.raises(ValueError, match="no arm"): + Maw(stress_period_data={0: [{"ifno": 0, "mawsetting": ("bogus", 1)}]}) + + def test_maw_round_trip(): from flopy4.mf6.converter.egress.unstructure import unstructure_component from flopy4.mf6.converter.ingress.structure import structure_component @@ -1692,34 +1716,40 @@ def test_dumps_prt_prp_release_setting(): def test_prt_prp_period_roundtrip(): - """PRT-PRP release settings survive a dump→load→structure_component cycle.""" + """PRT-PRP release settings, words (ALL, FIRST, LAST) and arms alike, + survive a dump→load→structure_component cycle.""" from flopy4.mf6.converter.egress.unstructure import unstructure_component from flopy4.mf6.converter.ingress.structure import structure_component from flopy4.mf6.prt.prp import Prp prp = Prp( - dims={"nper": 2}, + dims={"nper": 3}, stress_period_data={ - 0: [("FREQUENCY", 2)], + 0: [("FREQUENCY", 2), "all"], 1: [("FIRST",)], + 2: [{"keyword": "last"}, Prp.Steps(steps=(1, 3))], }, ) + assert prp.stress_period_data == { + 0: [Prp.Frequency(frequency=2), "ALL"], + 1: ["FIRST"], + 2: ["LAST", Prp.Steps(steps=(1, 3))], + } text = dumps(unstructure_component(prp)) - raw = loads(text) - prp2 = structure_component(raw, Prp) + assert "BEGIN PERIOD 2\n FIRST\nEND PERIOD 2" in text + prp2 = structure_component(loads(text), Prp) + assert prp2.stress_period_data == prp.stress_period_data - spd = prp2.stress_period_data - assert spd is not None - assert 0 in spd and 1 in spd - p0 = spd[0] - assert len(p0) == 1 - assert isinstance(p0[0], Prp.Frequency) - assert p0[0].frequency == 2 +def test_union_list_rejects_rows_no_arm_takes(): + """A list of unions' row that's none of its words or arms raises, + rather than being dropped.""" + from flopy4.mf6.prt.prp import Prp - p1 = spd[1] - assert len(p1) == 1 - assert isinstance(p1[0], Prp.First) + with pytest.raises(ValueError, match="no arm"): + Prp(dims={"nper": 1}, stress_period_data={0: ["bogus"]}) + with pytest.raises(ValueError, match="no arm"): + Prp(dims={"nper": 1}, stress_period_data={0: [("BOGUS", 1)]}) def test_oc_ocsetting_typed_dispatch(): diff --git a/test/mf6/test_mf6_codegen.py b/test/mf6/test_mf6_codegen.py index e1c26ec67..fa44887a8 100644 --- a/test/mf6/test_mf6_codegen.py +++ b/test/mf6/test_mf6_codegen.py @@ -1074,6 +1074,16 @@ def test_bounded_array_arm_is_variadic(tmp_path, all_dfns, name): assert "steps: tuple[int, ...] = field(default=(), array=True)" in spec.outpath.read_text() +def test_union_list_keyword_arms_are_words(tmp_path, all_dfns): + """A list of unions' keyword arms (PRP's ALL, FIRST and LAST) are words + in its alias, as a union field's are, not classes.""" + skip = {n for n in all_dfns if n != "prt-prp"} + (spec,) = make_modules(dfns=all_dfns, outdir=tmp_path, skip=skip, makedirs=True) + (union,) = spec.item_unions + assert union.arms == ["Literal['ALL', 'FIRST', 'LAST']", "Frequency", "Steps", "Fraction"] + assert "All" not in {ic.class_name for ic in spec.item_classes} + + @pytest.mark.parametrize( "field,fixed", [ diff --git a/test/test_dataframe_api.py b/test/test_dataframe_api.py index f3a804707..4ca958ac4 100644 --- a/test/test_dataframe_api.py +++ b/test/test_dataframe_api.py @@ -9,7 +9,8 @@ import pandas as pd import pytest -from flopy4.mf6.gwf import Chd, Drn, Wel +from flopy4.mf6.gwf import Chd, Drn, Maw, Wel +from flopy4.mf6.prt import Prp def test_chd_to_dataframe(): @@ -225,3 +226,35 @@ def test_from_dataframe_drops_missing_optional_columns(): rows = chd.stress_period_data[0] assert rows[0].boundname == "a" assert rows[1].boundname is None + + +def test_union_field_rows_round_trip(): + """A row's union field (MAW's mawsetting) is a column of arm instances, + not of dicts, so the frame round-trips.""" + maw = Maw( + stress_period_data={ + 0: [(0, "RATE", -1.0), (1, "STATUS", "INACTIVE")], + 1: [(0, "RATE", -2.0)], + } + ) + df = maw.to_dataframe() + assert list(df.columns) == ["kper", "ifno", "mawsetting"] + assert df["mawsetting"].iloc[0] == Maw.Rate(rate=-1.0) + maw2 = Maw() + maw2.from_dataframe(df) + assert maw2.stress_period_data == maw.stress_period_data + + +def test_union_list_round_trip(): + """A list of unions' rows (PRP's) name their arm, or are their word, in + a keyword column, and keep their ints though other arms leave gaps.""" + prp = Prp( + dims={"nper": 2}, + stress_period_data={0: [("FREQUENCY", 2), "ALL"], 1: [("STEPS", 1, 3)]}, + ) + df = prp.to_dataframe() + assert list(df["keyword"]) == ["FREQUENCY", "ALL", "STEPS"] + assert df["frequency"].iloc[0] == 2 and isinstance(df["frequency"].iloc[0], int) + prp2 = Prp(dims={"nper": 2}) + prp2.from_dataframe(df) + assert prp2.stress_period_data == prp.stress_period_data