Extract 'archive_context' from _install and _build_egg

1 files changed, 17 insertions, 26 deletions

diff --git a/ez_setup.py b/ez_setup.py
index 548abe71..e4a357eb 100644
--- a/ez_setup.py
+++ b/ez_setup.py
@@ -22,6 +22,7 @@ import optparse
 import subprocess
 import platform
 import textwrap
+import contextlib
 
 from distutils import log
 
@@ -40,21 +41,9 @@ def _python_cmd(*args):
     args = (sys.executable,) + args
     return subprocess.call(args) == 0
 
-def _install(archive_filename, install_args=()):
-    # extracting the archive
-    tmpdir = tempfile.mkdtemp()
-    log.warn('Extracting in %s', tmpdir)
-    old_wd = os.getcwd()
-    try:
-        os.chdir(tmpdir)
-        with zipfile.ZipFile(archive_filename) as archive:
-            archive.extractall()
-
-        # going in the directory
-        subdir = os.path.join(tmpdir, os.listdir(tmpdir)[0])
-        os.chdir(subdir)
-        log.warn('Now working in %s', subdir)
 
+def _install(archive_filename, install_args=()):
+    with archive_context(archive_filename):
         # installing
         log.warn('Installing Setuptools')
         if not _python_cmd('setup.py', 'install', *install_args):
@@ -62,37 +51,39 @@ def _install(archive_filename, install_args=()):
             log.warn('See the error message above.')
             # exitcode will be 2
             return 2
-    finally:
-        os.chdir(old_wd)
-        shutil.rmtree(tmpdir)
 
 
 def _build_egg(egg, archive_filename, to_dir):
+    with archive_context(archive_filename):
+        # building an egg
+        log.warn('Building a Setuptools egg in %s', to_dir)
+        _python_cmd('setup.py', '-q', 'bdist_egg', '--dist-dir', to_dir)
+    # returning the result
+    log.warn(egg)
+    if not os.path.exists(egg):
+        raise IOError('Could not build the egg.')
+
+
+@contextlib.contextmanager
+def archive_context(filename):
     # extracting the archive
     tmpdir = tempfile.mkdtemp()
     log.warn('Extracting in %s', tmpdir)
     old_wd = os.getcwd()
     try:
         os.chdir(tmpdir)
-        with zipfile.ZipFile(archive_filename) as archive:
+        with zipfile.ZipFile(filename) as archive:
             archive.extractall()
 
         # going in the directory
         subdir = os.path.join(tmpdir, os.listdir(tmpdir)[0])
         os.chdir(subdir)
         log.warn('Now working in %s', subdir)
-
-        # building an egg
-        log.warn('Building a Setuptools egg in %s', to_dir)
-        _python_cmd('setup.py', '-q', 'bdist_egg', '--dist-dir', to_dir)
+        yield
 
     finally:
         os.chdir(old_wd)
         shutil.rmtree(tmpdir)
-    # returning the result
-    log.warn(egg)
-    if not os.path.exists(egg):
-        raise IOError('Could not build the egg.')
 
 
 def _do_download(version, download_base, to_dir, download_delay):