diff --git a/foamlib/_cases/_async.py b/foamlib/_cases/_async.py index dae4011..83f7117 100644 --- a/foamlib/_cases/_async.py +++ b/foamlib/_cases/_async.py @@ -3,7 +3,7 @@ import os import sys from contextlib import asynccontextmanager -from typing import Any, Callable, Optional, TypeVar, Union +from typing import Any, Callable, Optional, TypeVar, Union, overload if sys.version_info >= (3, 9): from collections.abc import ( @@ -23,6 +23,8 @@ import aioshutil +from .._files import FoamFieldFile +from ._base import FoamCaseBase from ._run import FoamCaseRunBase from ._subprocess import run_async from ._util import ValuedGenerator, awaitableasynccontextmanager @@ -42,6 +44,20 @@ class AsyncFoamCase(FoamCaseRunBase): :param path: The path to the case directory. """ + class TimeDirectory(FoamCaseRunBase.TimeDirectory): + @property + def _case(self) -> "AsyncFoamCase": + return AsyncFoamCase(self.path.parent) + + async def cell_centers(self) -> FoamFieldFile: + """Write and return the cell centers.""" + calls = ValuedGenerator(self._cell_centers_calls()) + + for coro in calls: + await coro + + return calls.value + max_cpus = multiprocessing.cpu_count() """ Maximum number of CPUs to use for running `AsyncFoamCase`s concurrently. Defaults to the number of CPUs on the system. @@ -106,6 +122,23 @@ async def clean(self, *, check: bool = False) -> None: for coro in self._clean_calls(check=check): await coro + @overload + def __getitem__( + self, index: Union[int, float, str] + ) -> "AsyncFoamCase.TimeDirectory": ... + + @overload + def __getitem__(self, index: slice) -> Sequence["AsyncFoamCase.TimeDirectory"]: ... + + def __getitem__( + self, index: Union[int, slice, float, str] + ) -> Union["AsyncFoamCase.TimeDirectory", Sequence["AsyncFoamCase.TimeDirectory"]]: + ret = super().__getitem__(index) + if isinstance(ret, FoamCaseBase.TimeDirectory): + return AsyncFoamCase.TimeDirectory(ret) + else: + return [AsyncFoamCase.TimeDirectory(r) for r in ret] + async def run( self, cmd: Optional[Union[Sequence[Union[str, "os.PathLike[str]"]], str]] = None, diff --git a/foamlib/_cases/_base.py b/foamlib/_cases/_base.py index b86769f..bec3307 100644 --- a/foamlib/_cases/_base.py +++ b/foamlib/_cases/_base.py @@ -40,6 +40,10 @@ class TimeDirectory(Set[FoamFieldFile]): def __init__(self, path: Union["os.PathLike[str]", str]): self.path = Path(path).absolute() + @property + def _case(self) -> "FoamCaseBase": + return FoamCaseBase(self.path.parent) + @property def time(self) -> float: """The time that corresponds to this directory.""" @@ -58,7 +62,7 @@ def __getitem__(self, key: str) -> FoamFieldFile: def __contains__(self, obj: object) -> bool: if isinstance(obj, FoamFieldFile): - return obj.path.parent == self.path + return obj.path.parent == self.path and obj.path.is_file() elif isinstance(obj, str): return (self.path / obj).is_file() or ( self.path / f"{obj}.gz" diff --git a/foamlib/_cases/_run.py b/foamlib/_cases/_run.py index 2d5ba95..e2c2098 100644 --- a/foamlib/_cases/_run.py +++ b/foamlib/_cases/_run.py @@ -40,11 +40,34 @@ else: from typing_extensions import Self +from .._files import FoamFieldFile from ._base import FoamCaseBase from ._subprocess import DEVNULL, STDOUT class FoamCaseRunBase(FoamCaseBase): + class TimeDirectory(FoamCaseBase.TimeDirectory): + @abstractmethod + def cell_centers( + self, + ) -> Union[FoamFieldFile, Coroutine[None, None, FoamFieldFile]]: + raise NotImplementedError + + @property + @abstractmethod + def _case(self) -> "FoamCaseRunBase": + raise NotImplementedError + + def _cell_centers_calls(self) -> Generator[Any, None, FoamFieldFile]: + ret = self["C"] + + if ret not in self: + yield self._case.run( + ["postProcess", "-func", "writeCellCentres", "-time", self.name] + ) + + return ret + def __delitem__(self, key: Union[int, float, str]) -> None: shutil.rmtree(self[key].path) diff --git a/foamlib/_cases/_sync.py b/foamlib/_cases/_sync.py index 4208631..2c95ffb 100644 --- a/foamlib/_cases/_sync.py +++ b/foamlib/_cases/_sync.py @@ -3,13 +3,7 @@ import sys from pathlib import Path from types import TracebackType -from typing import ( - Any, - Callable, - Optional, - Type, - Union, -) +from typing import Any, Callable, Optional, Type, Union, overload if sys.version_info >= (3, 9): from collections.abc import Collection, Sequence @@ -21,6 +15,8 @@ else: from typing_extensions import Self +from .._files import FoamFieldFile +from ._base import FoamCaseBase from ._run import FoamCaseRunBase from ._subprocess import run_sync from ._util import ValuedGenerator @@ -37,6 +33,22 @@ class FoamCase(FoamCaseRunBase): :param path: The path to the case directory. """ + class TimeDirectory(FoamCaseRunBase.TimeDirectory): + @property + def _case(self) -> "FoamCase": + return FoamCase(self.path.parent) + + def cell_centers(self) -> FoamFieldFile: + """Write and return the cell centers.""" + calls = ValuedGenerator(self._cell_centers_calls()) + + for _ in calls: + pass + + print(calls.value) + + return calls.value + def __init__(self, path: Union["os.PathLike[str]", str] = Path()): super().__init__(path) @@ -67,6 +79,23 @@ def _copytree( ) -> None: shutil.copytree(src, dest, symlinks=symlinks, ignore=ignore) + @overload + def __getitem__( + self, index: Union[int, float, str] + ) -> "FoamCase.TimeDirectory": ... + + @overload + def __getitem__(self, index: slice) -> Sequence["FoamCase.TimeDirectory"]: ... + + def __getitem__( + self, index: Union[int, slice, float, str] + ) -> Union["FoamCase.TimeDirectory", Sequence["FoamCase.TimeDirectory"]]: + ret = super().__getitem__(index) + if isinstance(ret, FoamCaseBase.TimeDirectory): + return FoamCase.TimeDirectory(ret) + else: + return [FoamCase.TimeDirectory(r) for r in ret] + def __enter__(self) -> Self: return self diff --git a/tests/test_cases/test_cavity.py b/tests/test_cases/test_cavity.py index ea29dda..ec62ecf 100644 --- a/tests/test_cases/test_cavity.py +++ b/tests/test_cases/test_cavity.py @@ -56,3 +56,10 @@ def test_double_clean(cavity: FoamCase) -> None: cavity.clean() cavity.clean(check=True) cavity.run(parallel=False) + + +def test_cell_centers(cavity: FoamCase) -> None: + cavity.block_mesh() + C = cavity[0].cell_centers() + assert isinstance(C.internal_field, list) + assert len(C.internal_field) == 400 diff --git a/tests/test_cases/test_cavity_async.py b/tests/test_cases/test_cavity_async.py index f3b1bb6..151d74f 100644 --- a/tests/test_cases/test_cavity_async.py +++ b/tests/test_cases/test_cavity_async.py @@ -61,6 +61,14 @@ async def test_double_clean(cavity: AsyncFoamCase) -> None: await cavity.run(parallel=False) +@pytest.mark.asyncio +async def test_cell_centers(cavity: AsyncFoamCase) -> None: + await cavity.block_mesh() + C = await cavity[0].cell_centers() + assert isinstance(C.internal_field, list) + assert len(C.internal_field) == 400 + + def test_map(cavity: AsyncFoamCase) -> None: async def f(x: Sequence[float]) -> float: async with cavity.clone() as clone: