From 56b8f5078a9ce1e24fb8120a38e78cbfeddea896 Mon Sep 17 00:00:00 2001 From: lokijuhy Date: Wed, 10 Mar 2021 14:36:27 +0100 Subject: [PATCH] clean up failed SAD dir before erroring --- datatc/self_aware_data.py | 39 ++++++++++++++++++++++------------- tests/test_self_aware_data.py | 14 +++++++++++++ 2 files changed, 39 insertions(+), 14 deletions(-) diff --git a/datatc/self_aware_data.py b/datatc/self_aware_data.py index 75f07cf..e077a78 100644 --- a/datatc/self_aware_data.py +++ b/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 diff --git a/tests/test_self_aware_data.py b/tests/test_self_aware_data.py index 04edae3..a7510c3 100644 --- a/tests/test_self_aware_data.py +++ b/tests/test_self_aware_data.py @@ -1,4 +1,5 @@ import unittest +import glob import os import pandas as pd from pathlib import Path @@ -254,3 +255,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)