foamlib 0.4.4__py3-none-any.whl → 0.5.1__py3-none-any.whl

foamlib/_cases/_recipes.py

@@ -0,0 +1,299 @@
+import os
+import shlex
+import shutil
+import sys
+from contextlib import contextmanager
+from pathlib import Path
+from typing import (
+    IO,
+    Any,
+    Optional,
+    Tuple,
+    Union,
+)
+
+if sys.version_info >= (3, 9):
+    from collections.abc import (
+        Callable,
+        Collection,
+        Generator,
+        Mapping,
+        Sequence,
+        Set,
+    )
+else:
+    from typing import AbstractSet as Set
+    from typing import (
+        Callable,
+        Collection,
+        Generator,
+        Mapping,
+        Sequence,
+    )
+
+from ._base import FoamCaseBase
+from ._subprocess import DEVNULL, STDOUT
+
+
+class _FoamCaseRecipes(FoamCaseBase):
+    def _clean_paths(self) -> Set[Path]:
+        has_decompose_par_dict = (self.path / "system" / "decomposeParDict").is_file()
+        has_block_mesh_dict = (self.path / "system" / "blockMeshDict").is_file()
+
+        paths = set()
+
+        for p in self.path.iterdir():
+            if p.is_dir():
+                try:
+                    t = float(p.name)
+                except ValueError:
+                    pass
+                else:
+                    if t != 0:
+                        paths.add(p)
+
+                if has_decompose_par_dict and p.name.startswith("processor"):
+                    paths.add(p)
+
+        if (self.path / "0.orig").is_dir() and (self.path / "0").is_dir():
+            paths.add(self.path / "0")
+
+        if has_block_mesh_dict and (self.path / "constant" / "polyMesh").exists():
+            paths.add(self.path / "constant" / "polyMesh")
+
+        if self._run_script() is not None:
+            paths.update(self.path.glob("log.*"))
+
+        return paths
+
+    def __delitem__(self, key: Union[int, float, str]) -> None:
+        shutil.rmtree(self[key].path)
+
+    def _clone_ignore(
+        self,
+    ) -> Callable[[Union["os.PathLike[str]", str], Collection[str]], Collection[str]]:
+        clean_paths = self._clean_paths()
+
+        def ignore(
+            path: Union["os.PathLike[str]", str], names: Collection[str]
+        ) -> Collection[str]:
+            paths = {Path(path) / name for name in names}
+            return {p.name for p in paths.intersection(clean_paths)}
+
+        return ignore
+
+    def _clean_script(self) -> Optional[Path]:
+        """Return the path to the (All)clean script, or None if no clean script is found."""
+        clean = self.path / "clean"
+        all_clean = self.path / "Allclean"
+
+        if clean.is_file():
+            script = clean
+        elif all_clean.is_file():
+            script = all_clean
+        else:
+            return None
+
+        if sys.argv and Path(sys.argv[0]).absolute() == script.absolute():
+            return None
+
+        return script if Path(sys.argv[0]).absolute() != script.absolute() else None
+
+    def _run_script(self, *, parallel: Optional[bool] = None) -> Optional[Path]:
+        """Return the path to the (All)run script, or None if no run script is found."""
+        run = self.path / "run"
+        run_parallel = self.path / "run-parallel"
+        all_run = self.path / "Allrun"
+        all_run_parallel = self.path / "Allrun-parallel"
+
+        if run.is_file() or all_run.is_file():
+            if run_parallel.is_file() or all_run_parallel.is_file():
+                if parallel:
+                    script = (
+                        run_parallel if run_parallel.is_file() else all_run_parallel
+                    )
+                elif parallel is False:
+                    script = run if run.is_file() else all_run
+                else:
+                    raise ValueError(
+                        "Both (All)run and (All)run-parallel scripts are present. Please specify parallel argument."
+                    )
+            else:
+                script = run if run.is_file() else all_run
+        elif parallel is not False and (
+            run_parallel.is_file() or all_run_parallel.is_file()
+        ):
+            script = run_parallel if run_parallel.is_file() else all_run_parallel
+        else:
+            return None
+
+        if sys.argv and Path(sys.argv[0]).absolute() == script.absolute():
+            return None
+
+        return script
+
+    def _env(self, *, shell: bool) -> Optional[Mapping[str, str]]:
+        sip_workaround = os.environ.get(
+            "FOAM_LD_LIBRARY_PATH", ""
+        ) and not os.environ.get("DYLD_LIBRARY_PATH", "")
+
+        if not shell or sip_workaround:
+            env = os.environ.copy()
+
+            if not shell:
+                env["PWD"] = str(self.path)
+
+            if sip_workaround:
+                env["DYLD_LIBRARY_PATH"] = env["FOAM_LD_LIBRARY_PATH"]
+
+            return env
+        else:
+            return None
+
+    @contextmanager
+    def _output(
+        self, cmd: Union[Sequence[Union[str, "os.PathLike[str]"]], str], *, log: bool
+    ) -> Generator[Tuple[Union[int, IO[bytes]], Union[int, IO[bytes]]], None, None]:
+        if log:
+            if isinstance(cmd, str):
+                name = shlex.split(cmd)[0]
+            else:
+                if isinstance(cmd[0], os.PathLike):
+                    name = Path(cmd[0]).name
+                else:
+                    name = cmd[0]
+
+            with (self.path / f"log.{name}").open("ab") as stdout:
+                yield stdout, STDOUT
+        else:
+            yield DEVNULL, DEVNULL
+
+    def _copy_cmds(
+        self, dest: Union["os.PathLike[str]", str]
+    ) -> Generator[Tuple[str, Sequence[Any], Mapping[str, Any]], None, None]:
+        yield (
+            "_copytree",
+            (
+                self.path,
+                dest,
+            ),
+            {"symlinks": True},
+        )
+
+    def _clean_cmds(
+        self, *, script: bool = True, check: bool = False
+    ) -> Generator[Tuple[str, Sequence[Any], Mapping[str, Any]], None, None]:
+        script_path = self._clean_script() if script else None
+
+        if script_path is not None:
+            yield ("_run", ([script_path],), {"cpus": 0, "check": check, "log": False})
+        else:
+            for p in self._clean_paths():
+                if p.is_dir():
+                    yield ("_rmtree", (p,), {})
+                else:
+                    p.unlink()
+
+    def _clone_cmds(
+        self, dest: Union["os.PathLike[str]", str]
+    ) -> Generator[Tuple[str, Sequence[Any], Mapping[str, Any]], None, None]:
+        if self._clean_script() is not None:
+            yield ("copy", (dest,), {})
+            yield ("clean", (), {})
+        else:
+            yield (
+                "_copytree",
+                (
+                    self.path,
+                    dest,
+                ),
+                {"symlinks": True, "ignore": self._clone_ignore()},
+            )
+
+    def _restore_0_dir_cmds(
+        self,
+    ) -> Generator[Tuple[str, Sequence[Any], Mapping[str, Any]], None, None]:
+        yield ("_rmtree", (self.path / "0",), {"ignore_errors": True})
+        yield (
+            "_copytree",
+            (
+                self.path / "0.orig",
+                self.path / "0",
+            ),
+            {"symlinks": True},
+        )
+
+    def _run_cmds(
+        self,
+        cmd: Optional[Union[Sequence[Union[str, "os.PathLike[str]"]], str]] = None,
+        *,
+        script: bool = True,
+        parallel: Optional[bool] = None,
+        cpus: Optional[int] = None,
+        check: bool = True,
+    ) -> Generator[Tuple[str, Sequence[Any], Mapping[str, Any]], None, None]:
+        if cmd is not None:
+            if parallel:
+                if cpus is None:
+                    cpus = max(self._nprocessors, 1)
+            else:
+                parallel = False
+                if cpus is None:
+                    cpus = 1
+
+            yield ("_run", (cmd,), {"parallel": parallel, "cpus": cpus, "check": check})
+
+        else:
+            script_path = self._run_script(parallel=parallel) if script else None
+
+            if script_path is not None:
+                if parallel or parallel is None:
+                    if cpus is None:
+                        if self._nprocessors > 0:
+                            cpus = self._nprocessors
+                        elif (self.path / "system" / "decomposeParDict").is_file():
+                            cpus = self._nsubdomains
+                        else:
+                            cpus = 1
+                else:
+                    if cpus is None:
+                        cpus = 1
+
+                yield (
+                    "_run",
+                    ([script_path],),
+                    {"parallel": False, "cpus": cpus, "check": check},
+                )
+
+            else:
+                if not self and (self.path / "0.orig").is_dir():
+                    yield ("restore_0_dir", (), {})
+
+                if (self.path / "system" / "blockMeshDict").is_file():
+                    yield ("block_mesh", (), {"check": check})
+
+                if parallel is None:
+                    parallel = (
+                        (cpus is not None and cpus > 1)
+                        or self._nprocessors > 0
+                        or (self.path / "system" / "decomposeParDict").is_file()
+                    )
+
+                if parallel:
+                    if (
+                        self._nprocessors == 0
+                        and (self.path / "system" / "decomposeParDict").is_file()
+                    ):
+                        yield ("decompose_par", (), {"check": check})
+
+                    if cpus is None:
+                        cpus = max(self._nprocessors, 1)
+                else:
+                    if cpus is None:
+                        cpus = 1
+
+                yield (
+                    "_run",
+                    ([self.application],),
+                    {"parallel": parallel, "cpus": cpus, "check": check},
+                )

foamlib/_cases/_subprocess.py

@@ -0,0 +1,86 @@
+import asyncio
+import os
+import subprocess
+import sys
+from typing import IO, Optional, Union
+
+if sys.version_info >= (3, 9):
+    from collections.abc import Mapping, Sequence
+else:
+    from typing import Mapping, Sequence
+
+CalledProcessError = subprocess.CalledProcessError
+CompletedProcess = subprocess.CompletedProcess
+
+DEVNULL = subprocess.DEVNULL
+PIPE = subprocess.PIPE
+STDOUT = subprocess.STDOUT
+
+
+def run_sync(
+    cmd: Union[Sequence[Union[str, "os.PathLike[str]"]], str],
+    *,
+    check: bool = True,
+    cwd: Optional["os.PathLike[str]"] = None,
+    env: Optional[Mapping[str, str]] = None,
+    stdout: Optional[Union[int, IO[bytes]]] = None,
+    stderr: Optional[Union[int, IO[bytes]]] = None,
+) -> "CompletedProcess[bytes]":
+    if not isinstance(cmd, str) and sys.version_info < (3, 8):
+        cmd = [str(arg) for arg in cmd]
+
+    return subprocess.run(
+        cmd,
+        cwd=cwd,
+        env=env,
+        stdout=stdout,
+        stderr=stderr,
+        shell=isinstance(cmd, str),
+        check=check,
+    )
+
+
+async def run_async(
+    cmd: Union[Sequence[Union[str, "os.PathLike[str]"]], str],
+    *,
+    check: bool = True,
+    cwd: Optional["os.PathLike[str]"] = None,
+    env: Optional[Mapping[str, str]] = None,
+    stdout: Optional[Union[int, IO[bytes]]] = None,
+    stderr: Optional[Union[int, IO[bytes]]] = None,
+) -> "CompletedProcess[bytes]":
+    if isinstance(cmd, str):
+        proc = await asyncio.create_subprocess_shell(
+            cmd,
+            cwd=cwd,
+            env=env,
+            stdout=stdout,
+            stderr=stderr,
+        )
+
+    else:
+        if sys.version_info < (3, 8):
+            cmd = [str(arg) for arg in cmd]
+        proc = await asyncio.create_subprocess_exec(
+            *cmd,
+            cwd=cwd,
+            env=env,
+            stdout=stdout,
+            stderr=stderr,
+        )
+
+    output, error = await proc.communicate()
+
+    assert proc.returncode is not None
+
+    if check and proc.returncode != 0:
+        raise CalledProcessError(
+            returncode=proc.returncode,
+            cmd=cmd,
+            output=output,
+            stderr=error,
+        )
+
+    return CompletedProcess(
+        cmd, returncode=proc.returncode, stdout=output, stderr=error
+    )

foamlib/_cases/_sync.py

@@ -1,10 +1,13 @@
+import os
 import shutil
-import subprocess
 import sys
+import tempfile
 from pathlib import Path
+from types import TracebackType
 from typing import (
     Callable,
     Optional,
+    Type,
     Union,
 )
 
@@ -13,12 +16,16 @@ if sys.version_info >= (3, 9):
 else:
     from typing import Collection, Sequence
 
-from .._util import is_sequence
-from ._base import FoamCaseBase
-from ._util import check_returncode
+if sys.version_info >= (3, 11):
+    from typing import Self
+else:
+    from typing_extensions import Self
+
+from ._recipes import _FoamCaseRecipes
+from ._subprocess import run_sync
 
 
-class FoamCase(FoamCaseBase):
+class FoamCase(_FoamCaseRecipes):
     """
     An OpenFOAM case.
 
@@ -29,22 +36,51 @@ class FoamCase(FoamCaseBase):
     :param path: The path to the case directory.
     """
 
+    def __init__(self, path: Union["os.PathLike[str]", str] = Path()):
+        super().__init__(path)
+        self._tmp: Optional[bool] = None
+
     @staticmethod
-    def _rmtree(path: Path, *, ignore_errors: bool = False) -> None:
+    def _rmtree(
+        path: Union["os.PathLike[str]", str], *, ignore_errors: bool = False
+    ) -> None:
         shutil.rmtree(path, ignore_errors=ignore_errors)
 
     @staticmethod
     def _copytree(
-        src: Path,
-        dest: Path,
+        src: Union["os.PathLike[str]", str],
+        dest: Union["os.PathLike[str]", str],
         *,
         symlinks: bool = False,
         ignore: Optional[
-            Callable[[Union[Path, str], Collection[str]], Collection[str]]
+            Callable[[Union["os.PathLike[str]", str], Collection[str]], Collection[str]]
         ] = None,
     ) -> None:
         shutil.copytree(src, dest, symlinks=symlinks, ignore=ignore)
 
+    def __enter__(self) -> "FoamCase":
+        if self._tmp is None:
+            raise RuntimeError(
+                "Cannot use a non-copied/cloned case as a context manager"
+            )
+        return self
+
+    def __exit__(
+        self,
+        exc_type: Optional[Type[BaseException]],
+        exc_val: Optional[BaseException],
+        exc_tb: Optional[TracebackType],
+    ) -> None:
+        if self._tmp is not None:
+            if self._tmp:
+                self._rmtree(self.path.parent)
+            else:
+                self._rmtree(self.path)
+        else:
+            raise RuntimeError(
+                "Cannot use a non-copied/cloned case as a context manager"
+            )
+
     def clean(
         self,
         *,
@@ -62,43 +98,39 @@ class FoamCase(FoamCaseBase):
 
     def _run(
         self,
-        cmd: Union[Sequence[Union[str, Path]], str, Path],
+        cmd: Union[Sequence[Union[str, "os.PathLike[str]"]], str],
         *,
         parallel: bool = False,
         cpus: int = 1,
         check: bool = True,
+        log: bool = True,
     ) -> None:
-        shell = not is_sequence(cmd)
-
-        if parallel:
-            if shell:
-                cmd = f"mpiexec -np {cpus} {cmd} -parallel"
-            else:
-                assert is_sequence(cmd)
-                cmd = ["mpiexec", "-np", str(cpus), *cmd, "-parallel"]
-
-        if sys.version_info < (3, 8):
-            if shell:
-                cmd = str(cmd)
-            else:
-                cmd = (str(arg) for arg in cmd)
-
-        proc = subprocess.run(
-            cmd,
-            cwd=self.path,
-            env=self._env(shell=shell),
-            stdout=subprocess.DEVNULL,
-            stderr=subprocess.PIPE if check else subprocess.DEVNULL,
-            text=True,
-            shell=shell,
-        )
-
-        if check:
-            check_returncode(proc.returncode, cmd, proc.stderr)
+        with self._output(cmd, log=log) as (stdout, stderr):
+            if parallel:
+                if isinstance(cmd, str):
+                    cmd = [
+                        "mpiexec",
+                        "-n",
+                        str(cpus),
+                        "/bin/sh",
+                        "-c",
+                        f"{cmd} -parallel",
+                    ]
+                else:
+                    cmd = ["mpiexec", "-n", str(cpus), *cmd, "-parallel"]
+
+            run_sync(
+                cmd,
+                check=check,
+                cwd=self.path,
+                env=self._env(shell=isinstance(cmd, str)),
+                stdout=stdout,
+                stderr=stderr,
+            )
 
     def run(
         self,
-        cmd: Optional[Union[Sequence[Union[str, Path]], str, Path]] = None,
+        cmd: Optional[Union[Sequence[Union[str, "os.PathLike[str]"]], str]] = None,
         *,
         script: bool = True,
         parallel: Optional[bool] = None,
@@ -134,24 +166,46 @@ class FoamCase(FoamCaseBase):
         for name, args, kwargs in self._restore_0_dir_cmds():
             getattr(self, name)(*args, **kwargs)
 
-    def copy(self, dest: Union[Path, str]) -> "FoamCase":
+    def copy(self, dst: Optional[Union["os.PathLike[str]", str]] = None) -> "Self":
         """
         Make a copy of this case.
 
-        :param dest: The destination path.
+        Use as a context manager to automatically delete the copy when done.
+
+        :param dst: The destination path. If None, copy to a temporary directory.
         """
-        for name, args, kwargs in self._copy_cmds(dest):
+        if dst is None:
+            dst = Path(tempfile.mkdtemp(), self.name)
+            tmp = True
+        else:
+            tmp = False
+
+        for name, args, kwargs in self._copy_cmds(dst):
             getattr(self, name)(*args, **kwargs)
 
-        return FoamCase(dest)
+        ret = type(self)(dst)
+        ret._tmp = tmp
+
+        return ret
 
-    def clone(self, dest: Union[Path, str]) -> "FoamCase":
+    def clone(self, dst: Optional[Union["os.PathLike[str]", str]] = None) -> "Self":
         """
         Clone this case (make a clean copy).
 
-        :param dest: The destination path.
+        Use as a context manager to automatically delete the clone when done.
+
+        :param dst: The destination path. If None, clone to a temporary directory.
         """
-        for name, args, kwargs in self._clone_cmds(dest):
+        if dst is None:
+            dst = Path(tempfile.mkdtemp(), self.name)
+            tmp = True
+        else:
+            tmp = False
+
+        for name, args, kwargs in self._clone_cmds(dst):
             getattr(self, name)(*args, **kwargs)
 
-        return FoamCase(dest)
+        ret = type(self)(dst)
+        ret._tmp = tmp
+
+        return ret

foamlib/_cases/_util.py

@@ -1,35 +1,49 @@
-import subprocess
+import functools
 import sys
-from pathlib import Path
-from typing import Optional, Union
-from warnings import warn
+from types import TracebackType
+from typing import (
+    Any,
+    AsyncContextManager,
+    Callable,
+    Generic,
+    Optional,
+    Type,
+    TypeVar,
+)
 
 if sys.version_info >= (3, 9):
-    from collections.abc import Sequence
+    from collections.abc import Generator
 else:
-    from typing import Sequence
+    from typing import Generator
 
 
-class CalledProcessError(subprocess.CalledProcessError):
-    """Exception raised when a process fails and `check=True`."""
+R = TypeVar("R")
 
-    def __str__(self) -> str:
-        msg = super().__str__()
-        if self.stderr:
-            msg += f"\n{self.stderr}"
-        return msg
 
+class _AwaitableAsyncContextManager(Generic[R]):
+    def __init__(self, cm: "AsyncContextManager[R]"):
+        self._cm = cm
 
-class CalledProcessWarning(Warning):
-    """Warning raised when a process prints to stderr and `check=True`."""
+    def __await__(self) -> Generator[Any, Any, R]:
+        return self._cm.__aenter__().__await__()
 
+    async def __aenter__(self) -> R:
+        return await self._cm.__aenter__()
 
-def check_returncode(
-    retcode: int,
-    cmd: Union[Sequence[Union[str, Path]], str, Path],
-    stderr: Optional[str],
-) -> None:
-    if retcode != 0:
-        raise CalledProcessError(retcode, cmd, None, stderr)
-    elif stderr:
-        warn(f"Command {cmd} printed to stderr.\n{stderr}", CalledProcessWarning)
+    async def __aexit__(
+        self,
+        exc_type: Optional[Type[BaseException]],
+        exc_val: Optional[BaseException],
+        exc_tb: Optional[TracebackType],
+    ) -> Optional[bool]:
+        return await self._cm.__aexit__(exc_type, exc_val, exc_tb)
+
+
+def awaitableasynccontextmanager(
+    cm: Callable[..., "AsyncContextManager[R]"],
+) -> Callable[..., _AwaitableAsyncContextManager[R]]:
+    @functools.wraps(cm)
+    def f(*args: Any, **kwargs: Any) -> _AwaitableAsyncContextManager[R]:
+        return _AwaitableAsyncContextManager(cm(*args, **kwargs))
+
+    return f