Add support for manipulating DimensionSets

foamlib/_dictionaries.py

@@ -8,6 +8,7 @@
     Mapping,
     MutableMapping,
 )
+from collections import namedtuple
 from contextlib import suppress
 
 from ._subprocesses import run_process, CalledProcessError
@@ -23,6 +24,19 @@ class FoamDictionary(
     MutableMapping[str, Union["FoamDictionary.Value", "FoamDictionary"]]
 ):
     Value = Union[str, int, float, bool, Sequence["Value"]]
+    DimensionSet = namedtuple(
+        "DimensionSet",
+        [
+            "mass",
+            "length",
+            "time",
+            "temperature",
+            "moles",
+            "current",
+            "luminous_intensity",
+        ],
+        defaults=(0, 0, 0, 0, 0, 0, 0),
+    )
 
     def __init__(self, _file: "FoamFile", _keywords: Sequence[str]) -> None:
         self._file = _file
@@ -110,6 +124,16 @@ def _parse_field(value: str) -> Value:
         else:
             raise ValueError(f"Cannot parse '{value}' as a field")
 
+    @staticmethod
+    def _parse_dimensions(value: str) -> DimensionSet:
+        if value.startswith("["):
+            assert value.endswith("]")
+            return FoamDictionary.DimensionSet(
+                *(FoamDictionary._parse_number(v) for v in value[1:-1].split())
+            )
+        else:
+            raise ValueError(f"Cannot parse '{value}' as a dimension set")
+
     @staticmethod
     def _parse(value: str) -> Value:
         with suppress(ValueError):
@@ -121,6 +145,9 @@ def _parse(value: str) -> Value:
         with suppress(ValueError):
             return FoamDictionary._parse_number(value)
 
+        with suppress(ValueError):
+            return FoamDictionary._parse_dimensions(value)
+
         with suppress(ValueError):
             return FoamDictionary._parse_sequence(value)
 
@@ -133,7 +160,9 @@ def _str_mapping(mapping: Any) -> str:
         elif isinstance(mapping, Mapping):
             m = {
                 k: FoamDictionary._str(
-                    v, assume_field=(k == "internalField" or k == "value")
+                    v,
+                    assume_field=(k == "internalField" or k == "value"),
+                    assume_dimensions=(k == "dimensions"),
                 )
                 for k, v in mapping.items()
             }
@@ -189,11 +218,27 @@ def _str_field(value: Any) -> str:
                         )
                     return f"nonuniform List<{kind}> {len(value)}{s}"
 
+    def _str_dimensions(value: Any) -> str:
+        if (
+            isinstance(value, Sequence)
+            and not isinstance(value, str)
+            and len(value) == 7
+        ):
+            return f"[{' '.join(str(v) for v in value)}]"
+        else:
+            raise TypeError(f"Not a valid dimension set: {type(value)}")
+
     @staticmethod
-    def _str(value: Union[Value, "FoamDictionary"], assume_field: bool = False) -> str:
+    def _str(
+        value: Any, assume_field: bool = False, assume_dimensions: bool = False
+    ) -> str:
         with suppress(TypeError):
             return FoamDictionary._str_mapping(value)
 
+        if isinstance(value, FoamDictionary.DimensionSet) or assume_dimensions:
+            with suppress(TypeError):
+                return FoamDictionary._str_dimensions(value)
+
         if assume_field:
             with suppress(TypeError):
                 return FoamDictionary._str_field(value)
@@ -217,7 +262,9 @@ def __getitem__(self, key: str) -> Union[Value, "FoamDictionary"]:
 
     def __setitem__(self, key: str, value: Any) -> None:
         value = self._str(
-            value, assume_field=(key == "internalField" or key == "value")
+            value,
+            assume_field=(key == "internalField" or key == "value"),
+            assume_dimensions=(key == "dimensions"),
         )
 
         if len(value) < 1000:
@@ -256,6 +303,20 @@ def __init__(self, path: Union[str, Path]) -> None:
         elif not self.path.is_file():
             raise FileNotFoundError(self.path)
 
+    @property
+    def dimensions(self) -> FoamDictionary.DimensionSet:
+        """
+        Alias of `self["dimensions"]`.
+        """
+        ret = self["dimensions"]
+        if not isinstance(ret, FoamDictionary.DimensionSet):
+            raise TypeError("dimensions is not a DimensionSet")
+        return ret
+
+    @dimensions.setter
+    def dimensions(self, value: Any) -> None:
+        self["dimensions"] = value
+
     @property
     def internal_field(self) -> FoamDictionary.Value:
         """

tests/test_dictionaries.py

@@ -118,3 +118,14 @@ def test_field(pitz: FoamCase) -> None:
         assert u == pytest.approx(u_arr)
 
     pitz.run()
+
+
+def test_dimensions(pitz: FoamCase) -> None:
+    assert pitz[0]["p"].dimensions == FoamDictionary.DimensionSet(length=2, time=-2)
+    assert pitz[0]["U"].dimensions == FoamDictionary.DimensionSet(length=1, time=-1)
+
+    pitz[0]["p"].dimensions = FoamDictionary.DimensionSet(mass=1, length=1, time=-2)
+
+    assert pitz[0]["p"].dimensions == FoamDictionary.DimensionSet(
+        mass=1, length=1, time=-2
+    )