Merge pull request #1617 from ashgillman/REF-related-files

REF: define related filetypes once.

Author: Shoshana Berleant

nipype/utils/filemanip.py

@@ -29,6 +29,12 @@
 fmlogger = logging.getLogger("filemanip")
 
 
+related_filetype_sets = [
+    ('.hdr', '.img', '.mat'),
+    ('.BRIK', '.HEAD'),
+]
+
+
 class FileNotFoundError(Exception):
     pass
 
@@ -215,7 +221,8 @@ def hash_timestamp(afile):
 
 
 def copyfile(originalfile, newfile, copy=False, create_new=False,
-             hashmethod=None, use_hardlink=False):
+             hashmethod=None, use_hardlink=False,
+             copy_related_files=True):
     """Copy or link ``originalfile`` to ``newfile``.
 
     If ``use_hardlink`` is True, and the file can be hard-linked, then a
@@ -236,6 +243,9 @@ def copyfile(originalfile, newfile, copy=False, create_new=False,
     use_hardlink : Bool
         specifies whether to hard-link files, when able
         (Default=False), taking precedence over copy
+    copy_related_files : Bool
+        specifies whether to also operate on related files, as defined in
+        ``related_filetype_sets``
 
     Returns
     -------
@@ -328,38 +338,36 @@ def copyfile(originalfile, newfile, copy=False, create_new=False,
             fmlogger.warn(e.message)
 
     # Associated files
-    if originalfile.endswith(".img"):
-        hdrofile = originalfile[:-4] + ".hdr"
-        hdrnfile = newfile[:-4] + ".hdr"
-        matofile = originalfile[:-4] + ".mat"
-        if os.path.exists(matofile):
-            matnfile = newfile[:-4] + ".mat"
-            copyfile(matofile, matnfile, copy, hashmethod=hashmethod,
-                     use_hardlink=use_hardlink)
-        copyfile(hdrofile, hdrnfile, copy, hashmethod=hashmethod,
-                 use_hardlink=use_hardlink)
-    elif originalfile.endswith(".BRIK"):
-        hdrofile = originalfile[:-5] + ".HEAD"
-        hdrnfile = newfile[:-5] + ".HEAD"
-        copyfile(hdrofile, hdrnfile, copy, hashmethod=hashmethod,
-                 use_hardlink=use_hardlink)
+    if copy_related_files:
+        related_file_pairs = (get_related_files(f, include_this_file=False)
+                              for f in (originalfile, newfile))
+        for alt_ofile, alt_nfile in zip(*related_file_pairs):
+            if os.path.exists(alt_ofile):
+                copyfile(alt_ofile, alt_nfile, copy, hashmethod=hashmethod,
+                         use_hardlink=use_hardlink, copy_related_files=False)
 
     return newfile
 
 
-def get_related_files(filename):
-    """Returns a list of related files for Nifti-Pair, Analyze (SPM) and AFNI
-       files
+def get_related_files(filename, include_this_file=True):
+    """Returns a list of related files, as defined in
+    ``related_filetype_sets``, for a filename. (e.g., Nifti-Pair, Analyze (SPM)
+    and AFNI files).
+
+    Parameters
+    ----------
+    filename : str
+        File name to find related filetypes of.
+    include_this_file : bool
+        If true, output includes the input filename.
     """
     related_files = []
-    if filename.endswith(".img") or filename.endswith(".hdr"):
-        path, name, ext = split_filename(filename)
-        for ext in ['.hdr', '.img', '.mat']:
-            related_files.append(os.path.join(path, name + ext))
-    elif filename.endswith(".BRIK") or filename.endswith(".HEAD"):
-        path, name, ext = split_filename(filename)
-        for ext in ['.BRIK', '.HEAD']:
-            related_files.append(os.path.join(path, name + ext))
+    path, name, this_type = split_filename(filename)
+    for type_set in related_filetype_sets:
+        if this_type in type_set:
+            for related_type in type_set:
+                if include_this_file or related_type != this_type:
+                    related_files.append(os.path.join(path, name + related_type))
     if not len(related_files):
         related_files = [filename]
     return related_files

nipype/utils/tests/test_filemanip.py

@@ -8,7 +8,8 @@
 from tempfile import mkstemp, mkdtemp
 import warnings
 
-from ...testing import assert_equal, assert_true, assert_false, TempFATFS
+from ...testing import (assert_equal, assert_true, assert_false,
+                        assert_in, assert_not_in, TempFATFS)
 from ...utils.filemanip import (save_json, load_json,
                                 fname_presuffix, fnames_presuffix,
                                 hash_rename, check_forhash,
@@ -259,6 +260,30 @@ def test_copyfallback():
         os.unlink(orig_hdr)
 
 
+def test_get_related_files():
+    orig_img, orig_hdr = _temp_analyze_files()
+
+    related_files = get_related_files(orig_img)
+    yield assert_in, orig_img, related_files
+    yield assert_in, orig_hdr, related_files
+
+    related_files = get_related_files(orig_hdr)
+    yield assert_in, orig_img, related_files
+    yield assert_in, orig_hdr, related_files
+
+
+def test_get_related_files_noninclusive():
+    orig_img, orig_hdr = _temp_analyze_files()
+
+    related_files = get_related_files(orig_img, include_this_file=False)
+    yield assert_not_in, orig_img, related_files
+    yield assert_in, orig_hdr, related_files
+
+    related_files = get_related_files(orig_hdr, include_this_file=False)
+    yield assert_in, orig_img, related_files
+    yield assert_not_in, orig_hdr, related_files
+
+
 def test_filename_to_list():
     x = filename_to_list('foo.nii')
     yield assert_equal, x, ['foo.nii']