diff --git a/docs/dev/sdd.md b/docs/dev/sdd.md index b4406ae4..55208d04 100644 --- a/docs/dev/sdd.md +++ b/docs/dev/sdd.md @@ -127,17 +127,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 diff --git a/flopy4/mf6/__init__.py b/flopy4/mf6/__init__.py index f12fb2d4..e92f50cc 100644 --- a/flopy4/mf6/__init__.py +++ b/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) diff --git a/flopy4/mf6/codec/reader/__init__.py b/flopy4/mf6/codec/reader/__init__.py index 64f08ef5..c61d3c1d 100644 --- a/flopy4/mf6/codec/reader/__init__.py +++ b/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)) diff --git a/flopy4/mf6/codec/writer/__init__.py b/flopy4/mf6/codec/writer/__init__.py index f6af3ba4..be7d4c15 100644 --- a/flopy4/mf6/codec/writer/__init__.py +++ b/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) diff --git a/flopy4/mf6/component.py b/flopy4/mf6/component.py index efffa7c9..4a32d89e 100644 --- a/flopy4/mf6/component.py +++ b/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: diff --git a/flopy4/mf6/constants.py b/flopy4/mf6/constants.py index b9e58ed5..bcd2aaaf 100644 --- a/flopy4/mf6/constants.py +++ b/flopy4/mf6/constants.py @@ -1,4 +1,5 @@ import numpy as np +MF6 = "mf6" FILL_DEFAULT = np.nan FILL_DNODATA = 1e30 diff --git a/flopy4/mf6/context.py b/flopy4/mf6/context.py index 7e98598d..f79b9c6e 100644 --- a/flopy4/mf6/context.py +++ b/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) diff --git a/flopy4/mf6/converter.py b/flopy4/mf6/converter.py index 16db8e35..f8b21d1c 100644 --- a/flopy4/mf6/converter.py +++ b/flopy4/mf6/converter.py @@ -323,3 +323,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)