Clean up directories on failure. Closes pre-commit#58.

asottile · asottile · commit bcb00726a107 · 2014-04-03T22:54:27.000-07:00
diff --git a/pre_commit/languages/node.py b/pre_commit/languages/node.py
@@ -3,6 +3,7 @@
 from pre_commit.languages import helpers
 from pre_commit.languages import python
 from pre_commit.prefixed_command_runner import CalledProcessError
+from pre_commit.util import clean_path_on_failure
 
 
 NODE_ENV = 'node_env'
@@ -30,22 +31,30 @@ def install_environment(repo_cmd_runner):
     if repo_cmd_runner.exists(NODE_ENV):
         return
 
-    repo_cmd_runner.run(['virtualenv', '{{prefix}}{0}'.format(python.PY_ENV)])
-
-    with python.in_env(repo_cmd_runner) as python_env:
-        python_env.run('pip install nodeenv')
-
-        # Try and use the system level node executable first
-        try:
-            python_env.run('nodeenv -n system {{prefix}}{0}'.format(NODE_ENV))
-        except CalledProcessError:
-            # TODO: log failure here
-            # cleanup
-            # TODO: local.path(NODE_ENV).delete()
-            python_env.run('nodeenv --jobs 4 {{prefix}}{0}'.format(NODE_ENV))
+    with clean_path_on_failure(repo_cmd_runner.path(python.PY_ENV)):
+        repo_cmd_runner.run(
+            ['virtualenv', '{{prefix}}{0}'.format(python.PY_ENV)],
+        )
 
-        with in_env(repo_cmd_runner) as node_env:
-            node_env.run('cd {prefix} && npm install -g')
+        with python.in_env(repo_cmd_runner) as python_env:
+            python_env.run('pip install nodeenv')
+
+            with clean_path_on_failure(repo_cmd_runner.path(NODE_ENV)):
+                # Try and use the system level node executable first
+                try:
+                    python_env.run(
+                        'nodeenv -n system {{prefix}}{0}'.format(NODE_ENV),
+                    )
+                except CalledProcessError:
+                    # TODO: log failure here
+                    # cleanup
+                    # TODO: local.path(NODE_ENV).delete()
+                    python_env.run(
+                        'nodeenv --jobs 4 {{prefix}}{0}'.format(NODE_ENV),
+                    )
+
+                with in_env(repo_cmd_runner) as node_env:
+                    node_env.run('cd {prefix} && npm install -g')
 
 
 def run_hook(repo_cmd_runner, hook, file_args):
diff --git a/pre_commit/languages/python.py b/pre_commit/languages/python.py
@@ -2,6 +2,7 @@
 import contextlib
 
 from pre_commit.languages import helpers
+from pre_commit.util import clean_path_on_failure
 
 
 PY_ENV = 'py_env'
@@ -25,9 +26,10 @@ def install_environment(repo_cmd_runner):
         return
 
     # Install a virtualenv
-    repo_cmd_runner.run(['virtualenv', '{{prefix}}{0}'.format(PY_ENV)])
-    with in_env(repo_cmd_runner) as env:
-        env.run('cd {prefix} && pip install .')
+    with clean_path_on_failure(repo_cmd_runner.path(PY_ENV)):
+        repo_cmd_runner.run(['virtualenv', '{{prefix}}{0}'.format(PY_ENV)])
+        with in_env(repo_cmd_runner) as env:
+            env.run('cd {prefix} && pip install .')
 
 
 def run_hook(repo_cmd_runner, hook, file_args):
diff --git a/pre_commit/languages/ruby.py b/pre_commit/languages/ruby.py
@@ -2,6 +2,7 @@
 import contextlib
 
 from pre_commit.languages import helpers
+from pre_commit.util import clean_path_on_failure
 
 
 RVM_ENV = 'rvm_env'
@@ -23,9 +24,10 @@ def install_environment(repo_cmd_runner):
     if repo_cmd_runner.exists(RVM_ENV):
         return
 
-    repo_cmd_runner.run(['__rvm-env.sh', '{{prefix}}{0}'.format(RVM_ENV)])
-    with in_env(repo_cmd_runner) as env:
-        env.run('cd {prefix} && bundle install')
+    with clean_path_on_failure(repo_cmd_runner.path(RVM_ENV)):
+        repo_cmd_runner.run(['__rvm-env.sh', '{{prefix}}{0}'.format(RVM_ENV)])
+        with in_env(repo_cmd_runner) as env:
+            env.run('cd {prefix} && bundle install')
 
 
 def run_hook(repo_cmd_runner, hook, file_args):
diff --git a/pre_commit/repository.py b/pre_commit/repository.py
@@ -9,6 +9,7 @@
 from pre_commit.ordereddict import OrderedDict
 from pre_commit.prefixed_command_runner import PrefixedCommandRunner
 from pre_commit.util import cached_property
+from pre_commit.util import clean_path_on_failure
 
 
 class Repository(object):
@@ -62,9 +63,10 @@ def create(self):
                 # Project already exists, no reason to re-create it
                 return
 
-            local['git']['clone', '--no-checkout', self.repo_url, self.sha]()
-            with self.in_checkout():
-                local['git']['checkout', self.sha]()
+            with clean_path_on_failure(unicode(local.path(self.sha))):
+                local['git']['clone', '--no-checkout', self.repo_url, self.sha]()
+                with self.in_checkout():
+                    local['git']['checkout', self.sha]()
 
     def require_installed(self, cmd_runner):
         if self.__installed:
diff --git a/pre_commit/util.py b/pre_commit/util.py
@@ -1,6 +1,9 @@
 
+import contextlib
 import functools
 import os
+import os.path
+import shutil
 import sys
 
 
@@ -49,3 +52,14 @@ def wrapper(argv=None):
             argv = sys.argv[1:]
         return func(argv)
     return wrapper
+
+
+@contextlib.contextmanager
+def clean_path_on_failure(path):
+    """Cleans up the directory on an exceptional failure."""
+    try:
+        yield
+    except BaseException:
+        if os.path.exists(path):
+            shutil.rmtree(path)
+        raise
diff --git a/tests/conftest.py b/tests/conftest.py
@@ -13,12 +13,17 @@
 
 
 @pytest.yield_fixture
-def empty_git_dir(tmpdir):
+def in_tmpdir(tmpdir):
     with local.cwd(tmpdir.strpath):
-        local['git']['init']()
         yield tmpdir.strpath
 
 
+@pytest.yield_fixture
+def empty_git_dir(in_tmpdir):
+    local['git']['init']()
+    yield in_tmpdir
+
+
 def add_and_commit():
     local['git']['add', '.']()
     local['git']['commit', '-m', 'random commit {0}'.format(time.time())]()
diff --git a/tests/prefixed_command_runner_test.py b/tests/prefixed_command_runner_test.py
@@ -3,7 +3,6 @@
 import mock
 import pytest
 import subprocess
-from plumbum import local
 
 from pre_commit.prefixed_command_runner import _replace_cmd
 from pre_commit.prefixed_command_runner import CalledProcessError
@@ -110,23 +109,20 @@ def test_from_command_runner_preserves_popen(popen_mock, makedirs_mock):
     makedirs_mock.assert_called_once_with('foo/bar/')
 
 
-def test_create_path_if_not_exists(tmpdir):
-    with local.cwd(tmpdir.strpath):
-        instance = PrefixedCommandRunner('foo')
-        assert not os.path.exists('foo')
-        instance._create_path_if_not_exists()
-        assert os.path.exists('foo')
+def test_create_path_if_not_exists(in_tmpdir):
+    instance = PrefixedCommandRunner('foo')
+    assert not os.path.exists('foo')
+    instance._create_path_if_not_exists()
+    assert os.path.exists('foo')
 
 
-def test_exists_does_not_exist(tmpdir):
-    with local.cwd(tmpdir.strpath):
-        assert not PrefixedCommandRunner('.').exists('foo')
+def test_exists_does_not_exist(in_tmpdir):
+    assert not PrefixedCommandRunner('.').exists('foo')
 
 
-def test_exists_does_exist(tmpdir):
-    with local.cwd(tmpdir.strpath):
-        os.mkdir('foo')
-        assert PrefixedCommandRunner('.').exists('foo')
+def test_exists_does_exist(in_tmpdir):
+    os.mkdir('foo')
+    assert PrefixedCommandRunner('.').exists('foo')
 
 
 def test_raises_on_error(popen_mock, makedirs_mock):
diff --git a/tests/util_test.py b/tests/util_test.py
@@ -1,11 +1,14 @@
 
 import mock
 import pytest
+import os
+import os.path
 import random
 import sys
 from plumbum import local
 
 from pre_commit.util import cached_property
+from pre_commit.util import clean_path_on_failure
 from pre_commit.util import entry
 from pre_commit.util import memoize_by_cwd
 
@@ -84,3 +87,35 @@ def test_no_arguments_passed_uses_argv(entry_func):
     with mock.patch.object(sys, 'argv', argv):
         ret = entry_func()
         assert ret == argv[1:]
+
+
+def test_clean_on_failure_noop(in_tmpdir):
+    with clean_path_on_failure('foo'): pass
+
+
+def test_clean_path_on_failure_does_nothing_when_not_raising(in_tmpdir):
+    with clean_path_on_failure('foo'):
+        os.mkdir('foo')
+    assert os.path.exists('foo')
+
+
+def test_clean_path_on_failure_cleans_for_normal_exception(in_tmpdir):
+    class MyException(Exception): pass
+
+    with pytest.raises(MyException):
+        with clean_path_on_failure('foo'):
+            os.mkdir('foo')
+            raise MyException
+
+    assert not os.path.exists('foo')
+
+
+def test_clean_path_on_failure_cleans_for_system_exit(in_tmpdir):
+    class MySystemExit(SystemExit): pass
+
+    with pytest.raises(MySystemExit):
+        with clean_path_on_failure('foo'):
+            os.mkdir('foo')
+            raise MySystemExit
+
+    assert not os.path.exists('foo')
