diff --git a/fsl/utils/filetree/filetree.py b/fsl/utils/filetree/filetree.py index 9ba5fede6ad1d7f4ee6a9b9a24980941b6d7fc0a..0b97268f2dd82580550e1c1d8ce3fda5498b6cfd 100644 --- a/fsl/utils/filetree/filetree.py +++ b/fsl/utils/filetree/filetree.py @@ -185,7 +185,7 @@ class FileTree(object): """ return tuple(self.extract_variables(short_name, fn) for fn in self.get_all(short_name, glob_vars=glob_vars)) - def get_all_trees(self, short_name: str, glob_vars=()) -> Tuple["FileTree"]: + def get_all_trees(self, short_name: str, glob_vars=(), set_parent=True) -> Tuple["FileTree"]: """ Gets all the trees that generate the existing files matching the pattern @@ -195,23 +195,31 @@ class FileTree(object): :param glob_vars: sequence of undefined variables that can take any possible values when looking for matches on the disk. Any defined variables in `glob_vars` will be ignored. If glob_vars is set to 'all', all undefined variables will be used to look up matches. + :param set_parent: Update the variables of the top-level rather than current tree if True. + Ony relevant if `self` is a sub-tree. :return: sequence of FileTrees used to generate each file on disk matching the pattern of `short_name` """ - return tuple(self.update(**vars) for vars in self.get_all_vars(short_name, glob_vars=glob_vars)) + return tuple(self.update(set_parent=set_parent, **vars) + for vars in self.get_all_vars(short_name, glob_vars=glob_vars)) - def update(self, **variables) -> "FileTree": + def update(self, set_parent=True, **variables) -> "FileTree": """ Creates a new FileTree with updated variables + :param set_parent: Update the variables of the top-level rather than current tree if True. + Ony relevant if `self` is a sub-tree. :param variables: new values for the variables Setting a variable to None will cause the variable to be unset :return: New FileTree with same templates for directory names and filenames, but updated variables """ new_tree = deepcopy(self) - new_tree.variables.update(variables) + set_tree = new_tree + while set_parent and set_tree.parent is not None: + set_tree = set_tree.parent + set_tree.variables.update(variables) for key, value in variables.items(): if value is None: - del new_tree.variables[key] + del set_tree.variables[key] return new_tree def extract_variables(self, short_name: str, filename: str) -> Dict[str, str]: diff --git a/tests/test_filetree/test_read.py b/tests/test_filetree/test_read.py index 954f671e791ca8edf3d3ef0a8932488e3669f7ab..cba2625e19601722301e5b3b68f26fa472152637 100644 --- a/tests/test_filetree/test_read.py +++ b/tests/test_filetree/test_read.py @@ -58,6 +58,17 @@ def test_parent_tree(): with pytest.raises(KeyError): tree.get('sub_var/basename') + # test updating in parent tree + sub0_tree = tree.sub_trees['sub0'] + same_path(sub0_tree.update(subvar='test').get('../subvar/basename'), 'subvar_test') + with pytest.raises(KeyError): + sub0_tree.update(subvar='test', set_parent=False).get('../subvar/basename') + + sub0_tree = tree.update(subvar='grot').sub_trees['sub0'] + same_path(sub0_tree.update(subvar='test').get('../subvar/basename'), 'subvar_test') + same_path(sub0_tree.update(subvar='test', set_parent=False).get('../subvar/basename'), 'subvar_grot') + same_path(sub0_tree.get('../subvar/basename'), 'subvar_grot') + def test_custom_tree(): directory = op.split(__file__)[0]