Correctly write Oc package

Author: deltamarnix

flopy4/mf6/filters.py

@@ -15,13 +15,13 @@ def fieldkind(field: Attribute) -> str:
     - 'dim' for integer fields describing a dimension's size
     - 'attr' for all other fields
     """
-    if meta := field.metadata is None:
+    if (meta := field.metadata) is None:
         raise TypeError(f"Field {field.name} has no metadata")
-    if xatmeta := meta.get("xattree", None) is None:
+    if (xatmeta := meta.get("xattree", None)) is None:
         raise TypeError(f"Field {field.name} has no xattree metadata")
-    if kind := xatmeta.get("kind", None) is None:
+    if "kind" not in xatmeta:
         raise TypeError(f"Field {field.name} has no kind")
-    return kind
+    return xatmeta.get("kind", "attr")
 
 
 @pass_context
@@ -32,7 +32,9 @@ def fieldvalue(ctx, field: Attribute):
 
 def arraydelayed(value: xr.DataArray):
     """Yield chunks (lines) from a Dask array."""
-    for chunk in value.data.to_delayed():
+    # TODO: Determine a good chunk size,
+    # because if the underlying array is only numpy, it will stay one block.
+    for chunk in value.chunk():
         yield chunk.compute()
 
 

flopy4/mf6/templates/blocks.jinja

@@ -1,3 +1,5 @@
+{% import 'macros.jinja' as macros with context %}
+
 {% for block_name, block_ in blocks.items() %}
 BEGIN {{ block_name }}
 {% for field in block_.values() %}

flopy4/mf6/templates/macros.jinja

@@ -2,31 +2,31 @@
 {% set kind = field|fieldkind %}
 {% if kind == 'attr' %}
 {% if field.type|isdict %}
-{{ macros.record(field) }}
+{{ record(field) }}
 {% else %}
-{{ macros.scalar(field) }}
+{{ scalar(field) }}
 {% endif %}
 {% elif kind in ['array', 'coord'] %}
-{{ macros.array(field) }}
+{{ array(field) }}
 {% elif kind == 'dim' %}
-{{ macros.scalar(field) }}
+{{ scalar(field) }}
 {% elif kind == 'child' %}
 {# TODO #}
 {% endif %}
 {% endmacro %}
 
 {% macro scalar(field) %}
-{{ field.name }} {% endif %}{{ field.value }}
+{{ field.name }} {{ field.value }}
 {% endmacro %}
 
 {% macro keystring(field) %} {# union #}
 {% for item in field.value.items() %}
-{{ macros.field(item) }}
+{{ field(item) }}
 {% endfor %}
 {% endmacro %}
 
 {% macro record(field) %}
-{% for item in field.values() %}{% if item.tagged %}{{ item.name }} {% endif %}{{ macros.field(field) }}{%
+{% for item in field.values() %}{% if item.tagged %}{{ item.name }} {% endif %}{{ field(field) }}{%
 endfor
 %}
 {% endmacro %}
@@ -43,8 +43,9 @@ CONSTANT
 {% if layered %}
 {{ field.name }}
 {% for val in field.value %}{{ val }}{% endfor %}
+{% endif %}
 {% elif how == "internal" %}
-{{ field.name }} {{ macro.internal_array(field) }}
+{{ field.name }} {{ internal_array(field) }}
 {% elif how == "external" %}
 {{ field.name}} OPEN/CLOSE {{ field.value }}
 {% endif %}

flopy4/mf6/utils/cbc_reader.py

@@ -12,7 +12,7 @@
 from attrs import define
 from flopy.discretization import StructuredGrid
 
-from flopy4.structured_grid import StructuredGridWrapper
+from flopy4.discretization.structured_grid import StructuredGridWrapper
 
 from .grid_utils import get_coords
 

test/test_jinja.py

@@ -0,0 +1,49 @@
+from pathlib import Path
+
+from flopy.discretization import StructuredGrid
+from flopy.discretization.modeltime import ModelTime
+from jinja2 import Environment, PackageLoader
+
+from flopy4.mf6 import filters
+from flopy4.mf6.gwf import Chd, Gwf, Npf, Oc
+from flopy4.mf6.ims import Ims
+from flopy4.mf6.simulation import Simulation
+from flopy4.mf6.spec import blocks_dict, fields_dict
+
+
+def test_simulation_to_jinja():
+    name = "quickstart"
+    workspace = Path(__file__).parent / name
+    time = ModelTime(perlen=[1.0], nstp=[1])
+    grid = StructuredGrid(nlay=1, nrow=10, ncol=10)
+    sim = Simulation(name=name, path=workspace, tdis=time)
+    ims = Ims(parent=sim)
+    gwf_name = "mymodel"
+    gwf = Gwf(parent=sim, name=gwf_name, save_flows=True, dis=grid)
+    npf = Npf(parent=gwf, save_specific_discharge=True)
+    chd = Chd(
+        parent=gwf,
+        head={"*": {(0, 0, 0): 1.0, (0, 9, 9): 0.0}},
+    )
+    oc = Oc(
+        parent=gwf,
+        budget_file=f"{gwf.name}.bud",
+        head_file=f"{gwf.name}.hds",
+        save_head={"*": "all"},
+        save_budget={"*": "all"},
+    )
+
+    env = Environment(
+        loader=PackageLoader("flopy4.mf6"),
+        trim_blocks=True,
+        lstrip_blocks=True,
+    )
+    env.filters["fieldkind"] = filters.fieldkind
+    env.filters["fieldvalue"] = filters.fieldvalue
+    env.filters["arraydelayed"] = filters.arraydelayed
+    env.filters["array2string"] = filters.array2string
+
+    fields = fields_dict(Oc)
+    blocks = blocks_dict(Oc)
+    result = env.get_template("blocks.jinja").render(fields=fields, blocks=blocks, data=oc.data)
+    assert result != ""