closer

Author: wpbonelli

flopy4/structured_grid.py renamed to flopy4/discretization/structured_grid.py

@@ -0,0 +1,12 @@
+"""
+Core IO framework. Program interfaces can plug in bespoke
+load/write methods for particular components and formats.
+
+Most of this module is stolen/simplified from astropy, at:
+- https://github.com/astropy/astropy/tree/main/astropy/io.
+"""
+
+from flopy4.io.framework import IOMethod, Loader, Writer
+from flopy4.io.registry import DEFAULT_REGISTRY
+
+__all__ = ["IOMethod", "Loader", "Writer", "DEFAULT_REGISTRY"]

flopy4/io/__init__.py

@@ -0,0 +1,48 @@
+from typing import Literal
+
+from flopy4.io.registry import Registry
+
+IOMethodName = Literal["load", "write"]
+
+
+class IO:
+    def __init__(self, instance, cls, method_name: IOMethodName, registry: Registry = None):
+        self._registry = registry
+        self._instance = instance
+        self._cls = cls
+        self._method_name: IOMethodName = method_name
+
+    @property
+    def registry(self):
+        return self._registry
+
+    def list_formats(self, out=None):
+        formats = self._registry.get_formats(self._cls, self._method_name)
+
+        if out is None:
+            formats.pprint(max_lines=-1, max_width=-1)
+        else:
+            out.write("\n".join(formats.pformat(max_lines=-1, max_width=-1)))
+
+        return out
+
+
+class IOMethod(property):
+    def __get__(self, instance, owner_cls):
+        return self.fget(instance, owner_cls)
+
+
+class Loader(IO):
+    def __init__(self, instance, cls):
+        super().__init__(instance, cls, "load", registry=None)
+
+    def __call__(self, *args, **kwargs) -> None:
+        return self.registry.load(self._cls, *args, **kwargs)
+
+
+class Writer(IO):
+    def __init__(self, instance, cls):
+        super().__init__(instance, cls, "write", registry=None)
+
+    def __call__(self, *args, **kwargs) -> None:
+        return self.registry.write(self._cls, *args, **kwargs)

flopy4/io/framework.py

@@ -0,0 +1,45 @@
+class Registry:
+    def __init__(self):
+        self._loaders = {}
+        self._writers = {}
+
+    def get_loader(self, cls, format=None):
+        return next(
+            iter(
+                [
+                    fn
+                    for ((fmt, cls_), fn) in self._loaders.items()
+                    if fmt == format and issubclass(cls, cls_)
+                ]
+            )
+        )
+
+    def get_writer(self, cls, format=None):
+        return next(
+            iter(
+                [
+                    fn
+                    for ((fmt, cls_), fn) in self._writers.items()
+                    if fmt == format and issubclass(cls, cls_)
+                ]
+            )
+        )
+
+    def register_loader(self, cls, format, function):
+        if format in self._loaders:
+            raise ValueError(f"Reader for format {format} already registered.")
+        self._loaders[cls, format] = (cls, function)
+
+    def register_writer(self, cls, format, function):
+        if format in self._writers:
+            raise ValueError(f"Writer for format {format} already registered.")
+        self._writers[cls, format] = (cls, function)
+
+    def load(self, cls, *args, format=None, **kwargs):
+        return self.get_loader(cls, format)(*args, **kwargs)
+
+    def write(self, cls, *args, format=None, **kwargs):
+        return self.get_writer(cls, format)(*args, **kwargs)
+
+
+DEFAULT_REGISTRY = Registry()

flopy4/io/registry.py

@@ -0,0 +1,43 @@
+import sys
+
+import numpy as np
+from jinja2 import Environment, PackageLoader
+
+from flopy4.io import DEFAULT_REGISTRY
+from flopy4.mf6 import filters
+from flopy4.mf6.component import Component
+from flopy4.mf6.spec import blocks_dict, fields_dict
+
+JINJA_ENV = Environment(
+    loader=PackageLoader("flopy4.mf6"),
+    trim_blocks=True,
+    lstrip_blocks=True,
+)
+JINJA_ENV.filters["fieldkind"] = filters.fieldkind
+JINJA_ENV.filters["fieldvalue"] = filters.fieldvalue
+JINJA_ENV.filters["arraydelayed"] = filters.arraydelayed
+JINJA_ENV.filters["array2string"] = filters.array2string
+JINJA_TEMPLATE_NAME = "blocks.jinja"
+
+
+def _load_ascii(self) -> None:
+    # TODO
+    pass
+
+
+def _write_ascii(self) -> None:
+    cls = type(self)
+    fields = fields_dict(cls)
+    blocks = blocks_dict(cls)
+    template = JINJA_ENV.get_template(JINJA_TEMPLATE_NAME)
+    iterator = template.generate(fields=fields, blocks=blocks, data=unstructure(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)
+
+
+# TODO: where to do this? probably not here..on plugin discovery?
+DEFAULT_REGISTRY.register_loader(Component, "ascii", _load_ascii)
+DEFAULT_REGISTRY.register_writer(Component, "ascii", _write_ascii)

flopy4/mf6/codec.py

@@ -3,7 +3,7 @@
 
 from xattree import xattree
 
-from flopy4.mf6.io import ComponentReader, ComponentWriter, IOMethod
+from flopy4.io import IOMethod, Loader, Writer
 
 COMPONENTS = {}
 """MF6 component registry."""
@@ -14,12 +14,17 @@ class Component(ABC, MutableMapping):
     """
     Base class for MF6 components.
 
+    Notes
+    -----
+    All subclasses of `Component` must be decorated with `xattree`.
+
     We use the `children` attribute provided by `xattree`. We know
-    children are also `Component`s, but mypy does not. How to fix?
+    children are also `Component`s, but mypy does not. TODO: fix??
     """
 
     @classmethod
     def __attrs_init_subclass__(cls):
+        # add class to the component registry
         COMPONENTS[cls.__name__.lower()] = cls
 
     def __getitem__(self, key):
@@ -37,10 +42,10 @@ def __iter__(self):
     def __len__(self):
         return len(self.children)  # type: ignore
 
-    _read = IOMethod(ComponentReader)  # type: ignore
-    _write = IOMethod(ComponentWriter)  # type: ignore
+    _load = IOMethod(Loader)  # type: ignore
+    _write = IOMethod(Writer)  # type: ignore
 
-    def read(self, format=None) -> None:
+    def load(self, format=None) -> None:
         self._read(format=format)
         for child in self.children.values():  # type: ignore
             child.read(format=format)

flopy4/mf6/component.py

@@ -1,52 +1,41 @@
-import types
-from typing import Union, get_args, get_origin
-
 import numpy as np
 import xarray as xr
-import xattree
+from attrs import Attribute
 from jinja2 import pass_context
 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}")
+
+
+def fieldkind(field: Attribute) -> str:
+    """
+    Get a field's `xattree` kind. Kind is either:
+
+    - 'child' for child fields
+    - 'array' for array fields
+    - 'coord' for coordinate array fields
+    - 'dim' for integer fields describing a dimension's size
+    - 'attr' for all other fields
+    """
+    if meta := field.metadata is None:
+        raise TypeError(f"Field {field.name} has no metadata")
+    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:
+        raise TypeError(f"Field {field.name} has no kind")
+    return kind
 
 
 @pass_context
-def fieldvalue(ctx, field: Xattribute):
+def fieldvalue(ctx, field: Attribute):
+    """Get a field's value from the data tree via the template context."""
     return ctx["data"][field.name]
 
 
 def arraydelayed(value: xr.DataArray):
-    for block in value.data.to_delayed():
-        block_data = block.compute()
-        yield block_data
+    """Yield chunks (lines) from a Dask array."""
+    for chunk in value.data.to_delayed():
+        yield chunk.compute()
 
 
 def array2string(value: NDArray) -> str:
+    """Convert an array to a string."""
     return np.array2string(value, separator=" ")[1:-1]  # remove brackets