diff --git a/mglib/storage.py b/mglib/storage.py index 6976398..4c29bc2 100644 --- a/mglib/storage.py +++ b/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( diff --git a/mglib/test/utils.py b/mglib/test/utils.py index 945776f..db0dee9 100644 --- a/mglib/test/utils.py +++ b/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 diff --git a/mglib/utils.py b/mglib/utils.py index 0640640..ee80f70 100644 --- a/mglib/utils.py +++ b/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 + diff --git a/test/test_storage.py b/test/test_storage.py index 05b76b2..c12dcaa 100644 --- a/test/test_storage.py +++ b/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() )