diff --git a/dfttoolkit/utils/run_utils.py b/dfttoolkit/utils/run_utils.py index cbe52a8..9400aca 100644 --- a/dfttoolkit/utils/run_utils.py +++ b/dfttoolkit/utils/run_utils.py @@ -2,12 +2,20 @@ from functools import wraps -def no_repeat(calc_dir="./", force=False): +def no_repeat( + original_func=None, + *, + output_file: str = "aims.out", + calc_dir: str = "./", + force: bool = False, +): """ Don't repeat the calculation if aims.out exists in the calculation directory. Parameters ---------- + output_file : str, default='aims.out' + The name of the output file to check for. calc_dir : str, default="./" The directory where the calculation is performed force : bool, default=False @@ -27,11 +35,14 @@ def wrapper(*args, **kwargs): raise ValueError(f"{calc_dir} is not a directory.") if force: return func(*args, **kwargs) - if not os.path.isfile(f"{calc_dir}/aims.out"): + if not os.path.isfile(f"{calc_dir}/{output_file}"): return func(*args, **kwargs) else: print(f"aims.out already exists in {calc_dir}. Skipping calculation.") return wrapper + if original_func: + return _no_repeat(original_func) + return _no_repeat diff --git a/tests/test_output.py b/tests/test_output.py index cd73927..0f5c9f6 100644 --- a/tests/test_output.py +++ b/tests/test_output.py @@ -15,9 +15,9 @@ def aims_out(self, request, run_aims): cwd = os.path.dirname(os.path.realpath(__file__)) - # Run the FHI-aims calculations if the run_aims option is specified but not if + # Run the FHI-aims calculations if the run-aims option is specified but not if # they already exist. - # Force it to run if the run_aims option is "change_bin" + # Force it to run if the run-aims option is "change_bin" # run_aims fixture is defined in conftest.py if request.param == 1 and run_aims is not False: binary = aims_bin_path_prompt(run_aims, cwd) diff --git a/tests/test_utils.py b/tests/test_utils.py index b86b7fb..8d57230 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -1,9 +1,52 @@ import pytest -from dfttools.utils import run_utils +import os +from dfttoolkit.utils.run_utils import no_repeat -class TestRunUtils: +class TestNoRepeat: - def test_no_repeat(self): - # TODO - pass + calc_dir = ( + f"{os.path.dirname(os.path.realpath(__file__))}/fixtures/default_aims_calcs/1" + ) + + def test_specified_dir_not_found(self): + @no_repeat(calc_dir="bogus") + def to_be_decorated(): + return True + + with pytest.raises(ValueError) as excinfo: + to_be_decorated() + + assert "bogus is not a directory." == str(excinfo.value) + + @pytest.mark.parametrize( + ("calc_dir", "force", "expected"), + [ + ("./", False, True), + ( + f"{os.path.dirname(os.path.realpath(__file__))}/fixtures/default_aims_calcs/1", + True, + True, + ), + ], + ) + def test_default_dir_no_args(self, calc_dir, force, expected): + @no_repeat(calc_dir=calc_dir, force=force) + def to_be_decorated(): + return True + + assert to_be_decorated() == expected + + def test_no_repeat(self, capfd): + @no_repeat(calc_dir=self.calc_dir) + def to_be_decorated(): + return True + + to_be_decorated() + + out, err = capfd.readouterr() + assert ( + out + == f"aims.out already exists in {self.calc_dir}. Skipping calculation.\n" + ) + assert err == ""