diff --git a/luigi_tools/parameter.py b/luigi_tools/parameter.py index 9d1a23a..a928929 100644 --- a/luigi_tools/parameter.py +++ b/luigi_tools/parameter.py @@ -14,6 +14,7 @@ """This module provides some specific luigi parameters.""" import warnings +from pathlib import Path import luigi @@ -83,6 +84,34 @@ def __init__(self, *args, **kwargs): self.parsing = self.__class__.EXPLICIT_PARSING +class PathParameter(luigi.Parameter): + """Class to parse file path parameters. + + Args: + absolute (bool): the given path is converted to an absolute path. + create (bool): a folder is automatically created to the given path. + exists (bool): raise a :class:`ValueError` if the path does not exist. + """ + + def __init__(self, *args, absolute=False, create=False, exists=False, **kwargs): + super().__init__(*args, **kwargs) + + self.absolute = absolute + self.create = create + self.exists = exists + + def normalize(self, x): + """Normalize the given value to a :class:`pathlib.Path` object.""" + path = Path(x) + if self.absolute: + path = path.absolute() + if self.create: + path.mkdir(parents=True, exist_ok=True) + if self.exists and not path.exists(): + raise ValueError(f"The path {path} does not exist.") + return path + + class OptionalParameter: """Mixin to make a parameter class optional.""" @@ -179,3 +208,9 @@ class OptionalTupleParameter(OptionalParameter, luigi.TupleParameter): """Class to parse optional tuple parameters.""" expected_type = tuple + + +class OptionalPathParameter(OptionalParameter, PathParameter): + """Class to parse optional path parameters.""" + + expected_type = str diff --git a/tests/test_parameter.py b/tests/test_parameter.py index cd86ecb..b216fa1 100644 --- a/tests/test_parameter.py +++ b/tests/test_parameter.py @@ -541,3 +541,70 @@ class TaskBoolParameterFail(luigi.Task): """""" a = luigi_tools.parameter.BoolParameter(default=True, parsing="implicit") + + +@pytest.mark.parametrize("default", [None, "not_existing_dir"]) +@pytest.mark.parametrize("absolute", [True, False]) +@pytest.mark.parametrize("create", [True, False]) +@pytest.mark.parametrize("exists", [True, False]) +def test_path_parameter(tmpdir, default, absolute, create, exists): + class TaskPathParameter(luigi.Task): + + a = luigi_tools.parameter.PathParameter( + default=str(tmpdir / default) if default is not None else str(tmpdir), + absolute=absolute, + create=create, + exists=exists, + ) + b = luigi_tools.parameter.OptionalPathParameter( + default=str(tmpdir / default) if default is not None else str(tmpdir), + absolute=absolute, + create=create, + exists=exists, + ) + c = luigi_tools.parameter.OptionalPathParameter(default=None) + d = luigi_tools.parameter.OptionalPathParameter(default="not empty default") + + def run(self): + # Use the parameter as a Path object + new_file = self.a / "test.file" + new_optional_file = self.b / "test_optional.file" + if default is not None and not create: + new_file.parent.mkdir(parents=True) + new_file.touch() + new_optional_file.touch() + assert new_file.exists() + assert new_optional_file.exists() + assert self.c is None + assert self.d is None + + def output(self): + return luigi.LocalTarget("not_existing_file") + + # Test with default values + with set_luigi_config({"TaskPathParameter": {"d": ""}}): + if default is not None and not create and exists: + with pytest.raises(ValueError, match="The path .* does not exist"): + luigi.build([TaskPathParameter()], local_scheduler=True) + else: + assert luigi.build([TaskPathParameter()], local_scheduler=True) + + # Test with values from config + with set_luigi_config( + { + "TaskPathParameter": { + "a": str(tmpdir / (default + "_from_config")) + if default is not None + else str(tmpdir), + "b": str(tmpdir / (default + "_from_config")) + if default is not None + else str(tmpdir), + "d": "", + } + } + ): + if default is not None and not create and exists: + with pytest.raises(ValueError, match="The path .* does not exist"): + luigi.build([TaskPathParameter()], local_scheduler=True) + else: + assert luigi.build([TaskPathParameter()], local_scheduler=True)