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..7d28ae2 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.""" @@ -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 index f0a0056..4d32ead 100644 --- 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)