Rebased copy remote files branch. Updated copy remote file code to work correctly with tree dir structures. Added tests for recursion on/off and directory structure tests. Updated docstrings. Resolves #40

Author: Dan

pssh/pssh_client.py

@@ -843,13 +843,14 @@ def copy_file(self, local_file, remote_file, recurse=False):
                                 {'recurse' : recurse})
                 for host in self.hosts]
 
-    def _copy_file(self, host, local_file, remote_file, recurse):
+    def _copy_file(self, host, local_file, remote_file, recurse=False):
         """Make sftp client, copy file"""
         self._make_ssh_client(host)
         return self.host_clients[host].copy_file(local_file, remote_file,
                                                  recurse=recurse)
 
-    def copy_file_to_local(self, remote_file, local_file, recurse=False):
+    def copy_remote_file(self, remote_file, local_file, recurse=False,
+                         suffix_separator='_'):
         """Copy remote file to local file in parallel
 
         :param remote_file: remote filepath to copy to local host
@@ -858,25 +859,34 @@ def copy_file_to_local(self, remote_file, local_file, recurse=False):
         :type local_file: str
         :param recurse: whether or not to recurse
         :type recurse: bool
+        :param suffix_separator: (Optional) Separator string between \
+        filename and host, defaults to ``_``. Eg for a ``local_file`` value of \
+        ``my_file`` and default seaparator the resulting filename will be \
+        ``my_file_my_host`` for the file from host ``my_host``
+        :type suffix_separator: str
         .. note ::
           Local directories in `local_file` that do not exist will be
           created as long as permissions allow.
 
         .. note ::
           File names will be de-duplicated by appending the hostname to the
-          filepath.
+          filepath separated by ``suffix_separator``.
 
         :rtype: List(:mod:`gevent.Greenlet`) of greenlets for remote copy \
         commands
         """
-        return [self.pool.spawn(self._copy_file_to_local, host, remote_file, local_file, recurse)
-                for host in self.hosts]
+        return [self.pool.spawn(
+            self._copy_remote_file, host, remote_file,
+            local_file, recurse, suffix_separator=suffix_separator)
+            for host in self.hosts]
 
-    def _copy_file_to_local(self, host, remote_file, local_file, recurse):
+    def _copy_remote_file(self, host, remote_file, local_file, recurse,
+                          suffix_separator='_'):
         """Make sftp client, copy file to local"""
+        file_w_suffix = suffix_separator.join([local_file, host])
         self._make_ssh_client(host)
-        return self.host_clients[host].copy_file_to_local(
-                remote_file, '_'.join([local_file, host]), recurse=recurse)
+        return self.host_clients[host].copy_remote_file(
+                remote_file, file_w_suffix, recurse=recurse)
 
     def _make_ssh_client(self, host):
         if not host in self.host_clients or not self.host_clients[host]:

pssh/ssh_client.py

@@ -298,14 +298,14 @@ def _mkdir(self, sftp, directory):
 
     def mkdir(self, sftp, directory):
         """Make directory via SFTP channel.
-        
+
         Parent paths in the directory are created if they do not exist.
-        
+
         :param sftp: SFTP client object
         :type sftp: :mod:`paramiko.SFTPClient`
         :param directory: Remote directory to create
         :type directory: str
-        
+
         Catches and logs at error level remote IOErrors on creating directory.
         """
         try:
@@ -360,7 +360,7 @@ def copy_file(self, local_file, remote_file, recurse=False):
             raise ValueError("Recurse must be true if local_file is a "
                              "directory.")
         sftp = self._make_sftp()
-        destination = self._parent_path_split(remote_file)
+        destination = self._parent_paths_split(remote_file)
         try:
             sftp.stat(destination)
         except IOError:
@@ -375,70 +375,71 @@ def copy_file(self, local_file, remote_file, recurse=False):
         logger.info("Copied local file %s to remote destination %s:%s",
                     local_file, self.host, remote_file)
 
-    def _copy_dir_to_local(self, remote_dir, local_dir):
-        """Copies the remote directory to the local host."""
-        sftp = self._make_sftp()
-        file_list = sftp.listdir(remote_dir)
-        for file_name in file_list:
-            remote_path = os.path.join(remote_dir, file_name)
-            local_path = os.path.join(local_dir, file_name)
-            self.copy_file_to_local(remote_path, local_path, recurse=True)
-
-    def copy_file_to_local(self, remote_file, local_file, recurse=False):
+    def copy_remote_file(self, remote_file, local_file, recurse=False,
+                         sftp=None):
         """Copy remote file to local host via SFTP/SCP
 
-        Copy is done natively using SFTP/SCP version 2 protocol, no scp command \
+        Copy is done natively using SFTP/SCP version 2, no scp command \
         is used or required.
 
-        :param remote_file: Remote filepath to copy the file from.
+        :param remote_file: Remote filepath to copy from
         :type remote_file: str
-        :param local_file: Local filepath where the file will be copied.
+        :param local_file: Local filepath where file(s) will be copied to
         :type local_file: str
-        :param recurse: Whether or not to recursively copy directories.
+        :param recurse: Whether or not to recursively copy directories
         :type recurse: bool
 
         :raises: :mod:`ValueError` when a directory is supplied to remote_file \
         and recurse is not set
-        :raises: :mod:`OSError` on OS errors creating directories or file
-        :raises: :mod:`IOError` on IO errors creating directories or file
+        :raises: :mod:`IOError` on I/O errors creating directories or file
+        :raises: :mod:`OSError` on OS errors like permission denied
         """
-        sftp = self._make_sftp()
+        sftp = self._make_sftp() if not sftp else sftp
         try:
-            sftp.listdir(remote_file)
-        except (OSError, IOError):
-            remote_dir_exists = False
+            file_list = sftp.listdir(remote_file)
+        except IOError:
+            # remote_file is not dir
+            pass
         else:
-            remote_dir_exists = True
-        if remote_dir_exists and recurse:
-            return self._copy_dir_to_local(remote_file, local_file)
-        elif remote_dir_exists and not recurse:
-            raise ValueError("Recurse must be true if remote_file is a "
-                             "directory.")
-        destination = self._parent_path_split(local_file)
+            if not recurse:
+                raise ValueError("Recurse must be true if remote_file is a "
+                                 "directory.")
+            return self._copy_remote_dir(file_list, remote_file,
+                                         local_file, sftp)
+        destination = self._parent_paths_split(local_file)
         self._make_local_dir(destination)
         try:
             sftp.get(remote_file, local_file)
-        except Exception, error:
-            logger.error("Error occured copying file %s from remote destination %s:%s - %s",
+        except Exception as error:
+            logger.error("Error occured copying file %s from remote destination"
+                         " %s:%s - %s",
                          local_file, self.host, remote_file, error)
-            raise error
-        else:
-            logger.info("Copied local file %s from remote destination %s:%s",
-                        local_file, self.host, remote_file)
+            raise
+        logger.info("Copied local file %s from remote destination %s:%s",
+                    local_file, self.host, remote_file)
+
+    def _copy_remote_dir(self, file_list, remote_dir, local_dir, sftp):
+        for file_name in file_list:
+            remote_path = os.path.join(remote_dir, file_name)
+            local_path = os.path.join(local_dir, file_name)
+            self.copy_remote_file(remote_path, local_path, sftp=sftp,
+                                  recurse=True)
 
     def _make_local_dir(self, dirpath):
-        if not os.path.exists(dirpath):
-            try:
-                os.makedirs(dirpath)
-            except OSError:
-                logger.error("Unable to create local directory structure for "
-                             "directory %s", dirpath)
-                raise
+        if os.path.exists(dirpath):
+            return
+        try:
+            os.makedirs(dirpath)
+        except OSError:
+            logger.error("Unable to create local directory structure for "
+                         "directory %s", dirpath)
+            raise
 
-    def _parent_path_split(self, file_path):
+    def _parent_paths_split(self, file_path):
         try:
-            destination = [_dir for _dir in file_path.split(os.path.sep)
-                            if _dir][:-1][0]
+            destination = os.path.sep.join(
+                [_dir for _dir in file_path.split(os.path.sep)
+                 if _dir][:-1])
         except IndexError:
             destination = ''
         if file_path.startswith(os.path.sep) or not destination:

tests/test_pssh_client.py

@@ -451,31 +451,63 @@ def test_pssh_client_copy_file_failure(self):
         for path in [local_test_path, remote_test_path]:
             shutil.rmtree(path)
 
-    def test_pssh_copy_file_to_local(self):
+    def test_pssh_copy_remote_file(self):
         """Test parallel copy file to local host"""
         test_file_data = 'test'
-        remote_filename = 'test_file'
-        local_test_dir, local_filename = 'local_test_dir', 'test_file_copy'
-        local_filename = os.path.sep.join([local_test_dir, local_filename])
-        test_file = open(remote_filename, 'w')
-        test_file.writelines([test_file_data + os.linesep])
-        test_file.close()
-        server = start_server({ self.fake_cmd : self.fake_resp },
-                              self.listen_socket)
+        local_test_path = 'directory_test_local_remote_copied'
+        remote_test_path = 'directory_test_remote_copy'
+        local_copied_dir = '_'.join([local_test_path, self.host])
+        new_local_copied_dir = '.'.join([local_test_path, self.host])
+        for path in [local_test_path, remote_test_path, local_copied_dir,
+                     new_local_copied_dir]:
+            try:
+                shutil.rmtree(path)
+            except OSError:
+                try:
+                    os.unlink(path)
+                except Exception:
+                    pass
+                pass
+        os.mkdir(remote_test_path)
+        local_file_paths = []
+        for i in range(0, 10):
+            remote_file_path_dir = os.path.join(remote_test_path, 'dir_foo' + str(i))
+            os.mkdir(remote_file_path_dir)
+            remote_file_path = os.path.join(remote_file_path_dir, 'foo' + str(i))
+            local_file_path = os.path.join(local_copied_dir, 'dir_foo' + str(i), 'foo' + str(i))
+            local_file_paths.append(local_file_path)
+            test_file = open(remote_file_path, 'w')
+            test_file.write(test_file_data)
+            test_file.close()
         client = ParallelSSHClient([self.host], port=self.listen_port,
                                    pkey=self.user_key)
-        cmds = client.copy_file_to_local(remote_filename, local_filename)
-        cmds[0].get()
-        local_filename += '_' + self.host
-        self.assertTrue(os.path.isdir(local_test_dir),
-                        msg="SFTP create local directory failed")
-        self.assertTrue(os.path.isfile(local_filename),
-                        msg="SFTP copy failed")
-        for filepath in [remote_filename, local_filename]:
-            os.unlink(filepath)
-        shutil.rmtree(local_test_dir)
-        del client
-        server.join()
+        cmds = client.copy_remote_file(remote_test_path, local_test_path)
+        for cmd in cmds:
+            self.assertRaises(ValueError, cmd.get)
+        cmds = client.copy_remote_file(remote_test_path, local_test_path,
+                                       recurse=True)
+        for cmd in cmds:
+            cmd.get()
+        try:
+            self.assertTrue(os.path.isdir(local_copied_dir))
+            for path in local_file_paths:
+                self.assertTrue(os.path.isfile(path))
+        except Exception:
+            shutil.rmtree(remote_test_path)
+        finally:
+            shutil.rmtree(local_copied_dir)
+        cmds = client.copy_remote_file(remote_test_path, local_test_path,
+                                       suffix_separator='.', recurse=True)
+        for cmd in cmds:
+            cmd.get()
+        new_local_copied_dir = '.'.join([local_test_path, self.host])
+        try:
+            for path in local_file_paths:
+                path = path.replace(local_copied_dir, new_local_copied_dir)
+                self.assertTrue(os.path.isfile(path))
+        finally:
+            shutil.rmtree(new_local_copied_dir)
+            shutil.rmtree(remote_test_path)
 
     def test_pssh_pool_size(self):
         """Test setting pool size to non default values"""