copyfile uses get_related_files, tests for get_related_files

Author: ashgillman

nipype/utils/filemanip.py

@@ -207,7 +207,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
@@ -228,6 +229,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
     -------
@@ -320,30 +324,36 @@ def copyfile(originalfile, newfile, copy=False, create_new=False,
             fmlogger.warn(e.message)
 
     # Associated files
-    _, _, this_type = split_filename(originalfile)
-    for type_set in related_filetype_sets:
-        if this_type in type_set:
-            for alt_type in type_set:
-                alt_ofile = originalfile[:-len(this_type)] + alt_type
-                alt_nfile = newfile[:-len(this_type)] + alt_type
-                if os.path.exists(alt_ofile) and not os.path.exists(alt_nfile):
-                    copyfile(alt_ofile, alt_nfile, 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 = []
-    path, name, ext = split_filename(filename)
+    path, name, this_type = split_filename(filename)
     for type_set in related_filetype_sets:
-        if ext in type_set:
-            for new_ext in type_set:
-                related_files.append(os.path.join(path, name + new_ext))
+        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,13 +8,14 @@
 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,
-                                    copyfile, copyfiles,
-                                    filename_to_list, list_to_filename,
-                                    split_filename, get_related_files)
+                                fname_presuffix, fnames_presuffix,
+                                hash_rename, check_forhash,
+                                copyfile, copyfiles,
+                                filename_to_list, list_to_filename,
+                                split_filename, get_related_files)
 
 import numpy as np
 
@@ -249,6 +250,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']