refactor and flesh out codec (#192)

misc cleanup, switch from strings/paths to file handles in load and write functions (more conventional), try json/toml for fun, rename "ascii" format -> "mf6" as it entails ascii files possibly referencing binary or netcdf files

Author: wpbonelli

docs/dev/sdd.md

@@ -135,17 +135,7 @@ Input file IO is implemented in three layers:
 
 The `flopy4.uio` module provides a pluggable IO framework adapted from [`astropy`](https://github.com/astropy/astropy/tree/main/astropy/io). A global `Registry` maintains mappings from `(component_class, format)` pairs to load and write functions. The `Component` base class implements user-facing `load` and `write` methods via descriptors which dispatch functions in the registry.
 
-Loaders and writers can be registered for any component class and format. The registry supports inheritance: a loader/writer registered for a base class is available to all subclasses.
-
-```python
-from flopy4.uio import DEFAULT_REGISTRY
-from flopy4.mf6.component import Component
-
-DEFAULT_REGISTRY.register_writer(Component, "ascii", write_ascii)
-DEFAULT_REGISTRY.register_writer(Component, "netcdf", write_netcdf)
-```
-
-The user may then select a format at call time, e.g. `component.write(format="netcdf")`.
+Loaders and writers can be registered for any component class and format. The registry supports inheritance: a loader/writer registered for a base class is available to all subclasses. The user may then select a format at call time.
 
 #### Conversion
 

flopy4/mf6/__init__.py

@@ -1,16 +1,50 @@
-from flopy4.mf6 import (  # noqa: F401
-    gwf,
-    ims,
-    simulation,
-    tdis,
-)
-from flopy4.mf6.codec import dump
+from json import dump as dump_json
+from json import load as load_json
+from pathlib import Path
+
+from tomli import load as load_toml
+from tomli_w import dump as dump_toml
+
+from flopy4.mf6.codec import dump as dump_mf6
+from flopy4.mf6.codec import load as load_mf6
 from flopy4.mf6.component import Component
-from flopy4.mf6.converter import COMPONENT_CONVERTER
+from flopy4.mf6.converter import structure, unstructure
 from flopy4.uio import DEFAULT_REGISTRY
 
-# register io methods
-# TODO: call format "mf6" or something? since it might include binary files
-DEFAULT_REGISTRY.register_writer(
-    Component, "ascii", lambda c: dump(COMPONENT_CONVERTER.unstructure(c), c.path)
-)
+
+def _load_mf6(path: Path) -> Component:
+    with open(path, "r") as fp:
+        return structure(load_mf6(fp), path)
+
+
+def _load_json(path: Path) -> Component:
+    with open(path, "r") as fp:
+        return structure(load_json(fp), path)
+
+
+def _load_toml(path: Path) -> Component:
+    with open(path, "rb") as fp:
+        return structure(load_toml(fp), path)
+
+
+def _write_mf6(component: Component) -> None:
+    with open(component.path, "w") as fp:
+        dump_mf6(unstructure(component), fp)
+
+
+def _write_json(component: Component) -> None:
+    with open(component.path, "w") as fp:
+        dump_json(unstructure(component), fp, indent=4)
+
+
+def _write_toml(component: Component) -> None:
+    with open(component.path, "wb") as fp:
+        dump_toml(unstructure(component), fp)
+
+
+DEFAULT_REGISTRY.register_loader(Component, "mf6", _load_mf6)
+DEFAULT_REGISTRY.register_loader(Component, "json", _load_json)
+DEFAULT_REGISTRY.register_loader(Component, "toml", _load_toml)
+DEFAULT_REGISTRY.register_writer(Component, "mf6", _write_mf6)
+DEFAULT_REGISTRY.register_writer(Component, "json", _write_json)
+DEFAULT_REGISTRY.register_writer(Component, "toml", _write_toml)

flopy4/mf6/codec/reader/__init__.py

@@ -1,12 +1,13 @@
-from os import PathLike
-from pathlib import Path
-from typing import Any
+from typing import IO, Any
 
 from flopy4.mf6.codec.reader.parser import make_basic_parser
 from flopy4.mf6.codec.reader.transformer import BasicTransformer
 
+BASIC_PARSER = make_basic_parser()
+BASIC_TRANSFORMER = BasicTransformer()
 
-def load(path: str | PathLike) -> Any:
+
+def load(fp: IO[str]) -> Any:
     """
     Load and parse an MF6 input file.
 
@@ -20,10 +21,7 @@ def load(path: str | PathLike) -> Any:
     Any
         Parsed MF6 input file structure
     """
-    path = Path(path)
-    with open(path, "r") as f:
-        data = f.read()
-    return loads(data)
+    return loads(fp.read())
 
 
 def loads(data: str) -> Any:
@@ -41,6 +39,4 @@ def loads(data: str) -> Any:
         Parsed MF6 input file structure
     """
 
-    parser = make_basic_parser()
-    transformer = BasicTransformer()
-    return transformer.transform(parser.parse(data))
+    return BASIC_TRANSFORMER.transform(BASIC_PARSER.parse(data))

flopy4/mf6/codec/writer/__init__.py

@@ -1,5 +1,5 @@
 import sys
-from os import PathLike
+from typing import IO
 
 import numpy as np
 from jinja2 import Environment, PackageLoader
@@ -32,8 +32,8 @@ def dumps(data) -> str:
         return template.render(blocks=data)
 
 
-def dump(data, path: str | PathLike) -> None:
+def dump(data, fp: IO[str]) -> None:
     template = _JINJA_ENV.get_template(_JINJA_TEMPLATE_NAME)
     iterator = template.generate(blocks=data)
-    with np.printoptions(**_PRINT_OPTIONS), open(path, "w") as f:  # type: ignore
-        f.writelines(iterator)
+    with np.printoptions(**_PRINT_OPTIONS):  # type: ignore
+        fp.writelines(iterator)

flopy4/mf6/component.py

@@ -83,13 +83,14 @@ class Component(ABC, MutableMapping):
     _load = IO(Loader)  # type: ignore
     _write = IO(Writer)  # type: ignore
 
-    filename: str = field(default=None)
+    filename: str | None = field(default=None)
 
     dfn: ClassVar[Dfn]
 
     @property
     def path(self) -> Path:
         """Get the path to the component's input file."""
+        self.filename = self.filename or self.default_filename()
         return Path.cwd() / self.filename
 
     def default_filename(self) -> str:

flopy4/mf6/constants.py

@@ -1,4 +1,5 @@
 import numpy as np
 
+MF6 = "mf6"
 FILL_DEFAULT = np.nan
 FILL_DNODATA = 1e30

flopy4/mf6/context.py

@@ -5,6 +5,7 @@
 from xattree import xattree
 
 from flopy4.mf6.component import Component
+from flopy4.mf6.constants import MF6
 from flopy4.mf6.spec import field
 
 
@@ -23,12 +24,13 @@ def __attrs_post_init__(self):
 
     @property
     def path(self) -> Path:
+        self.filename = self.filename or self.default_filename()
         return self.workspace / self.filename
 
-    def load(self, format="ascii"):
+    def load(self, format=MF6):
         with cd(self.workspace):
             super().load(format=format)
 
-    def write(self, format="ascii"):
+    def write(self, format=MF6):
         with cd(self.workspace):
             super().write(format=format)

flopy4/mf6/converter.py

@@ -326,3 +326,15 @@ def final(arr):
             set_(a, v, nn)
 
     return final(a)
+
+
+def structure(data: dict[str, Any], path: Path) -> Component:
+    component = COMPONENT_CONVERTER.structure(data, Component)
+    if isinstance(component, Context):
+        component.workspace = path.parent
+    component.filename = path.name
+    return component
+
+
+def unstructure(component: Component) -> dict[str, Any]:
+    return COMPONENT_CONVERTER.unstructure(component)