Add recursive option to sort dicoms

shimmingtoolbox/cli/sort_dicoms.py

@@ -14,30 +14,25 @@
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
+DEFAULT_RECURSIVE_DEPTH = 5
+
 
 @click.command(context_settings=CONTEXT_SETTINGS)
 @click.option('-i', "--input", "path_input", help="Input folder.")
+@click.option('-r', "--recursive", 'is_recursive', is_flag=True, default=False, show_default=True,
+              help=f"Specifies to look into sub-folders. See also the --recursive-depth option.")
+@click.option('--recursive-depth', 'recursive_depth', type=click.INT, default=DEFAULT_RECURSIVE_DEPTH,
+              show_default=True, help="Depth of the recursive search.")
 @click.option('-o', "--output", "path_output", help="Output folder.")
 @click.option('-v', '--verbose', type=click.Choice(['info', 'debug']), default='info', help="Be more verbose")
-def sort_dicoms(path_input, path_output, verbose):
+def sort_dicoms(path_input, is_recursive, recursive_depth, path_output, verbose):
     set_all_loggers(verbose)
 
-    # Create a list containing all the DICOMs in the input folder
-    list_dicoms = []
-    for name in os.listdir(path_input):
-        fname_tmp = os.path.join(path_input, name)
-        # If it's not a file
-        if not os.path.isfile(fname_tmp):
-            continue
-        # If it's not a DICOM file
-        if not is_dicom(fname_tmp):
-            continue
-
-        list_dicoms.append(fname_tmp)
+    list_dicoms = get_dicom_paths(path_input, is_recursive, 0, max_depth=recursive_depth)
 
     # Make sure there is at least one DICOM
     if not list_dicoms:
-        raise RuntimeError(f"{path_input} does not contain dicom files")
+        raise RuntimeError(f"{path_input} does not contain dicom files, use the -r option to look into sub-folders.")
 
     # Create output directory
     create_output_dir(path_output)
@@ -60,3 +55,24 @@ def sort_dicoms(path_input, path_output, verbose):
         shutil.copyfile(fname_dcm, fname_output)
 
     logger.info("Successfully sorted the DICOMs")
+
+
+def get_dicom_paths(path, is_recursive, subfolder_depth=0, max_depth=DEFAULT_RECURSIVE_DEPTH):
+    # Create a list containing all the DICOMs in the input folder
+
+    list_dicoms = []
+    for name in os.listdir(path):
+        fname_tmp = os.path.join(path, name)
+        # If it's not a file
+        if not os.path.isfile(fname_tmp):
+            if os.path.isdir(fname_tmp):
+                if is_recursive and subfolder_depth < max_depth:
+                    list_dicoms += get_dicom_paths(fname_tmp, is_recursive, subfolder_depth + 1, max_depth=max_depth)
+            continue
+        # If it's not a DICOM file
+        if not is_dicom(fname_tmp):
+            continue
+
+        list_dicoms.append(fname_tmp)
+
+    return list_dicoms

test/cli/test_cli_sort_dicoms.py

@@ -2,9 +2,11 @@
 # coding: utf-8
 
 from click.testing import CliRunner
+import copy
 import os
 import pathlib
 import pytest
+import shutil
 import tempfile
 
 from shimmingtoolbox.cli.sort_dicoms import sort_dicoms
@@ -29,3 +31,21 @@ def test_sort_dicoms_cli_no_dicom():
     path = os.path.join(__dir_testing__)
     with pytest.raises(RuntimeError, match="does not contain dicom files"):
         CliRunner().invoke(sort_dicoms, ['-i', path], catch_exceptions=False)
+
+
+def test_sort_dicoms_cli_recursive():
+    path = os.path.join(__dir_testing__, 'dicom_unsorted')
+    with tempfile.TemporaryDirectory(prefix='st_' + pathlib.Path(__file__).stem) as tmp:
+        path_subfolder = copy.deepcopy(tmp)
+        for i_subfolder in range(1, 5):
+            path_subfolder = os.path.join(path_subfolder, 'subfolder' + str(i_subfolder))
+            os.mkdir(path_subfolder)
+
+        shutil.copytree(path, os.path.join(path_subfolder, 'dicoms'))
+
+        path_output = os.path.join(tmp, 'sorted')
+        result = CliRunner().invoke(sort_dicoms, ['-i', tmp, '-r', '-o', path_output], catch_exceptions=False)
+        assert result.exit_code == 0
+        outputs = os.listdir(path_output)
+        assert '06-a_gre_DYNshim' in outputs
+        assert '07-a_gre_DYNshim' in outputs