Merge pull request #25 from uzh-dqbm-cmi/bugfix_catch_dir_creation_fail

Bugfix: Catch SAD dir creation fail to not leave empty directory

datatc/self_aware_data.py

@@ -3,6 +3,7 @@
 import glob
 import inspect
 import os
+import shutil
 from pathlib import Path
 from typing import Any, Callable, Dict, List, Tuple, Type, Union
 
@@ -541,10 +542,15 @@ def save(cls, sad: SelfAwareData, parent_path: str, file_name: str,  **kwargs) -
         new_transform_dir_path = Path(parent_path, transform_dir_name)
         os.makedirs(new_transform_dir_path)
 
-        data_interface = MagicDataInterface.select_data_interface(data_file_type)
-        data_interface.save(sad.data, 'data', new_transform_dir_path, **kwargs)
-        cls.file_component_interfaces['func'].save(transformer_func, 'func', new_transform_dir_path)
-        cls.file_component_interfaces['code'].save(code, 'code', new_transform_dir_path)
+        try:
+            data_interface = MagicDataInterface.select_data_interface(data_file_type)
+            data_interface.save(sad.data, 'data', new_transform_dir_path, **kwargs)
+            cls.file_component_interfaces['func'].save(transformer_func, 'func', new_transform_dir_path)
+            cls.file_component_interfaces['code'].save(code, 'code', new_transform_dir_path)
+        except Exception:
+            # clean up the failed SAD dir before raising the error
+            shutil.rmtree(new_transform_dir_path)
+            raise
 
         print('created new file {}'.format(new_transform_dir_path))
         return new_transform_dir_path
@@ -660,16 +666,21 @@ def save(cls, sad: SelfAwareData, parent_path: str, file_name: str,  **kwargs) -
         new_transform_dir_path = Path(parent_path, transform_dir_name)
         os.makedirs(new_transform_dir_path)
 
-        data_interface = MagicDataInterface.select_data_interface(data_file_type)
-        data_interface.save(sad.data, 'data', new_transform_dir_path, **kwargs)
-
-        cls.file_component_interfaces['sad'].save(sad, 'sad', new_transform_dir_path)
-
-        provenance = {
-            'interface_version': cls.version,
-            'transform_steps': sad.get_info()
-        }
-        cls.file_component_interfaces['provenance'].save(provenance, 'provenance', new_transform_dir_path)
+        try:
+            data_interface = MagicDataInterface.select_data_interface(data_file_type)
+            data_interface.save(sad.data, 'data', new_transform_dir_path, **kwargs)
+
+            cls.file_component_interfaces['sad'].save(sad, 'sad', new_transform_dir_path)
+
+            provenance = {
+                'interface_version': cls.version,
+                'transform_steps': sad.get_info()
+            }
+            cls.file_component_interfaces['provenance'].save(provenance, 'provenance', new_transform_dir_path)
+        except Exception:
+            # clean up the failed SAD dir before raising the error
+            shutil.rmtree(new_transform_dir_path)
+            raise
 
         print('created new SAD directory {}'.format(new_transform_dir_path))
         return new_transform_dir_path

tests/test_self_aware_data.py

@@ -1,4 +1,5 @@
 import unittest
+import glob
 import os
 import pandas as pd
 from pathlib import Path
@@ -264,3 +265,16 @@ def test_load_interface_version_0(self):
         self.assertTrue(set(transform_step_0_info.keys()) == expected_transform_info_keys)
         for key in expected_info:
             self.assertEqual(transform_step_0_info[key], expected_info[key])
+
+    def test_save_fail_leaves_no_dir(self):
+        raw_sad = SelfAwareData(self.raw_df)
+        my_sad = raw_sad.transform(self.transform_func, enforce_clean_git=False)
+        sad_path = Path(self.test_dir, 'sad_fail.bad_file_ext')
+
+        with self.assertRaises(Exception) as context:
+            sad_file_path = my_sad.save(sad_path, index=False)
+
+        glob_path = Path(self.test_dir, '*')
+        subpaths = glob.glob(glob_path.__str__())
+        sad_dirs = [os.path.basename(file_path) for file_path in subpaths if 'sad_dir' in file_path]
+        self.assertEqual(len(sad_dirs), 0)