Skip to content

Commit

Permalink
Merge pull request #249 from gerlero/typing
Browse files Browse the repository at this point in the history
Update typing
  • Loading branch information
gerlero authored Nov 2, 2024
2 parents 06a2e7b + d8fad94 commit 92aafa3
Show file tree
Hide file tree
Showing 5 changed files with 108 additions and 101 deletions.
164 changes: 84 additions & 80 deletions foamlib/_files/_files.py
Original file line number Diff line number Diff line change
Expand Up @@ -183,12 +183,12 @@ def __getitem__(
def __setitem__(
self, keywords: str | tuple[str, ...] | None, data: FoamFileBase.Data
) -> None:
with self:
if not keywords:
keywords = ()
elif not isinstance(keywords, tuple):
keywords = (keywords,)
if not keywords:
keywords = ()
elif not isinstance(keywords, tuple):
keywords = (keywords,)

with self:
try:
write_header = (
not self and "FoamFile" not in self and keywords != ("FoamFile",)
Expand Down Expand Up @@ -223,83 +223,87 @@ 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:
Expand Down
10 changes: 2 additions & 8 deletions foamlib/_files/_parsing.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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]

Expand Down
24 changes: 15 additions & 9 deletions foamlib/_files/_serialization.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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()
Expand Down
9 changes: 6 additions & 3 deletions foamlib/_files/_util.py
Original file line number Diff line number Diff line change
@@ -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
Expand All @@ -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)
2 changes: 1 addition & 1 deletion tests/test_files/test_parsing.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"

0 comments on commit 92aafa3

Please sign in to comment.