Finished adding local_copy logic and passed all unit tests

Author: pintohutch

nipype/interfaces/io.py

@@ -209,7 +209,7 @@ class DataSinkInputSpec(DynamicTraitedSpec, BaseInterfaceInputSpec):
     bucket = traits.Generic(mandatory=False,
                             desc='Boto3 S3 bucket for manual override of bucket')
     # Set this if user wishes to have local copy of files as well
-    local_dir = traits.Str(desc='Copy files locally as well as to S3 bucket')
+    local_copy = traits.Str(desc='Copy files locally as well as to S3 bucket')
 
     # Set call-able inputs attributes
     def __setattr__(self, key, value):
@@ -392,6 +392,10 @@ def _check_s3_base_dir(self):
         s3_str = 's3://'
         base_directory = self.inputs.base_directory
 
+        if not isdefined(base_directory):
+            s3_flag = False
+            return s3_flag
+
         # Explicitly lower-case the "s3"
         if base_directory.lower().startswith(s3_str):
             base_dir_sp = base_directory.split('/')
@@ -616,7 +620,7 @@ def _upload_to_s3(self, src, dst):
                 else:
                     iflogger.info('Overwriting previous S3 file...')
 
-            except ClientError as exc:
+            except ClientError:
                 iflogger.info('New file to S3')
 
             # Copy file up to S3 (either encrypted or not)
@@ -653,18 +657,21 @@ def _list_outputs(self):
         # Check if base directory reflects S3 bucket upload
         try:
             s3_flag = self._check_s3_base_dir()
-            s3dir = self.inputs.base_directory
-            if isdefined(self.inputs.container):
-                s3dir = os.path.join(s3dir, self.inputs.container)
+            if s3_flag:
+                s3dir = self.inputs.base_directory
+                if isdefined(self.inputs.container):
+                    s3dir = os.path.join(s3dir, self.inputs.container)
+            else:
+                s3dir = '<N/A>'
         # If encountering an exception during bucket access, set output
         # base directory to a local folder
         except Exception as exc:
+            s3dir = '<N/A>'
+            s3_flag = False
             if not isdefined(self.inputs.local_copy):
                 local_out_exception = os.path.join(os.path.expanduser('~'),
                                                    's3_datasink_' + self.bucket.name)
                 outdir = local_out_exception
-            else:
-                outdir = self.inputs.local_copy
             # Log local copying directory
             iflogger.info('Access to S3 failed! Storing outputs locally at: '\
                           '%s\nError: %s' %(outdir, exc))
@@ -673,8 +680,8 @@ def _list_outputs(self):
         if isdefined(self.inputs.container):
             outdir = os.path.join(outdir, self.inputs.container)
 
-        # If doing a localy output
-        if not outdir.lower().startswith('s3://'):
+        # If sinking to local folder
+        if outdir != s3dir:
             outdir = os.path.abspath(outdir)
             # Create the directory if it doesn't exist
             if not os.path.exists(outdir):
@@ -714,18 +721,19 @@ def _list_outputs(self):
                 if not os.path.isfile(src):
                     src = os.path.join(src, '')
                 dst = self._get_dst(src)
+                if s3_flag:
+                    s3dst = os.path.join(s3tempoutdir, dst)
+                    s3dst = self._substitute(s3dst)
                 dst = os.path.join(tempoutdir, dst)
-                s3dst = os.path.join(s3tempoutdir, dst)
                 dst = self._substitute(dst)
                 path, _ = os.path.split(dst)
 
                 # If we're uploading to S3
                 if s3_flag:
-                    dst = dst.replace(outdir, self.inputs.base_directory)
-                    self._upload_to_s3(src, dst)
-                    out_files.append(dst)
+                    self._upload_to_s3(src, s3dst)
+                    out_files.append(s3dst)
                 # Otherwise, copy locally src -> dst
-                else:
+                if not s3_flag or isdefined(self.inputs.local_copy):
                     # Create output directory if it doesnt exist
                     if not os.path.exists(path):
                         try:
@@ -787,6 +795,8 @@ class S3DataSinkInputSpec(DynamicTraitedSpec, BaseInterfaceInputSpec):
     _outputs = traits.Dict(traits.Str, value={}, usedefault=True)
     remove_dest_dir = traits.Bool(False, usedefault=True,
                                   desc='remove dest directory when copying dirs')
+    # Set this if user wishes to have local copy of files as well
+    local_copy = traits.Str(desc='Copy files locally as well as to S3 bucket')
 
     def __setattr__(self, key, value):
         if key not in self.copyable_trait_names():

nipype/interfaces/tests/test_io.py

@@ -177,7 +177,7 @@ def _check_for_fakes3():
 
     # Check for fakes3
     try:
-        ret_code = subprocess.check_call(['which', 'fakes3'])
+        ret_code = subprocess.check_call(['which', 'fakes3'], stdout=open(os.devnull, 'wb'))
         if ret_code == 0:
             fakes3_found = True
     except subprocess.CalledProcessError as exc:
@@ -188,7 +188,29 @@ def _check_for_fakes3():
     # Return if found
     return fakes3_found
 
-@skipif(noboto3)
+def _make_dummy_input():
+    '''
+    '''
+
+    # Import packages
+    import tempfile
+
+    # Init variables
+    input_dir = tempfile.mkdtemp()
+    input_path = os.path.join(input_dir, 'datasink_test_s3.txt')
+
+    # Create input file
+    with open(input_path, 'wb') as f:
+        f.write('ABCD1234')
+
+    # Return path
+    return input_path
+
+# Check for fakes3
+fakes3 = _check_for_fakes3()
+
+
+@skipif(noboto3 or not fakes3)
 # Test datasink writes to s3 properly
 def test_datasink_to_s3():
     '''
@@ -208,13 +230,7 @@ def test_datasink_to_s3():
     output_dir = 's3://' + bucket_name
     # Local temporary filepaths for testing
     fakes3_dir = tempfile.mkdtemp()
-    input_dir = tempfile.mkdtemp()
-    input_path = os.path.join(input_dir, 'datasink_test_s3.txt')
-
-    # Check for fakes3
-    fakes3_found = _check_for_fakes3()
-    if not fakes3_found:
-        return
+    input_path = _make_dummy_input()
 
     # Start up fake-S3 server
     proc = Popen(['fakes3', '-r', fakes3_dir, '-p', '4567'], stdout=open(os.devnull, 'wb'))
@@ -230,10 +246,6 @@ def test_datasink_to_s3():
     # Create bucket
     bucket = resource.create_bucket(Bucket=bucket_name)
 
-    # Create input file
-    with open(input_path, 'wb') as f:
-        f.write('ABCD1234')
-
     # Prep datasink
     ds.inputs.base_directory = output_dir
     ds.inputs.container = container
@@ -249,15 +261,59 @@ def test_datasink_to_s3():
     dst_md5 = obj.e_tag.replace('"', '')
     src_md5 = hashlib.md5(open(input_path, 'rb').read()).hexdigest()
 
-    # Make sure md5sums match
-    yield assert_equal, src_md5, dst_md5
-
     # Kill fakes3
     proc.kill()
 
     # Delete fakes3 folder and input file
     shutil.rmtree(fakes3_dir)
-    shutil.rmtree(input_dir)
+    shutil.rmtree(os.path.dirname(input_path))
+
+    # Make sure md5sums match
+    yield assert_equal, src_md5, dst_md5
+
+# Test the local copy attribute
+def test_datasink_localcopy():
+    '''
+    Function to validate DataSink will make local copy via local_copy
+    attribute
+    '''
+
+    # Import packages
+    import hashlib
+    import tempfile
+
+    # Init variables
+    local_dir = tempfile.mkdtemp()
+    container = 'outputs'
+    attr_folder = 'text_file'
+
+    # Make dummy input file and datasink
+    input_path = _make_dummy_input()
+    ds = nio.DataSink()
+
+    # Set up datasink
+    ds.inputs.container = container
+    ds.inputs.local_copy = local_dir
+    setattr(ds.inputs, attr_folder, input_path)
+
+    # Expected local copy path
+    local_copy = os.path.join(local_dir, container, attr_folder,
+                              os.path.basename(input_path))
+
+    # Run the datasink
+    ds.run()
+
+    # Check md5sums of both
+    src_md5 = hashlib.md5(open(input_path, 'rb').read()).hexdigest()
+    dst_md5 = hashlib.md5(open(local_copy, 'rb').read()).hexdigest()
+
+    # Delete temp diretories
+    shutil.rmtree(os.path.dirname(input_path))
+    shutil.rmtree(local_dir)
+
+    # Perform test
+    yield assert_equal, src_md5, dst_md5
+
 
 @skipif(noboto)
 def test_s3datasink():
@@ -300,7 +356,7 @@ def test_datasink_substitutions():
     shutil.rmtree(indir)
     shutil.rmtree(outdir)
 
-@skipif(noboto)
+@skipif(noboto or not fakes3)
 def test_s3datasink_substitutions():
     indir = mkdtemp(prefix='-Tmp-nipype_ds_subs_in')
     outdir = mkdtemp(prefix='-Tmp-nipype_ds_subs_out')