From 44b95f23a36bc67258eaa89a7faeadea08bc481a Mon Sep 17 00:00:00 2001
From: Ruslan Kuprieiev <ruslan@iterative.ai>
Date: Mon, 15 Apr 2019 02:07:46 +0300
Subject: [PATCH] remote: ssh: support directories

Signed-off-by: Ruslan Kuprieiev <ruslan@iterative.ai>
---
 dvc/remote/ssh/__init__.py               | 15 +++++-
 dvc/remote/ssh/connection.py             | 69 ++++++++++++++++++------
 tests/unit/remote/ssh/test_connection.py |  4 +-
 3 files changed, 69 insertions(+), 19 deletions(-)

diff --git a/dvc/remote/ssh/__init__.py b/dvc/remote/ssh/__init__.py
index 2f2e8a9d6..6d1b96238 100644
--- a/dvc/remote/ssh/__init__.py
+++ b/dvc/remote/ssh/__init__.py
@@ -93,7 +93,7 @@ class RemoteSSH(RemoteBase):
         assert path_info["scheme"] == self.scheme
 
         with self.ssh(**path_info) as ssh:
-            return ssh.file_exists(path_info["path"])
+            return ssh.exists(path_info["path"])
 
     def get_file_checksum(self, path_info):
         if path_info["scheme"] != self.scheme:
@@ -102,6 +102,10 @@ class RemoteSSH(RemoteBase):
         with self.ssh(**path_info) as ssh:
             return ssh.md5(path_info["path"])
 
+    def isdir(self, path_info):
+        with self.ssh(**path_info) as ssh:
+            return ssh.isdir(path_info["path"])
+
     def copy(self, from_info, to_info):
         if (
             from_info["scheme"] != self.scheme
@@ -203,3 +207,12 @@ class RemoteSSH(RemoteBase):
     def list_cache_paths(self):
         with self.ssh(**self.path_info) as ssh:
             return list(ssh.walk_files(self.prefix))
+
+    def walk(self, path_info):
+        with self.ssh(**path_info) as ssh:
+            for entry in ssh.walk(path_info["path"]):
+                yield entry
+
+    def makedirs(self, path_info):
+        with self.ssh(**path_info) as ssh:
+            ssh.makedirs(path_info["path"])
diff --git a/dvc/remote/ssh/connection.py b/dvc/remote/ssh/connection.py
index 7db8b45b5..1241ef0ae 100644
--- a/dvc/remote/ssh/connection.py
+++ b/dvc/remote/ssh/connection.py
@@ -2,6 +2,7 @@ import errno
 import os
 import posixpath
 import logging
+from stat import S_ISDIR
 
 try:
     import paramiko
@@ -72,7 +73,7 @@ class SSHConnection:
 
         self._ssh.close()
 
-    def file_exists(self, path):
+    def exists(self, path):
         self._sftp_connect()
         try:
             return self._sftp.stat(path)
@@ -118,7 +119,7 @@ class SSHConnection:
                 "a file with the same name '{}' already exists".format(path)
             )
 
-        head, tail = os.path.split(path)
+        head, tail = posixpath.split(path)
 
         if head:
             self.makedirs(head)
@@ -126,8 +127,11 @@ class SSHConnection:
         if tail:
             self._sftp.mkdir(path)
 
-    def walk_files(self, directory):
-        from stat import S_ISLNK, S_ISDIR, S_ISREG
+    def walk(self, directory, topdown=True):
+        # NOTE: original os.walk() implementation [1] with default options was
+        # used as a template.
+        #
+        # [1] https://github.com/python/cpython/blob/master/Lib/os.py
 
         self._sftp_connect()
 
@@ -141,28 +145,61 @@ class SSHConnection:
                 cause=exc,
             )
 
+        dirs = []
+        nondirs = []
         for entry in dir_entries:
-            path = posixpath.join(directory, entry.filename)
-
-            if S_ISLNK(entry.st_mode):
-                path = self._sftp.readlink(directory)
-                entry = self._sftp.stat(path)
-
+            name = entry.filename
             if S_ISDIR(entry.st_mode):
-                for inner_path in self.walk_files(path):
-                    yield inner_path
+                dirs.append(name)
+            else:
+                nondirs.append(name)
 
-            elif S_ISREG(entry.st_mode):
-                yield path
+        if topdown:
+            yield directory, dirs, nondirs
 
-    def remove(self, path):
-        self._sftp_connect()
+        for dname in dirs:
+            newpath = posixpath.join(directory, dname)
+            for entry in self.walk(newpath, topdown=topdown):
+                yield entry
+
+        if not topdown:
+            yield directory, dirs, nondirs
+
+    def walk_files(self, directory):
+        for root, dirs, files in self.walk(directory):
+            for fname in files:
+                yield posixpath.join(root, fname)
+
+    def _remove_file(self, path):
         try:
             self._sftp.remove(path)
         except IOError as exc:
             if exc.errno != errno.ENOENT:
                 raise
 
+    def _remove_dir(self, path):
+        for root, dirs, files in self.walk(path, topdown=False):
+            for fname in files:
+                path = posixpath.join(root, fname)
+                self._remove_file(path)
+
+            for dname in dirs:
+                path = posixpath.join(root, dname)
+                self._sftp.rmdir(dname)
+        try:
+            self._sftp.rmdir(path)
+        except IOError as exc:
+            if exc.errno != errno.ENOENT:
+                raise
+
+    def remove(self, path):
+        self._sftp_connect()
+
+        if self.isdir(path):
+            self._remove_dir(path)
+        else:
+            self._remove_file(path)
+
     def download(self, src, dest, no_progress_bar=False, progress_title=None):
         self._sftp_connect()
 
diff --git a/tests/unit/remote/ssh/test_connection.py b/tests/unit/remote/ssh/test_connection.py
index 94a1d3eb4..a570f3213 100644
--- a/tests/unit/remote/ssh/test_connection.py
+++ b/tests/unit/remote/ssh/test_connection.py
@@ -37,8 +37,8 @@ def test_isdir(server):
     assert isdir is True
 
 
-def test_file_exists(server):
+def test_exists(server):
     client = _client(server)
     path = "/path/to/file"
-    file_path = client.file_exists(path=path)
+    file_path = client.exists(path=path)
     assert file_path is False