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)