From 87ec85c35fdd7bceffdd2d7f4dae1edfe6fd344f Mon Sep 17 00:00:00 2001 From: Gabriel Gerlero Date: Sat, 2 Nov 2024 00:15:16 -0300 Subject: [PATCH] Update typing --- foamlib/_files/_files.py | 158 ++++++++++++++++--------------- foamlib/_files/_parsing.py | 10 +- foamlib/_files/_serialization.py | 24 +++-- foamlib/_files/_util.py | 9 +- tests/test_files/test_parsing.py | 2 +- 5 files changed, 107 insertions(+), 96 deletions(-) diff --git a/foamlib/_files/_files.py b/foamlib/_files/_files.py index c389180..b7257f7 100644 --- a/foamlib/_files/_files.py +++ b/foamlib/_files/_files.py @@ -223,83 +223,91 @@ def __setitem__( if ( kind in (Kind.FIELD, Kind.BINARY_FIELD) ) and self.class_ == "dictionary": - if not is_sequence(data): - class_ = "volScalarField" - elif (len(data) == 3 and not is_sequence(data[0])) or len(data[0]) == 3: - class_ = "volVectorField" - elif (len(data) == 6 and not is_sequence(data[0])) or len(data[0]) == 6: - class_ = "volSymmTensorField" - elif (len(data) == 9 and not is_sequence(data[0])) or len(data[0]) == 9: - class_ = "volTensorField" - else: - class_ = "volScalarField" - - self.class_ = class_ - self[keywords] = data + if isinstance(data, (int, float)): + self.class_ = "volScalarField" + + elif is_sequence(data) and data: + if isinstance(data[0], (int, float)): + if len(data) == 3: + self.class_ = "volVectorField" + elif len(data) == 6: + self.class_ = "volSymmTensorField" + elif len(data) == 9: + self.class_ = "volTensorField" + elif ( + is_sequence(data[0]) + and data[0] + and isinstance(data[0][0], (int, float)) + ): + if len(data[0]) == 3: + self.class_ = "volVectorField" + elif len(data[0]) == 6: + self.class_ = "volSymmTensorField" + elif len(data[0]) == 9: + self.class_ = "volTensorField" + + parsed = self._get_parsed(missing_ok=True) + + start, end = parsed.entry_location(keywords, missing_ok=True) + + before = b"" + if parsed.contents[:start] and not parsed.contents[:start].endswith( + b"\n" + ): + before = b"\n" + if ( + parsed.contents[:start] + and len(keywords) <= 1 + and not parsed.contents[:start].endswith(b"\n\n") + ): + before = b"\n\n" + + after = b"" + if parsed.contents[end:].startswith(b"}"): + after = b" " * (len(keywords) - 2) + if not parsed.contents[end:] or not parsed.contents[end:].startswith( + b"\n" + ): + after = b"\n" + after + + indentation = b" " * (len(keywords) - 1) + + if isinstance(data, Mapping): + if isinstance(data, (FoamFile, FoamFile.SubDict)): + data = data.as_dict() + + parsed.put( + keywords, + ..., + before + + indentation + + dumps(keywords[-1]) + + b"\n" + + indentation + + b"{\n" + + indentation + + b"}" + + after, + ) + + for k, v in data.items(): + self[(*keywords, k)] = v + + elif keywords: + parsed.put( + keywords, + data, + before + + indentation + + dumps(keywords[-1]) + + b" " + + dumps(data, kind=kind) + + b";" + + after, + ) else: - parsed = self._get_parsed(missing_ok=True) - - start, end = parsed.entry_location(keywords, missing_ok=True) - - before = b"" - if parsed.contents[:start] and not parsed.contents[:start].endswith( - b"\n" - ): - before = b"\n" - if ( - parsed.contents[:start] - and len(keywords) <= 1 - and not parsed.contents[:start].endswith(b"\n\n") - ): - before = b"\n\n" - - after = b"" - if parsed.contents[end:].startswith(b"}"): - after = b" " * (len(keywords) - 2) - if not parsed.contents[end:] or not parsed.contents[end:].startswith( - b"\n" - ): - after = b"\n" + after - - indentation = b" " * (len(keywords) - 1) - - if isinstance(data, Mapping): - if isinstance(data, (FoamFile, FoamFile.SubDict)): - data = data.as_dict() - - parsed.put( - keywords, - ..., - before - + indentation - + dumps(keywords[-1]) - + b"\n" - + indentation - + b"{\n" - + indentation - + b"}" - + after, - ) - - for k, v in data.items(): - self[(*keywords, k)] = v - - elif keywords: - parsed.put( - keywords, - data, - before - + indentation - + dumps(keywords[-1]) - + b" " - + dumps(data, kind=kind) - + b";" - + after, - ) - - else: - parsed.put(keywords, data, before + dumps(data, kind=kind) + after) + parsed.put(keywords, data, before + dumps(data, kind=kind) + after) def __delitem__(self, keywords: str | tuple[str, ...] | None) -> None: if not keywords: diff --git a/foamlib/_files/_parsing.py b/foamlib/_files/_parsing.py index ac55c59..8618712 100644 --- a/foamlib/_files/_parsing.py +++ b/foamlib/_files/_parsing.py @@ -258,11 +258,8 @@ def _flatten_result( return ret def __getitem__( - self, keywords: str | tuple[str, ...] + self, keywords: tuple[str, ...] ) -> FoamFileBase._DataEntry | EllipsisType: - if isinstance(keywords, str): - keywords = (keywords,) - _, data, _ = self._parsed[keywords] return data @@ -290,10 +287,7 @@ def put( if keywords != k and keywords == k[: len(keywords)]: del self._parsed[k] - def __delitem__(self, keywords: str | tuple[str, ...]) -> None: - if isinstance(keywords, str): - keywords = (keywords,) - + def __delitem__(self, keywords: tuple[str, ...]) -> None: start, end = self.entry_location(keywords) del self._parsed[keywords] diff --git a/foamlib/_files/_serialization.py b/foamlib/_files/_serialization.py index 53817bc..cfe2784 100644 --- a/foamlib/_files/_serialization.py +++ b/foamlib/_files/_serialization.py @@ -4,11 +4,12 @@ import itertools import sys from enum import Enum, auto +from typing import cast if sys.version_info >= (3, 9): - from collections.abc import Mapping + from collections.abc import Mapping, Sequence else: - from typing import Mapping + from typing import Mapping, Sequence from ._base import FoamFileBase from ._util import is_sequence @@ -65,21 +66,26 @@ def dumps( return b"uniform " + dumps(data, kind=Kind.SINGLE_ENTRY) if kind in (Kind.FIELD, Kind.BINARY_FIELD) and is_sequence(data): - if isinstance(data[0], (int, float)): + if data and isinstance(data[0], (int, float)): tensor_kind = b"scalar" - elif len(data[0]) == 3: - tensor_kind = b"vector" - elif len(data[0]) == 6: - tensor_kind = b"symmTensor" - elif len(data[0]) == 9: - tensor_kind = b"tensor" + elif is_sequence(data[0]) and data[0] and isinstance(data[0][0], (int, float)): + if len(data[0]) == 3: + tensor_kind = b"vector" + elif len(data[0]) == 6: + tensor_kind = b"symmTensor" + elif len(data[0]) == 9: + tensor_kind = b"tensor" + else: + return dumps(data) else: return dumps(data) if kind == Kind.BINARY_FIELD: if tensor_kind == b"scalar": + data = cast(Sequence[float], data) contents = b"(" + array.array("d", data).tobytes() + b")" else: + data = cast(Sequence[Sequence[float]], data) contents = ( b"(" + array.array("d", itertools.chain.from_iterable(data)).tobytes() diff --git a/foamlib/_files/_util.py b/foamlib/_files/_util.py index 01ac2ae..f802ad0 100644 --- a/foamlib/_files/_util.py +++ b/foamlib/_files/_util.py @@ -1,7 +1,7 @@ from __future__ import annotations import sys -from typing import Any +from typing import TYPE_CHECKING if sys.version_info >= (3, 9): from collections.abc import Sequence @@ -13,8 +13,11 @@ else: from typing_extensions import TypeGuard +if TYPE_CHECKING: + from ._base import FoamFileBase + def is_sequence( - value: Any, -) -> TypeGuard[Sequence[Any]]: + value: FoamFileBase.Data, +) -> TypeGuard[Sequence[FoamFileBase.Data]]: return isinstance(value, Sequence) and not isinstance(value, str) diff --git a/tests/test_files/test_parsing.py b/tests/test_files/test_parsing.py index 237c7a6..f434779 100644 --- a/tests/test_files/test_parsing.py +++ b/tests/test_files/test_parsing.py @@ -71,4 +71,4 @@ def test_parse_value() -> None: {"(air and water)": {"type": "constant", "sigma": 0.07}} ] assert Parsed(b"[]")[()] == FoamFile.DimensionSet() - assert Parsed(b"object f.1;")["object"] == "f.1" + assert Parsed(b"object f.1;")[("object",)] == "f.1"