closer

Author: wpbonelli

flopy4/mf6/component.py

@@ -2,7 +2,8 @@
 from collections.abc import MutableMapping
 
 from xattree import xattree
-from flopy4.io import Writer
+
+from flopy4.mf6.io import Writer
 
 COMPONENTS = {}
 """MF6 component registry."""

flopy4/mf6/filters.py

@@ -1,27 +1,52 @@
-from collections.abc import Iterable, Mapping
-from inspect import isclass
 import types
 from typing import Union, get_args, get_origin
+
 import numpy as np
 import xarray as xr
+import xattree
 from jinja2 import pass_context
-
-from attrs import Attribute
+from numpy.typing import NDArray
 from xattree import Xattribute
 
 
+def fieldkind(field: Xattribute) -> str:
+    """Get the kind of a field."""
+    if isinstance(field, (xattree.Array, xattree.Coord)):
+        return "array"
+    if isinstance(field, xattree.Dim):
+        return "scalar"
+    if isinstance(field, xattree.Child):
+        raise TypeError(f"Child field {field.name} unsupported in this context")
+    type_ = field.type
+    if type_ is None:
+        raise TypeError(f"Field {field.name} has no type")
+    if issubclass(type_, xattree.Scalar):
+        return "scalar"
+    args = get_args(type_)
+    origin = get_origin(type_)
+    if origin in (Union, types.UnionType):
+        if args[-1] is types.NoneType:  # Optional
+            type_ = args[0]
+            assert type_ is not None
+        else:
+            return "union"
+    if issubclass(type_, xattree.Scalar):
+        return "scalar"
+    if issubclass(type_, xattree.Attr):
+        return "record"
+    raise TypeError(f"Unsupported field type {type_} for field {field.name}")
+
+
 @pass_context
-def value(ctx, field: Xattribute):
-    """Return the kind of the field."""
-    # TODO
-    pass
+def fieldvalue(ctx, field: Xattribute):
+    return ctx["data"][field.name]
 
 
-def dask_expand(data: xr.DataArray):
-    for block in data.data.to_delayed():
+def arraydelayed(value: xr.DataArray):
+    for block in value.data.to_delayed():
         block_data = block.compute()
         yield block_data
 
 
-def nparray2string(data: np.ndarray):
-    return np.array2string(data, separator=" ")[1:-1]  # remove brackets
+def array2string(value: NDArray) -> str:
+    return np.array2string(value, separator=" ")[1:-1]  # remove brackets

flopy4/mf6/io.py

@@ -1,44 +1,44 @@
 import sys
+
 import numpy as np
 from jinja2 import Environment, PackageLoader
-from flopy4.mf6 import filters
-from flopy4.mf6.spec import blocks_dict
 
+from flopy4.mf6 import filters
+from flopy4.mf6.spec import blocks_dict, fields_dict
 
 env = Environment(
-loader=PackageLoader("flopy4.mf6"),
+    loader=PackageLoader("flopy4.mf6"),
     trim_blocks=True,
     lstrip_blocks=True,
 )
-env.filters["dask_expand"] = filters.dask_expand
-env.filters["nparray2string"] = filters.nparray2string
+env.filters["fieldkind"] = filters.fieldkind
+env.filters["fieldvalue"] = filters.fieldvalue
+env.filters["arraydelayed"] = filters.arraydelayed
+env.filters["array2string"] = filters.array2string
 
 
 class Writer:
-    def _write_ascii(self, path) -> None:
-        # TODO: factor out an ascii writer separately
-
-        block_spec = blocks_dict(type(self))
-        blocks = {}
-        for block_name, block in block_spec.items():
-            blocks[block_name] = {}
-            for field_name, field in block.items():
-                if field_name == "data":
-                    continue
-                blocks[block_name][field_name] = {
-                    "spec": field,
-                    "value": getattr(self, field_name),
-                }
+    # TODO remove type: ignore statements below.
+    # but idk how to properly type a mixin class.
+    # this one assumes the presence of attributes:
+    # - name
+    # - path
+    # - data
 
+    def _write_ascii(self) -> None:
+        cls = type(self)
+        fields = fields_dict(cls)
+        blocks = blocks_dict(cls)
         template = env.get_template("blocks.jinja")
-        iterator = template.generate(blocks=blocks)
-        with np.printoptions(
-            precision=4, linewidth=sys.maxsize, threshold=sys.maxsize
-        ):
-            with open(path, "w") as f:
+        iterator = template.generate(fields=fields, blocks=blocks, data=self.data)  # type: ignore
+        # are these printoptions always applicable?
+        with np.printoptions(precision=4, linewidth=sys.maxsize, threshold=sys.maxsize):
+            # TODO don't hardcode the filename, maybe a filename attribute?
+            with open(self.path / self.name, "w") as f:  # type: ignore
                 f.writelines(iterator)
 
     def write(self) -> None:
+        # TODO: factor out an ascii writer separately
         self._write_ascii()
         for child in self.children.values():  # type: ignore
             child.write()

flopy4/mf6/spec.py

@@ -3,11 +3,20 @@
 These include field decorators and introspection functions.
 """
 
-from attrs import NOTHING, Attribute, fields_dict
-from xattree import array as xattree_array
-from xattree import coord as xattree_coord
-from xattree import dim as xattree_dim
-from xattree import field as xattree_field
+from attrs import NOTHING, Attribute
+from xattree import (
+    array as xattree_array,
+)
+from xattree import (
+    coord as xattree_coord,
+)
+from xattree import (
+    dim as xattree_dim,
+)
+from xattree import (
+    field as xattree_field,
+)
+from xattree import fields_dict as xattree_fields_dict
 
 
 def field(
@@ -139,11 +148,24 @@ def blocks_dict(cls) -> dict[str, Block]:
     (field) name to `attrs.Attribute`.
     """
     fields = fields_dict(cls)
-    fields = {k: v for k, v in fields.items() if "block" in v.metadata}
     blocks: dict[str, Block] = {}
     for k, v in fields.items():
         block = v.metadata["block"]
         if block not in blocks:
             blocks[block] = {}
         blocks[block][k] = v
     return dict(sorted(blocks.items(), key=_block_sort_key))
+
+
+def fields(cls) -> list[Attribute]:
+    """Return an ordered list of fields for a component class."""
+    return list(fields_dict(cls).values())
+
+
+def fields_dict(cls) -> dict[str, Attribute]:
+    """
+    Return an ordered dictionary of fields for a component class,
+    whose keys are field names. Each field is an `attrs.Attribute`.
+    """
+    fields = xattree_fields_dict(cls)
+    return {k: v for k, v in fields.items() if "block" in v.metadata}

flopy4/mf6/templates/blocks.jinja

@@ -1,8 +1,8 @@
-{% for block in blocks %}
-BEGIN {{ block.name }} {% if block.index is not none %}{{ block.index }}{% endif %}
-{% for field in block %}
+{% for block_name, block_ in blocks.items() %}
+BEGIN {{ block_name }}
+{% for field in block_.values() %}
 {{ macros.field(field) }}
 {% endfor %}
-END {{ block.name }}
+END {{ block_name }}
 
 {% endfor %}

flopy4/mf6/templates/macros.jinja

@@ -1,13 +1,13 @@
 {% macro field(field) %}
-{% if field|kind == 'scalar' %}
+{% if field|fieldkind == 'scalar' %}
 {{ macros.scalar(field) }}
-{% elif field|kind == 'union' %}
+{% elif field|fieldkind == 'union' %}
 {{ macros.union(field) }}
-{% elif field|kind == 'record' %}
+{% elif field|fieldkind == 'record' %}
 {{ macros.record(field) }}
-{% elif field|kind == 'array' %}
+{% elif field|fieldkind == 'array' %}
 {{ macros.array(field) }}
-{% elif field|kind == 'list' %}
+{% elif field|fieldkind == 'list' %}
 {{ macros.list(field) }}
 {% endif %}
 {% endmacro %}
@@ -48,8 +48,8 @@ CONSTANT
 {% endmacro %}
 
 {% macro internal_array(field) %}
-{% for chunk in field.value|dask_expand %}
-{{ chunk|nparray2string }}
+{% for chunk in field|value|arraydelayed %}
+{{ chunk|array2string }}
 {% endfor %}
 {% endmacro %}