more lark

Author: wpbonelli

flopy4/mf6/__init__.py

@@ -4,10 +4,9 @@
 from typing import Optional
 
 import numpy as np
-from attr import field
 from attrs import define
 from numpy.typing import NDArray
-from xattree import ROOT, array, dim, xattree
+from xattree import ROOT, array, dim, field, xattree
 
 __all__ = [
     "Component",
@@ -54,6 +53,12 @@ class Exchange(Package):
 
 @xattree
 class Tdis(Package):
+    @define
+    class PeriodData:
+        perlen: float
+        nstp: int
+        tsmult: float
+
     nper: int = dim(
         name="per",
         default=1,
@@ -66,6 +71,11 @@ class Tdis(Package):
     start_date_time: Optional[datetime] = field(
         default=None, metadata={"block": "options"}
     )
+    # perioddata: NDArray[np.object_] = array(
+    #     PeriodData,
+    #     dims=("per",),
+    #     metadata={"block": "perioddata"},
+    # )
     perlen: NDArray[np.floating] = array(
         default=1.0,
         dims=("per",),

flopy4/mf6/io/__init__.py

@@ -85,7 +85,7 @@
 """
 
 
-def make_parser(cls: type, **kwargs) -> Lark:
+def make_parser(cls: type, extra_params=None, **kwargs) -> Lark:
     """
     Create a parser for the MODFLOW 6 input language with the given
     parameter and block specification.
@@ -100,7 +100,8 @@ def make_parser(cls: type, **kwargs) -> Lark:
     if not has_xats(cls):
         raise ValueError(f"Class '{cls.__name__}' is not a `xattree`.")
     spec = cls.__xattree__["spec"].flat
-    params = "|".join(['"' + n + '"i' for n in spec.keys()])
+    pnames = list(spec.keys()) + (extra_params or [])
+    params = "|".join(['"' + n + '"i' for n in pnames])
     blocks = set([xat.metadata.get("block", None) for xat in spec.values()])
     blocks.discard(None)
     # temp hack, TODO detect list blocks as blocks with a single

test/test_lark.py

@@ -1,9 +1,46 @@
+from pathlib import Path
+
+import flopy
 from lark import Lark
 
-from flopy4.mf6.gwf.oc import Oc
-from flopy4.mf6.io import make_parser
+from flopy4.mf6 import Tdis
+from flopy4.mf6.io import MF6Transformer, make_parser
 
 
 def test_make_parser():
-    parser = make_parser(Oc)
+    parser = make_parser(Tdis)
     assert isinstance(parser, Lark)
+
+
+def readme_example(name: str, workspace: Path) -> flopy.mf6.MFSimulation:
+    sim = flopy.mf6.MFSimulation(
+        sim_name=name, sim_ws=workspace, exe_name="mf6"
+    )
+    tdis = flopy.mf6.ModflowTdis(sim)
+    ims = flopy.mf6.ModflowIms(sim)
+    gwf = flopy.mf6.ModflowGwf(sim, modelname=name, save_flows=True)
+    dis = flopy.mf6.ModflowGwfdis(gwf, nrow=10, ncol=10)
+    ic = flopy.mf6.ModflowGwfic(gwf)
+    npf = flopy.mf6.ModflowGwfnpf(gwf, save_specific_discharge=True)
+    chd = flopy.mf6.ModflowGwfchd(
+        gwf, stress_period_data=[[(0, 0, 0), 1.0], [(0, 9, 9), 0.0]]
+    )
+    budget_file = name + ".bud"
+    head_file = name + ".hds"
+    oc = flopy.mf6.ModflowGwfoc(
+        gwf,
+        budget_filerecord=budget_file,
+        head_filerecord=head_file,
+        saverecord=[("HEAD", "ALL"), ("BUDGET", "ALL")],
+    )
+    return sim
+
+
+def test_parse(tmp_path):
+    sim = readme_example("test", Path(tmp_path))
+    sim.write_simulation()
+    parser = make_parser(Tdis)
+    transformer = MF6Transformer()
+    with open(tmp_path / "test.tdis", "r") as f:
+        tree = parser.parse(f.read())
+        data = transformer.transform(tree)