Update typing

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

Hide file tree

Showing 2 changed files with 22 additions and 15 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,23 +35,21 @@ 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."""
 
         def __init__(self, _file: "FoamFile", _keywords: Tuple[str, ...]) -> None:
             self._file = _file
             self._keywords = _keywords
 
-        def __getitem__(
-            self, keyword: str
-        ) -> Union["FoamFile.Data", "FoamFile.SubDict"]:
+        def __getitem__(self, keyword: str) -> FoamFileBase._MutableData:
             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 +94,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 +154,7 @@ def object_(self, value: str) -> None:
 
     def __getitem__(
         self, keywords: Optional[Union[str, Tuple[str, ...]]]
-    ) -> "FoamFile.Data":
+    ) -> FoamFileBase._MutableData:
         if not keywords:
             keywords = ()
         elif not isinstance(keywords, tuple):
@@ -173,7 +171,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 +391,7 @@ def value(self) -> None:
 
     def __getitem__(
         self, keywords: Optional[Union[str, Tuple[str, ...]]]
-    ) -> FoamFile.Data:
+    ) -> FoamFile._MutableData:
         if not keywords:
             keywords = ()
         elif not isinstance(keywords, tuple):