From 7c1cd10bfbf0fec8cfae9dfb3b1d0a9eaec1f477 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Micha=C5=82=20D?= <michal@datarevenue.com>
Date: Mon, 7 Sep 2020 13:40:13 +0200
Subject: [PATCH] Add recursive removal to local fs, expose it in s3 fs too.

#time 46m
---
 drfs/filesystems/local.py            | 12 +++++++++---
 drfs/filesystems/s3.py               |  3 +++
 drfs/tests/filesystems/test_local.py | 22 ++++++++++++++++++++++
 drfs/tests/filesystems/test_s3.py    | 12 ++++++++++++
 4 files changed, 46 insertions(+), 3 deletions(-)

diff --git a/drfs/filesystems/local.py b/drfs/filesystems/local.py
index ca969b2..04b35e5 100644
--- a/drfs/filesystems/local.py
+++ b/drfs/filesystems/local.py
@@ -38,9 +38,15 @@ def makedirs(self, *args, **kwargs):
         os.makedirs(*args, **kwargs)
 
     @allow_pathlib
-    def remove(self, path):
-        """Remove a file."""
-        os.remove(path)
+    def remove(self, path, recursive=False):
+        """Remove a file or a directory which may be non-empty."""
+        try:
+            os.remove(path)
+        except IsADirectoryError:
+            if not recursive:
+                self.rmdir(path)
+            else:
+                shutil.rmtree(path)
 
     @return_pathlib
     @allow_pathlib
diff --git a/drfs/filesystems/s3.py b/drfs/filesystems/s3.py
index e8c96ed..b7550d1 100644
--- a/drfs/filesystems/s3.py
+++ b/drfs/filesystems/s3.py
@@ -19,6 +19,9 @@ def makedirs(self, *args, **kwargs):
     def rmdir(self, path, **kwargs):
         pass
 
+    def rm(self, path, recursive=False, **kwargs):
+        return self.fs.rm(path, recursive=recursive, **kwargs)
+
     def put(self, filename, path, **kwargs):
         from drfs.path import asstr
         filename, path = asstr(filename), asstr(path)
diff --git a/drfs/tests/filesystems/test_local.py b/drfs/tests/filesystems/test_local.py
index 8585a16..c0f7bf7 100644
--- a/drfs/tests/filesystems/test_local.py
+++ b/drfs/tests/filesystems/test_local.py
@@ -61,3 +61,25 @@ def test_walk_scheme_local(tmpdir):
     assert all(isinstance(item, Path) for item in res)
     assert any('deep_test' in str(item) for item in res)
     assert not any(str(item).rstrip('/').endswith('dir1') for item in res)
+
+
+def test_remove(tmpdir):
+    fs = LocalFileSystem()
+    empty_dir = tmpdir / 'empty'
+    fs.makedirs(empty_dir)
+    dir1 = tmpdir / 'dir1'
+    file1 = dir1 / 'deep_test.txt'
+    fs.touch(file1)
+
+    assert fs.exists(empty_dir)
+    fs.remove(empty_dir)
+    assert not fs.exists(empty_dir)
+
+    with pytest.raises(OSError):
+        fs.remove(dir1)
+
+    assert fs.exists(file1)
+    fs.remove(dir1, recursive=True)
+    assert not fs.exists(file1)
+    assert not fs.exists(dir1)
+
diff --git a/drfs/tests/filesystems/test_s3.py b/drfs/tests/filesystems/test_s3.py
index 96faa2b..2cab8d9 100644
--- a/drfs/tests/filesystems/test_s3.py
+++ b/drfs/tests/filesystems/test_s3.py
@@ -38,3 +38,15 @@ def test_walk_scheme_s3(s3):
     assert all(str(item).startswith('s3://') for item in res)
     assert any('deep_test' in str(item) for item in res)
     assert not any(str(item).rstrip('/').endswith('dir1') for item in res)
+
+
+def test_remove_recursive(s3):
+    fs = S3FileSystem()
+    fs.touch("s3://test-bucket/dir1/deep_test.txt")
+    fs.touch("s3://test-bucket/dir1/subdir/deeper_test.txt")
+
+    fs.remove("s3://test-bucket/dir1", recursive=True)
+
+    assert not fs.exists("s3://test-bucket/dir1")
+    assert not fs.exists("s3://test-bucket/dir1/deep_test.txt")
+    assert not fs.exists("s3://test-bucket/dir1/subdir/deeper_test.txt")