wip

Author: wpbonelli

flopy4/mf6/gwf/__init__.py

@@ -19,6 +19,16 @@
 __all__ = ["Gwf", "Chd", "Dis", "Ic", "Npf", "Oc"]
 
 
+def convert_grid(value):
+    if isinstance(value, Grid):
+        return Dis.from_grid(value)
+    if isinstance(value, Dis):
+        return value
+    raise TypeError(
+        f"Expected Grid or Dis, got {type(value)}"
+    )
+
+
 @xattree
 class Gwf(Model):
     @define
@@ -40,7 +50,7 @@ def budget(self):
                 merge_to_dataset=True,
             )
 
-    dis: Dis = field(converter=lambda grid: Dis.from_grid(grid))
+    dis: Dis = field(converter=convert_grid)
     ic: Ic = field()
     oc: Oc = field()
     npf: Npf = field()

flopy4/mf6/simulation.py

@@ -12,10 +12,19 @@
 from flopy.discretization.modeltime import ModelTime
 
 
+def convert_time(value):
+    if isinstance(value, ModelTime):
+        return Tdis.from_time(value)
+    if isinstance(value, Tdis):
+        return value
+    raise TypeError(
+        f"Expected ModelTime or Tdis, got {type(value)}"
+    )
+
 @xattree
 class Simulation(Component):
     models: dict[str, Model] = field()
     exchanges: dict[str, Exchange] = field()
     solutions: dict[str, Solution] = field()
     sim_ws: Path = field(default=None)
-    tdis: Tdis = field(converter=lambda time: Tdis.from_time(time))
+    tdis: Tdis = field(converter=convert_time)

test/test_component.py

@@ -49,12 +49,12 @@ def test_init_gwf_explicit_dims():
     )
 
     assert isinstance(gwf.data, DataTree)
-    assert gwf.dis is not dis  # dimension order switched.. is this ok?
+    assert gwf.dis is dis  # dimension order switched.. is this ok?
     assert gwf.ic is ic
     assert gwf.oc is oc
     assert gwf.npf is npf
     assert gwf.chd[0] is chd
-    assert gwf.data.dis is not dis.data
+    assert gwf.data.dis is dis.data
     assert gwf.data.ic is ic.data
     assert gwf.data.oc is oc.data
     assert gwf.data.npf is npf.data
@@ -106,7 +106,7 @@ def test_init_gwf_dis_first():
     chd = Chd(parent=gwf, strict=False)
 
     assert isinstance(gwf.data, DataTree)
-    assert gwf.dis is not dis
+    assert gwf.dis is dis
     assert gwf.ic is ic
     assert gwf.oc is oc
     assert gwf.npf is npf
@@ -185,7 +185,7 @@ def test_init_sim_explicit_dims():
     assert isinstance(sim.data, DataTree)
     assert sim.data.tdis is tdis.data
     assert sim.data.gwf is gwf.data
-    assert gwf.dis is not dis  # gwf.dis has inherited dim nper
+    assert gwf.dis is dis  # gwf.dis has inherited dim nper
     assert gwf.ic is ic
     assert gwf.oc is oc
     assert gwf.npf is npf