consistent typing for lists, continued

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
 

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

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

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,17 +154,19 @@ 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]
                 if dim:
                     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"):