diff --git a/foamlib/_files/_base.py b/foamlib/_files/_base.py index 1ff395d..ed15494 100644 --- 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,7 +45,7 @@ def __post_init__(self) -> None: _Tensor, Sequence[_Tensor], "np.ndarray[Tuple[int, int], np.dtype[np.generic]]" ] - Data = Union[ + _DataEntry = Union[ str, int, float, @@ -53,13 +53,22 @@ def __post_init__(self) -> None: 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 index 7f124e6..41e64b7 100644 --- 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.""" @@ -43,15 +43,13 @@ 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):