working

Author: wpbonelli

flopy4/mf6/converters.py

@@ -1,5 +1,3 @@
-import types
-from typing import Union, get_args, get_origin
 
 import numpy as np
 from numpy.typing import NDArray
@@ -8,7 +6,7 @@
 from flopy4.mf6.constants import FILL_DNODATA
 
 
-def convert_array(value, self_, field):
+def convert_array(value, self_, field) -> NDArray:
     if not isinstance(value, dict):
         # if not a dict, assume it's a numpy array
         # and let xarray deal with it if it isn't
@@ -29,20 +27,8 @@ def convert_array(value, self_, field):
     if any(unresolved):
         raise ValueError(f"Couldn't resolve dims: {unresolved}")
 
-    # extract dtype
-    origin = get_origin(field.type)
-    args = get_args(field.type)
-    if origin in (Union, types.UnionType) and args[-1] is types.NoneType:
-        origin = args[0]  # Optional
-    if origin is NDArray:
-        dtype = args[1]
-    elif origin is np.ndarray:
-        dtype = args[0]
-    else:
-        raise ValueError(f"Expected NDArray, got {origin}")
-
     # create array
-    a = np.full(shape, fill_value=FILL_DNODATA, dtype=dtype)
+    a = np.full(shape, fill_value=FILL_DNODATA, dtype=field.dtype)
 
     def _get_nn(cellid):
         match len(cellid):
@@ -59,13 +45,13 @@ def _get_nn(cellid):
 
     # populate array. TODO: is there a way to do this
     # without hardcoding awareness of kper and cellid?
-    if "kper" in dims:
+    if "per" in dims:
         for kper, period in value.items():
             if kper == "*":
                 kper = 0
             match len(shape):
                 case 1:
-                    a[kper] = v
+                    a[kper] = value
                 case _:
                     for cellid, v in period.items():
                         nn = _get_nn(cellid)
@@ -75,6 +61,6 @@ def _get_nn(cellid):
     else:
         for cellid, v in value.items():
             nn = _get_nn(cellid)
-            a[kper, nn] = v
+            a[nn] = v
 
     return a

flopy4/mf6/gwf/__init__.py

@@ -4,12 +4,12 @@
 from attrs import define
 from xattree import field, xattree
 
-from flopy4.mf6 import Model
 from flopy4.mf6.gwf.chd import Chd
 from flopy4.mf6.gwf.dis import Dis
 from flopy4.mf6.gwf.ic import Ic
 from flopy4.mf6.gwf.npf import Npf
 from flopy4.mf6.gwf.oc import Oc
+from flopy4.mf6.model import Model
 
 __all__ = ["Gwf", "Chd", "Dis", "Ic", "Npf", "Oc"]
 

flopy4/mf6/gwf/chd.py

@@ -6,8 +6,8 @@
 from numpy.typing import NDArray
 from xattree import array, field, xattree
 
-from flopy4.mf6 import Package
 from flopy4.mf6.converters import convert_array
+from flopy4.mf6.package import Package
 
 
 # TODO get rid of multi, just infer from parent?

flopy4/mf6/gwf/dis.py

@@ -3,8 +3,8 @@
 from numpy.typing import NDArray
 from xattree import array, dim, field, xattree
 
-from flopy4.mf6 import Package
 from flopy4.mf6.converters import convert_array
+from flopy4.mf6.package import Package
 
 
 @xattree

flopy4/mf6/gwf/ic.py

@@ -3,8 +3,8 @@
 from numpy.typing import NDArray
 from xattree import array, field, xattree
 
-from flopy4.mf6 import Package
 from flopy4.mf6.converters import convert_array
+from flopy4.mf6.package import Package
 
 
 @xattree

flopy4/mf6/gwf/npf.py

@@ -6,8 +6,8 @@
 from numpy.typing import NDArray
 from xattree import array, field, xattree
 
-from flopy4.mf6 import Package
 from flopy4.mf6.converters import convert_array
+from flopy4.mf6.package import Package
 
 
 @xattree

flopy4/mf6/gwf/oc.py

@@ -6,8 +6,8 @@
 from numpy.typing import NDArray
 from xattree import array, field, xattree
 
-from flopy4.mf6 import Package
 from flopy4.mf6.converters import convert_array
+from flopy4.mf6.package import Package
 from flopy4.utils import to_path
 
 

pyproject.toml

@@ -51,7 +51,7 @@ dynamic = ["version"]
 
 [project.optional-dependencies]
 dev = ["flopy4[lint,test,build]"]
-lint = ["ruff", "pre-commit>=4.0.1,<5"]
+lint = ["ruff", "pre-commit>=4.0.1,<5", "mypy"]
 test = [
     "flopy4[lint]",
     "coverage",
@@ -139,3 +139,8 @@ test = { cmd = "pytest -v -n auto" }
 [tool.pixi.feature.lint.tasks]
 lint = { cmd = "ruff check ." }
 install = { cmd = "pre-commit install --install-hooks" }
+
+[tool.mypy]
+mypy_path = "flopy4"
+ignore_missing_imports = true
+warn_unreachable = true

test/test_component.py

@@ -3,8 +3,10 @@
 from flopy.discretization.modeltime import ModelTime
 from xarray import DataTree
 
-from flopy4.mf6 import COMPONENTS, Simulation, Tdis
+from flopy4.mf6.component import COMPONENTS
 from flopy4.mf6.gwf import Chd, Dis, Gwf, Ic, Npf, Oc
+from flopy4.mf6.simulation import Simulation
+from flopy4.mf6.tdis import Tdis
 
 
 def test_registry():