Improve parsing

gerlero · Mar 21, 2024 · 684b278 · 684b278
1 parent 584c5ff
commit 684b278
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 13 deletions.
diff --git a/foamlib/_dictionaries.py b/foamlib/_dictionaries.py
@@ -67,20 +67,21 @@ def __post_init__(self) -> None:
 _yes = Keyword("yes").set_parse_action(lambda s, loc, tks: True)
 _no = Keyword("no").set_parse_action(lambda s, loc, tks: False)
 _value = Forward()
-_list = (
-    Opt(common.integer).suppress()
-    + Literal("(").suppress()
-    + Group(ZeroOrMore(_value))
-    + Literal(")").suppress()
+_list = Opt(
+    Literal("List") + Literal("<") + common.identifier + Literal(">")
+).suppress() + (
+    (
+        Opt(common.integer).suppress()
+        + Literal("(").suppress()
+        + Group(ZeroOrMore(_value))
+        + Literal(")").suppress()
+    )
+    | (
+        common.integer + Literal("{").suppress() + _value + Literal("}").suppress()
+    ).set_parse_action(lambda s, loc, tks: [tks[1]] * tks[0])
 )
 _uniform_field = Keyword("uniform").suppress() + _value
-_nonuniform_field = (
-    Keyword("nonuniform").suppress()
-    + Literal("List<").suppress()
-    + common.identifier.suppress()
-    + Literal(">").suppress()
-    + _list
-)
+_nonuniform_field = Keyword("nonuniform").suppress() + _list
 _dimensions = (
     Literal("[").suppress() + common.number * 7 + Literal("]").suppress()
 ).set_parse_action(lambda s, loc, tks: FoamDimensionSet(*tks))

diff --git a/tests/test_dictionaries.py b/tests/test_dictionaries.py
@@ -137,14 +137,26 @@ def test_internal_field(pitz: FoamCase) -> None:
 
     pitz.clean()
 
+    p_arr = np.zeros(size)
+    U_arr = np.zeros((size, 3))
+
+    pitz[0]["p"].internal_field = p_arr
+    pitz[0]["U"].internal_field = U_arr
+
+    assert pitz[0]["p"].internal_field == pytest.approx(p_arr)
+    U = pitz[0]["U"].internal_field
+    assert isinstance(U, Sequence)
+    for u, u_arr in zip(U, U_arr):
+        assert u == pytest.approx(u_arr)
+
     p_arr = np.arange(size) * 1e-6
     U_arr = np.full((size, 3), [-1e-6, 1e-6, 0]) * np.arange(size)[:, np.newaxis]
 
     pitz[0]["p"].internal_field = p_arr
     pitz[0]["U"].internal_field = U_arr
 
     assert pitz[0]["p"].internal_field == pytest.approx(p_arr)
-    U = pitz[-1]["U"].internal_field
+    U = pitz[0]["U"].internal_field
     assert isinstance(U, Sequence)
     for u, u_arr in zip(U, U_arr):
         assert u == pytest.approx(u_arr)