generalize!

Author: wpbonelli

flopy4/mf6/codec/reader/transformer.py

@@ -1,9 +1,12 @@
+from collections import ChainMap
+from collections.abc import Mapping
 from pathlib import Path
 from typing import Any
 
 import numpy as np
 import xarray as xr
 from lark import Token, Transformer
+from modflow_devtools.dfn import _SCALAR_TYPES, Dfn, get_blocks, get_fields
 
 
 class BasicTransformer(Transformer):
@@ -59,7 +62,16 @@ def INT(self, token: Token) -> int:
 class TypedTransformer(Transformer):
     """Type-aware transformer for MF6 input files."""
 
-    def start(self, items: list[Any]) -> dict:
+    def __init__(self, visit_tokens=True, dfn: Dfn = None):
+        super().__init__(visit_tokens)
+        self.dfn = dfn
+        self.blocks = get_blocks(dfn) if dfn else None
+        self.fields = get_fields(dfn) if dfn else None
+
+    def start(self, items: list[Any]) -> Mapping:
+        return ChainMap(*items)
+
+    def block(self, items: list[Any]) -> dict:
         return items[0]
 
     def array(self, items: list[Any]) -> dict:
@@ -157,3 +169,19 @@ def try_create_dataarray(array_info: dict) -> dict:
             case "external":
                 pass
         return array_info
+
+    def __default__(self, data, children, meta):
+        if self.blocks is None or self.fields is None:
+            return super().__default__(data, children, meta)
+        if data.endswith("_block") and (block_name := data[:-6]) in self.blocks:
+            return {block_name: children[0]}
+        elif data.endswith("_vars"):
+            return {item[0].lower(): item[1] for item in children}
+        elif (field := self.fields.get(data, None)) is not None:
+            if field["type"] == "keyword":
+                return data, True
+            elif field["type"] in _SCALAR_TYPES and field.get("shape", None):
+                return data, TypedTransformer.try_create_dataarray(children[0])
+            else:
+                return data, children[0]
+        return super().__default__(data, children, meta)

test/test_parser.py

@@ -1,7 +1,5 @@
 import os
-from collections import ChainMap
 from pathlib import Path
-from typing import Any
 
 import numpy as np
 import xarray as xr
@@ -151,48 +149,22 @@ def test_transform_full_component():
 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(visit_tokens=False)
+    transformer = TypedTransformer(
+        visit_tokens=False,
+        dfn={
+            "options": {
+                "r2d2": {"name": "r2d2", "type": "keyword"},
+                "b": {"name": "b", "type": "string"},
+                "c": {"name": "c", "type": "integer"},
+                "p": {"name": "p", "type": "double"},
+            },
+            "arrays": {
+                "x": {"name": "x", "type": "double", "shape": None},
+                "y": {"name": "y", "type": "array", "shape": None},
+                "z": {"name": "z", "type": "array", "shape": None},
+            },
+        },
+    )
     result = transformer.transform(
         parser.parse("""
 BEGIN OPTIONS