testing location, documentpath, pagepath

mglib/storage.py

@@ -1,25 +1,11 @@
 import os
 import logging
 import shutil
+from mglib.utils import safe_to_delete
 
 logger = logging.getLogger(__name__)
 
 
-def safe_to_delete(place, safe_extensions):
-    if not os.path.exists(place):
-        logging.warning(
-            f"Trying to delete not exising folder"
-            f" {place}"
-        )
-        return False
-
-    for root, dirs, files in os.walk(place):
-        for name in files:
-            print(name)
-
-    return False
-
-
 class Storage:
     """
     Default Storage class which works with DocumentPath and PagePath
@@ -57,16 +43,18 @@ class Storage:
         # double check that there are only
         # .pdf, .txt, .hocr, .jpg files.
         if safe_to_delete(
-            abs_dirname_docs,
-            safe_extensions=['pdf']
+            abs_dirname_docs
         ):
             shutil.rmtree(abs_dirname_docs)
+            if os.path.exists(abs_dirname_docs):
+                os.rmdir(abs_dirname_docs)
 
         if safe_to_delete(
-            abs_dirname_results,
-            safe_extensions=['txt', 'jpg', 'hocr']
+            abs_dirname_results
         ):
             shutil.rmtree(abs_dirname_results)
+            if os.path.exists(abs_dirname_results):
+                os.rmdir(abs_dirname_results)
 
     def exists(self, _path):
         return os.path.exists(

mglib/test/utils.py

@@ -1,4 +1,5 @@
 import os
+import shutil
 
 
 class TemporaryNode:
@@ -41,14 +42,14 @@ class TemporaryNode:
 
     def __enter__(self):
         if not os.path.exists(self.location):
-            os.makedirs(self.location)
+            os.makedirs(self.location, exist_ok=True)
 
         return self
 
     def __exit__(self, exc_type, exc_value, traceback):
         if os.path.exists(self.location):
             if os.path.isdir(self.location):
-                os.rmdir(self.location)
+                shutil.rmtree(self.location)
             else:
                 os.remove(self.location)
 
@@ -57,12 +58,11 @@ class TemporaryNode:
             self.location,
             folder
         )
-        os.makedirs(new_location)
+        os.makedirs(new_location, exist_ok=True)
         return TemporaryNode(new_location)
 
     def exists(self):
         return os.path.exists(self.location)
 
     def add_file(self, file):
-        pass
         return self

mglib/utils.py

@@ -1,4 +1,9 @@
 import os
+import logging
+
+logger = logging.getLogger(__name__)
+
+SAFE_EXTENSIONS = ['.txt', '.jpg', '.hocr', '.pdf']
 
 
 def get_bool(key, default="NO"):
@@ -14,3 +19,21 @@ def get_bool(key, default="NO"):
         return True
 
     return False
+
+
+def safe_to_delete(place):
+    if not os.path.exists(place):
+        logging.warning(
+            f"Trying to delete not exising folder"
+            f" {place}"
+        )
+        return False
+
+    for root, dirs, files in os.walk(place):
+        for name in files:
+            base, ext = os.path.splitext(name)
+            if ext not in SAFE_EXTENSIONS:
+                raise Exception("Trying to delete unsefe location")
+
+    return True
+

test/test_storage.py

@@ -31,13 +31,13 @@ class TestStorage(unittest.TestCase):
                 file_name='doku.pdf'
             )
 
+            self.assertTrue(
+                f1.exists()
+            )
+
             storage.delete_document(doc_path)
 
             self.assertFalse(
-                docs.exists()
-            )
-
-            self.assertFalse(
-                res.exists()
+                f1.exists()
             )