Exits working directory if wrapped function throws an exception. (#175)

* Exits working directory if wrapped function throws an exception

* Adds jevandezande to README

README.md

@@ -92,3 +92,4 @@ If **autodE** is used in a publication please consider citing the [paper](https:
 - Kjell Jorner ([@kjelljorner](https://github.com/kjelljorner))
 - Thibault Lestang ([@tlestang](https://github.com/tlestang))
 - Domen Pregeljc ([@dpregeljc](https://github.com/dpregeljc))
+- Jonathon Vandezande ([@jevandezande](https://github.com/jevandezande))

autode/utils.py

@@ -143,13 +143,15 @@ def wrapped_function(*args, **kwargs):
                 os.mkdir(dir_path)
 
             os.chdir(dir_path)
-            result = func(*args, **kwargs)
-            os.chdir(here)
+            try:
+                result = func(*args, **kwargs)
+            finally:
+                os.chdir(here)
 
-            if len(os.listdir(dir_path)) == 0:
-                logger.warning(f'Worked in {dir_path} but made no files '
-                               f'- deleting')
-                os.rmdir(dir_path)
+                if len(os.listdir(dir_path)) == 0:
+                    logger.warning(f'Worked in {dir_path} but made no files '
+                                   f'- deleting')
+                    os.rmdir(dir_path)
 
             return result
 
@@ -204,20 +206,23 @@ def wrapped_function(*args, **kwargs):
             # Move directories and execute
             os.chdir(tmpdir_path)
 
-            logger.info('Function   ...running')
-            result = func(*args, **kwargs)
-            logger.info('           ...done')
+            try:
+                logger.info('Function   ...running')
+                result = func(*args, **kwargs)
+                logger.info('           ...done')
+
+                for filename in os.listdir(tmpdir_path):
 
-            for filename in os.listdir(tmpdir_path):
+                    if any([filename.endswith(ext) for ext in kept_file_exts]):
+                        logger.info(f'Copying back {filename}')
+                        shutil.copy(filename, here)
 
-                if any([filename.endswith(ext) for ext in kept_file_exts]):
-                    logger.info(f'Copying back {filename}')
-                    shutil.copy(filename, here)
+            finally:
+                os.chdir(here)
 
-            os.chdir(here)
+                logger.info('Removing temporary directory')
+                shutil.rmtree(tmpdir_path)
 
-            logger.info('Removing temporary directory')
-            shutil.rmtree(tmpdir_path)
             return result
 
         return wrapped_function

tests/test_utils.py

@@ -29,6 +29,20 @@ def make_test_files():
     os.rmdir('test')
 
 
+def test_reset_dir_on_error():
+    @utils.work_in("tmp_path")
+    def raise_error():
+        assert 0
+
+    here = os.getcwd()
+    try:
+        raise_error()
+    except AssertionError:
+        pass
+
+    assert here == os.getcwd()
+
+
 def test_monitored_external():
 
     echo = ['echo', 'test']
@@ -77,6 +91,20 @@ def test():
     os.remove('test.txt')
 
 
+def test_reset_tmp_dir_on_error():
+    @utils.work_in_tmp_dir()
+    def raise_error():
+        assert 0
+
+    here = os.getcwd()
+    try:
+        raise_error()
+    except AssertionError:
+        pass
+
+    assert here == os.getcwd()
+
+
 @work_in_tmp_dir(filenames_to_copy=[], kept_file_exts=[])
 def test_calc_output():