From 3c678cf9e7be6a3f811140bc4a5ee8132cd24c20 Mon Sep 17 00:00:00 2001
From: Michiel Cottaar <MichielCottaar@gmail.com>
Date: Thu, 14 Feb 2019 10:13:20 +0000
Subject: [PATCH] ENH: added get_all_trees to simplify iterating over variables

---
 fsl/utils/filetree/__init__.py   |  6 +++---
 fsl/utils/filetree/filetree.py   | 14 ++++++++++++++
 tests/test_filetree/test_read.py |  4 ++++
 3 files changed, 21 insertions(+), 3 deletions(-)

diff --git a/fsl/utils/filetree/__init__.py b/fsl/utils/filetree/__init__.py
index 73589296a..9d31383db 100644
--- a/fsl/utils/filetree/__init__.py
+++ b/fsl/utils/filetree/__init__.py
@@ -193,9 +193,9 @@ Assuming that the input T1w's already exist, we can then simply run BET for ever
     from fsl.utils.filetree import FileTree
     from fsl.wrappers.bet import bet
     tree = FileTree.read(<tree filename>)
-    variables = tree.get_all_vars('T1w')  # extract the set of variables for all existing T1w files
-    for single_variable_set in variables:
-        T1w_tree = tree.update(**single_variable_set)
+
+    # Iterates over set of variables that correspond to each T1-weighted image file matching the template
+    for T1w_tree in tree.get_all_trees('T1w', glob_vars='all'):
         # get retrieves the filenames based on the current set of variables
         # make_dir=True ensures that the output directory containing the "bet_output" actually exists
         bet(input=T1w_tree.get('T1w'), output=T1w_tree.get('bet_output', make_dir=True), mask=True)
diff --git a/fsl/utils/filetree/filetree.py b/fsl/utils/filetree/filetree.py
index 812cdf8c0..fb7aa082f 100644
--- a/fsl/utils/filetree/filetree.py
+++ b/fsl/utils/filetree/filetree.py
@@ -185,6 +185,20 @@ 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, global_vars=()) -> Tuple["FileTree"]:
+        """
+        Gets all the trees that generate the existing files matching the pattern
+
+        tree.get_all(short_name) == tuple(tree.get(short_name) for tree in tree.get_all_trees(short_name))
+
+        :param short_name: short name of the path template
+        :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.
+        :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=global_vars))
+
     def update(self, **variables) -> "FileTree":
         """
         Creates a new FileTree with updated variables
diff --git a/tests/test_filetree/test_read.py b/tests/test_filetree/test_read.py
index 40ffde224..954f671e7 100644
--- a/tests/test_filetree/test_read.py
+++ b/tests/test_filetree/test_read.py
@@ -90,6 +90,10 @@ def test_custom_tree():
         tree.get_all('opt_file')
     assert len(tree.get_all('opt_file', glob_vars=['opt'])) == 1
 
+    for short_name in ('sub_file', 'opt_file'):
+        for glob_vars in (['opt'], 'all'):
+            assert tree.get_all(short_name, glob_vars) == tuple(stree.get(short_name) for stree in tree.get_all_trees(short_name, glob_vars))
+
     for vars in ({'opt': None}, {'opt': 'test'}):
         filename = tree.update(**vars).get('sub_file')
         assert vars == tree.extract_variables('sub_file', filename)
-- 
GitLab