Update typing

gerlero · Nov 1, 2024 · f90c277 · f90c277
1 parent f6bddbc
commit f90c277
Show file tree

Hide file tree

Showing 3 changed files with 28 additions and 21 deletions.
diff --git a/foamlib/_files/_base.py b/foamlib/_files/_base.py
@@ -6,9 +6,9 @@
     import numpy as np
 
 if sys.version_info >= (3, 9):
-    from collections.abc import Mapping, Sequence
+    from collections.abc import Mapping, MutableMapping, Sequence
 else:
-    from typing import Mapping, Sequence
+    from typing import Mapping, MutableMapping, Sequence
 
 
 class FoamFileBase:
@@ -45,21 +45,30 @@ def __post_init__(self) -> None:
         _Tensor, Sequence[_Tensor], "np.ndarray[Tuple[int, int], np.dtype[np.generic]]"
     ]
 
-    Data = Union[
+    _DataEntry = Union[
         str,
         int,
         float,
         bool,
         Dimensioned,
         DimensionSet,
         Sequence["Data"],
-        Mapping[str, "Data"],
         _Tensor,
         _Field,
     ]
+
+    Data = Union[
+        _DataEntry,
+        Mapping[str, "Data"],
+    ]
     """
     A value that can be stored in an OpenFOAM file.
     """
 
+    _MutableData = Union[
+        _DataEntry,
+        MutableMapping[str, "_MutableData"],
+    ]
+
     _Dict = Dict[str, Union["Data", "_Dict"]]
     _File = Dict[Optional[str], Union["Data", "_Dict"]]
diff --git a/foamlib/_files/_files.py b/foamlib/_files/_files.py
@@ -22,7 +22,7 @@ class FoamFile(
     FoamFileBase,
     MutableMapping[
         Optional[Union[str, Tuple[str, ...]]],
-        Union["FoamFile.Data", "FoamFile.SubDict"],
+        FoamFileBase._MutableData,
     ],
     FoamFileIO,
 ):
@@ -35,7 +35,7 @@ class FoamFile(
     """
 
     class SubDict(
-        MutableMapping[str, Union["FoamFile.Data", "FoamFile.SubDict"]],
+        MutableMapping[str, FoamFileBase._MutableData],
     ):
         """An OpenFOAM dictionary within a file as a mutable mapping."""
 
@@ -45,13 +45,13 @@ def __init__(self, _file: "FoamFile", _keywords: Tuple[str, ...]) -> None:
 
         def __getitem__(
             self, keyword: str
-        ) -> Union["FoamFile.Data", "FoamFile.SubDict"]:
+        ) -> Union[FoamFileBase._DataEntry, "FoamFile.SubDict"]:
             return self._file[(*self._keywords, keyword)]
 
         def __setitem__(
             self,
             keyword: str,
-            data: "FoamFile.Data",
+            data: FoamFileBase.Data,
         ) -> None:
             self._file[(*self._keywords, keyword)] = data
 
@@ -96,8 +96,8 @@ def as_dict(self) -> FoamFileBase._Dict:
     def version(self) -> float:
         """Alias of `self["FoamFile", "version"]`."""
         ret = self["FoamFile", "version"]
-        if not isinstance(ret, float):
-            raise TypeError("version is not a float")
+        if not isinstance(ret, (int, float)):
+            raise TypeError("version is not a number")
         return ret
 
     @version.setter
@@ -156,7 +156,7 @@ def object_(self, value: str) -> None:
 
     def __getitem__(
         self, keywords: Optional[Union[str, Tuple[str, ...]]]
-    ) -> "FoamFile.Data":
+    ) -> Union[FoamFileBase._DataEntry, "FoamFile.SubDict"]:
         if not keywords:
             keywords = ()
         elif not isinstance(keywords, tuple):
@@ -173,7 +173,7 @@ def __getitem__(
         return deepcopy(value)
 
     def __setitem__(
-        self, keywords: Optional[Union[str, Tuple[str, ...]]], data: "FoamFile.Data"
+        self, keywords: Optional[Union[str, Tuple[str, ...]]], data: FoamFileBase.Data
     ) -> None:
         with self:
             if not keywords:
@@ -393,7 +393,7 @@ def value(self) -> None:
 
     def __getitem__(
         self, keywords: Optional[Union[str, Tuple[str, ...]]]
-    ) -> FoamFile.Data:
+    ) -> Union[FoamFileBase._DataEntry, FoamFile.SubDict]:
         if not keywords:
             keywords = ()
         elif not isinstance(keywords, tuple):

diff --git a/foamlib/_files/_parsing.py b/foamlib/_files/_parsing.py
@@ -207,11 +207,11 @@ def _unpack_binary_field(
 )
 
 
-class Parsed(Mapping[Tuple[str, ...], Union[FoamFileBase.Data, EllipsisType]]):
+class Parsed(Mapping[Tuple[str, ...], Union[FoamFileBase._DataEntry, EllipsisType]]):
     def __init__(self, contents: bytes) -> None:
         self._parsed: MutableMapping[
             Tuple[str, ...],
-            Tuple[int, Union[FoamFileBase.Data, EllipsisType], int],
+            Tuple[int, Union[FoamFileBase._DataEntry, EllipsisType], int],
         ] = {}
         for parse_result in _FILE.parse_string(
             contents.decode("latin-1"), parse_all=True
@@ -225,11 +225,11 @@ def __init__(self, contents: bytes) -> None:
     def _flatten_result(
         parse_result: ParseResults, *, _keywords: Tuple[str, ...] = ()
     ) -> Mapping[
-        Tuple[str, ...], Tuple[int, Union[FoamFileBase.Data, EllipsisType], int]
+        Tuple[str, ...], Tuple[int, Union[FoamFileBase._DataEntry, EllipsisType], int]
     ]:
         ret: MutableMapping[
             Tuple[str, ...],
-            Tuple[int, Union[FoamFileBase.Data, EllipsisType], int],
+            Tuple[int, Union[FoamFileBase._DataEntry, EllipsisType], int],
         ] = {}
         start = parse_result.locn_start
         assert isinstance(start, int)
@@ -257,7 +257,7 @@ def _flatten_result(
 
     def __getitem__(
         self, keywords: Union[str, Tuple[str, ...]]
-    ) -> Union[FoamFileBase.Data, EllipsisType]:
+    ) -> Union[FoamFileBase._DataEntry, EllipsisType]:
         if isinstance(keywords, str):
             keywords = (keywords,)
 
@@ -267,11 +267,9 @@ def __getitem__(
     def put(
         self,
         keywords: Tuple[str, ...],
-        data: Union[FoamFileBase.Data, EllipsisType],
+        data: Union[FoamFileBase._DataEntry, EllipsisType],
         content: bytes,
     ) -> None:
-        assert not isinstance(data, Mapping)
-
         start, end = self.entry_location(keywords, missing_ok=True)
 
         diff = len(content) - (end - start)