handle scalars, test component-specific parsing/transformation

Author: wpbonelli

flopy4/mf6/codec/reader/grammar/basic.lark

@@ -1,7 +1,7 @@
 start: [WS] [_NL*] (block [[WS] _NL*])+ [WS]
-block: "begin"i block_name _NL _content "end"i block_name _NL+
+block: "begin"i block_name _NL _list "end"i block_name _NL+
 block_name: CNAME [INT]
-_content: line*
+_list: line*
 line: [WS] item+ _NL+
 item: word | NUMBER
 word: /[a-zA-Z0-9._'~,-\\(\\)]+/

flopy4/mf6/codec/reader/grammar/array.lark renamed to flopy4/mf6/codec/reader/grammar/typed.lark

@@ -1,22 +1,28 @@
-start: readarray
-readarray: (single_array | layered_array)
-single_array: [netcdf] array
-layered_array: layered [netcdf] array+
+integer: _integer
+double: _number
+string: ESCAPED_STRING | record
+record: _token+ _NL
+list: record*
+array: (single_array | layered_array)
+single_array: [netcdf] readarray
+layered_array: layered [netcdf] readarray+
 layered: "layered"i
 netcdf: "netcdf"i
-array: control [data]
+readarray: control [data]
 control: constant | internal | external
 constant: "constant"i _number
 internal: "internal"i [factor] [iprn]
 external: "open/close"i filename [factor] [binary] [iprn]
 factor: "factor"i _number
-iprn: "iprn"i _signed_integer
-binary: "(" "binary"i ")"
+iprn: "iprn"i _integer
+binary: "(binary)"i
 filename: ESCAPED_STRING | _word
 data: _number+
-_number: SIGNED_NUMBER | NUMBER
-_signed_integer: SIGNED_INT | INT
+
 _word: /[a-zA-Z0-9._'~,-\\(\\)]+/
+_number: SIGNED_NUMBER | NUMBER
+_integer: SIGNED_INT | INT
+_token: _word | _number
 
 %import common.NEWLINE -> _NL
 %import common.WS

flopy4/mf6/codec/reader/parser.py

@@ -8,10 +8,3 @@ def make_basic_parser() -> Lark:
     with open(grammar_path, "r") as f:
         grammar = f.read()
     return Lark(grammar, parser="lalr", debug=True)
-
-
-def make_array_parser() -> Lark:
-    grammar_path = Path(__file__).parent / "grammar" / "array.lark"
-    with open(grammar_path, "r") as f:
-        grammar = f.read()
-    return Lark(grammar, parser="lalr", debug=True)

flopy4/mf6/codec/reader/transformer.py

@@ -27,6 +27,9 @@ def block(self, items: list[Any]) -> dict[str, Any]:
     def block_name(self, items: list[Any]) -> str:
         return " ".join([str(item) for item in items if item is not None])
 
+    def _list(self, items: list[Any]) -> list[Any]:
+        return items[0] if items else []
+
     def line(self, items: list[Any]) -> list[Any]:
         return items[1:]
 
@@ -53,19 +56,13 @@ def INT(self, token: Token) -> int:
         return int(token)
 
 
-class ArrayTransformer(Transformer):
-    """
-    Transformer for MF6 array input format. Returns xarray DataArrays
-    for internal/constant arrays and Path objects for external arrays,
-    inside a dictionary which also contains control information. This
-    is a first step towards a smarter parser/transformer for the full
-    MF6 input format specification.
-    """
+class TypedTransformer(Transformer):
+    """Type-aware transformer for MF6 input files."""
 
     def start(self, items: list[Any]) -> dict:
         return items[0]
 
-    def readarray(self, items: list[Any]) -> dict:
+    def array(self, items: list[Any]) -> dict:
         infos = items[0]
         if isinstance(infos, list):
             data = xr.concat([info["data"] for info in infos if "data" in info], dim="layer")
@@ -82,7 +79,7 @@ def single_array(self, items: list[Any]) -> dict:
         info = items[-1]
         if netcdf:
             info["netcdf"] = netcdf
-        return ArrayTransformer.try_create_dataarray(info)
+        return TypedTransformer.try_create_dataarray(info)
 
     def layered_array(self, items: list[Any]) -> list[dict]:
         netcdf = items[0]
@@ -92,10 +89,10 @@ def layered_array(self, items: list[Any]) -> list[dict]:
                 continue
             if netcdf:
                 info["netcdf"] = netcdf
-            infos.append(ArrayTransformer.try_create_dataarray(info))
+            infos.append(TypedTransformer.try_create_dataarray(info))
         return infos
 
-    def array(self, items: list[Any]) -> dict[str, Any]:
+    def readarray(self, items: list[Any]) -> dict[str, Any]:
         control = items[0]
         data = items[1] if len(items) > 1 else None
         if (value := control.get("value", None)) is not None:
@@ -116,7 +113,11 @@ def internal(self, items: list[Any]) -> dict[str, Any]:
         return result
 
     def external(self, items: list[Any]) -> dict[str, Any]:
-        return {"type": "external", "value": items[0]}
+        result = {"type": "external", "value": items[0]}
+        for item in items[1:]:
+            if item is not None:
+                result.update(item)
+        return result
 
     def factor(self, items: list[Any]) -> dict[str, float]:
         return {"factor": items[0]}
@@ -130,6 +131,15 @@ def binary(self, items: list[Any]) -> dict[str, bool]:
     def filename(self, items: list[Any]) -> Path:
         return Path(items[0])
 
+    def string(self, items: list[Any]) -> str:
+        return items[0].strip("\"'")
+
+    def integer(self, items: list[Any]) -> int:
+        return int(items[0])
+
+    def double(self, items: list[Any]) -> float:
+        return float(items[0])
+
     def data(self, items: list[Any]) -> np.ndarray:
         return np.array(items)
 
@@ -157,17 +167,12 @@ def ESCAPED_STRING(self, token: Token) -> str:
 
     @staticmethod
     def try_create_dataarray(array_info: dict) -> dict:
-        """Create an xarray DataArray from MF6 array information."""
         control = array_info["control"]
         match control["type"]:
             case "constant":
-                data = control["value"]
-                array_info["data"] = xr.DataArray(data=data)
+                array_info["data"] = xr.DataArray(data=control["value"])
             case "internal":
-                data = array_info["data"]
-                factor = control.get("factor", 1.0)
-                data = data * factor
-                array_info["data"] = xr.DataArray(data=data)
+                array_info["data"] = xr.DataArray(data=array_info["data"])
             case "external":
                 pass
         return array_info

test/test_parse_array.py

@@ -1,14 +1,27 @@
+import os
+from collections import ChainMap
 from pathlib import Path
+from typing import Any
 
 import numpy as np
 import xarray as xr
+from lark import Lark
 
-from flopy4.mf6.codec.reader.parser import make_array_parser
-from flopy4.mf6.codec.reader.transformer import ArrayTransformer
+from flopy4.mf6.codec.reader.transformer import TypedTransformer
+
+PROJ_ROOT_PATH = Path(__file__).parents[1]
+BASE_GRAMMAR_PATH = (
+    PROJ_ROOT_PATH / "flopy4" / "mf6" / "codec" / "reader" / "grammar" / "typed.lark"
+)
+
+
+def make_typed_parser(grammar: str):
+    with open(BASE_GRAMMAR_PATH, "r") as f:
+        return Lark(grammar + os.linesep + f.read(), parser="lalr", debug=True)
 
 
 def test_parse_internal_array():
-    parser = make_array_parser()
+    parser = make_typed_parser("start: array")
     tree = parser.parse("""
 INTERNAL FACTOR 1.0 IPRN 3
 1.2 3.7 9.3 4.2 2.2 9.9 1.0 
@@ -20,7 +33,7 @@ def test_parse_internal_array():
 
 
 def test_parse_layered_array():
-    parser = make_array_parser()
+    parser = make_typed_parser("start: array")
     tree = parser.parse("""
 LAYERED
 CONSTANT 1.0
@@ -34,32 +47,32 @@ def test_parse_layered_array():
 
 
 def test_parse_constant_array():
-    parser = make_array_parser()
+    parser = make_typed_parser("start: array")
     tree = parser.parse("""
 CONSTANT 1.0
     """)
     print(tree.pretty())
 
 
 def test_parse_external_array_no_quotation_marks():
-    parser = make_array_parser()
+    parser = make_typed_parser("start: array")
     tree = parser.parse("""
 OPEN/CLOSE some.file
     """)
     print(tree.pretty())
 
 
 def test_parse_external_array_with_quotation_marks():
-    parser = make_array_parser()
+    parser = make_typed_parser("start: array")
     tree = parser.parse("""
 OPEN/CLOSE "some.file"
     """)
     print(tree.pretty())
 
 
 def test_transform_internal_array():
-    parser = make_array_parser()
-    transformer = ArrayTransformer()
+    parser = make_typed_parser("start: array")
+    transformer = TypedTransformer()
     result = transformer.transform(
         parser.parse("""
 INTERNAL FACTOR 1.5 IPRN 3
@@ -76,8 +89,8 @@ def test_transform_internal_array():
 
 
 def test_transform_constant_array():
-    parser = make_array_parser()
-    transformer = ArrayTransformer()
+    parser = make_typed_parser("start: array")
+    transformer = TypedTransformer()
     result = transformer.transform(
         parser.parse("""
 CONSTANT 42.5
@@ -88,8 +101,8 @@ def test_transform_constant_array():
 
 
 def test_transform_external_array():
-    parser = make_array_parser()
-    transformer = ArrayTransformer()
+    parser = make_typed_parser("start: array")
+    transformer = TypedTransformer()
     result = transformer.transform(
         parser.parse("""
 OPEN/CLOSE "data/heads.dat" FACTOR 1.0 (BINARY)
@@ -100,8 +113,8 @@ def test_transform_external_array():
 
 
 def test_transform_layered_array():
-    parser = make_array_parser()
-    transformer = ArrayTransformer()
+    parser = make_typed_parser("start: array")
+    transformer = TypedTransformer()
     result = transformer.transform(
         parser.parse("""
 LAYERED
@@ -119,3 +132,93 @@ def test_transform_layered_array():
     assert result["data"].shape == (2, 8)
     assert result["data"].dims == ("layer", "dim_0")
     assert np.array_equal(result["data"][0], np.ones((8,)))
+
+
+def test_transform_full_component():
+    grammar = """
+start: block*
+block: options_block | arrays_block
+options_block: "begin"i "options"i options_vars "end"i "options"i
+options_vars: (r2d2 | b | c | p)*
+arrays_block: "begin"i "arrays"i arrays_vars "end"i "arrays"i
+arrays_vars: (x | y | z)*
+r2d2: "r2d2"i // keyword
+b: "b"i string
+c: "c"i integer
+p: "p"i double
+x: "x"i array
+y: "y"i array
+z: "z"i array
+"""
+    parser = make_typed_parser(grammar)
+
+    class BlockTransformer(TypedTransformer):
+        def start(self, items: list[Any]) -> dict:
+            return ChainMap(*items)
+
+        def block(self, items: list[Any]) -> dict:
+            return items[0]
+
+        def options_block(self, items: list[Any]) -> dict:
+            return {"options": items[0]}
+
+        def arrays_block(self, items: list[Any]) -> dict:
+            return {"arrays": items[0]}
+
+        def options_vars(self, items: list[Any]) -> dict:
+            return {item[0].lower(): item[1] for item in items}
+
+        def arrays_vars(self, items: list[Any]) -> dict:
+            return {item[0].lower(): item[1] for item in items}
+
+        def r2d2(self, _: list[Any]) -> bool:
+            return "r2d2", True
+
+        def b(self, items: list[Any]) -> tuple[str, str]:
+            return "b", items[0]
+
+        def c(self, items: list[Any]) -> tuple[str, int]:
+            return "c", items[0]
+
+        def p(self, items: list[Any]) -> tuple[str, float]:
+            return "p", items[0]
+
+        def x(self, items: list[Any]) -> tuple[str, dict]:
+            return "x", TypedTransformer.try_create_dataarray(items[0])
+
+        def y(self, items: list[Any]) -> tuple[str, dict]:
+            return "y", TypedTransformer.try_create_dataarray(items[0])
+
+        def z(self, items: list[Any]) -> tuple[str, dict]:
+            return "z", TypedTransformer.try_create_dataarray(items[0])
+
+    transformer = BlockTransformer()
+    result = transformer.transform(
+        parser.parse("""
+BEGIN OPTIONS
+    R2D2
+    B "nice said"
+    C 3
+    P 0.
+END OPTIONS
+BEGIN ARRAYS
+    X CONSTANT 1.0
+    Y INTERNAL 4.0 5.0 6.0
+    Z OPEN/CLOSE "data/z.dat" FACTOR 1.0 (BINARY)
+END ARRAYS
+""")
+    )
+    assert "options" in result
+    assert "arrays" in result
+    assert result["options"]["r2d2"] is True
+    assert result["options"]["b"] == "nice said"
+    assert result["options"]["c"] == 3
+    assert result["options"]["p"] == 0.0
+    assert result["arrays"]["x"]["control"]["type"] == "constant"
+    assert np.array_equal(result["arrays"]["x"]["data"], np.array(1.0))
+    assert result["arrays"]["y"]["control"]["type"] == "internal"
+    assert np.array_equal(result["arrays"]["y"]["data"], np.array([4.0, 5.0, 6.0]))
+    assert result["arrays"]["z"]["control"]["type"] == "external"
+    assert result["arrays"]["z"]["control"]["factor"] == 1.0
+    assert result["arrays"]["z"]["control"]["binary"] is True
+    assert result["arrays"]["z"]["data"] == Path("data/z.dat")