Source code for dpgen2.utils.chdir

import os
from contextlib import (
    contextmanager,
)
from functools import (
    wraps,
)
from pathlib import (
    Path,
)
from typing import (
    Callable,
)

from dflow.python import (
    OPIO,
)


[docs] @contextmanager def set_directory(path: Path): """Sets the current working path within the context. Parameters ---------- path : Path The path to the cwd Yields ------ None Examples -------- >>> with set_directory("some_path"): ... do_something() """ cwd = Path().absolute() path.mkdir(exist_ok=True, parents=True) try: os.chdir(path) yield finally: os.chdir(cwd)
[docs] def chdir(path_key: str): """Returns a decorator that can change the current working path. Parameters ---------- path_key : str key to OPIO Examples -------- >>> class SomeOP(OP): ... @chdir("path") ... def execute(self, ip: OPIO): ... do_something() """ def decorator(func: Callable): """Change the current working path.""" @wraps(func) def wrapper(self, ip: OPIO): with set_directory(Path(ip[path_key])): return func(self, ip) return wrapper return decorator