morphlib: Use new temp_dir helper context manager

Change-Id: Ie4e024a63f2ab1c7ea66f2cbedaef99c9adf5e69

3 files changed, 38 insertions, 54 deletions

diff --git a/morphlib/plugins/artifact_inspection_plugin.py b/morphlib/plugins/artifact_inspection_plugin.py
index b643c57c..fc433a01 100644
--- a/morphlib/plugins/artifact_inspection_plugin.py
+++ b/morphlib/plugins/artifact_inspection_plugin.py
@@ -20,8 +20,6 @@ import os
 import re
 import contextlib
 
-import fs.tempfs
-
 import morphlib
 
 from morphlib.bins import call_in_artifact_directory
@@ -111,16 +109,16 @@ class AutotoolsVersionGuesser(ProjectVersionGuesser):
         return None
 
     def _check_autoconf_package_version(self, repo, ref, filename, data):
-        with contextlib.closing(fs.tempfs.TempFS(
-                temp_dir=self.app.settings['tempdir'])) as tempdir:
-            with open(tempdir.getsyspath(filename), 'w') as f:
+        with morphlib.util.temp_dir(
+                dir=self.app.settings['tempdir']) as tempdir:
+            with open(os.path.join(tempdir, filename), 'w') as f:
                 f.write(data)
             exit_code, output, errors = self.app.runcmd_unchecked(
                 ['autoconf', filename],
                 ['grep', '^PACKAGE_VERSION='],
                 ['cut', '-d=', '-f2'],
                 ['sed', "s/'//g"],
-                cwd=tempdir.root_path)
+                cwd=tempdir)
             version = None
             if output:
                 output = output.strip()
diff --git a/morphlib/plugins/deploy_plugin.py b/morphlib/plugins/deploy_plugin.py
index 1be4a3d1..c0018edd 100644
--- a/morphlib/plugins/deploy_plugin.py
+++ b/morphlib/plugins/deploy_plugin.py
@@ -423,16 +423,13 @@ class DeployPlugin(cliapp.Plugin):
     def deploy_cluster(self, sb, build_command, cluster_morphology,
                        root_repo_dir, repo, commit, env_vars, deployments):
         # Create a tempdir for this deployment to work in
-        deploy_tempdir = tempfile.mkdtemp(
-            dir=os.path.join(self.app.settings['tempdir'], 'deployments'))
-        try:
+        tmp_basedir = os.path.join(self.app.settings['tempdir'], 'deployments')
+        with morphlib.util.temp_dir(dir=tmp_basedir) as deploy_tempdir:
             for system in cluster_morphology['systems']:
                 self.deploy_system(sb, build_command, deploy_tempdir,
                                    root_repo_dir, repo, commit, system,
                                    env_vars, deployments,
                                    parent_location='')
-        finally:
-            shutil.rmtree(deploy_tempdir)
 
     def _sanitise_morphology_paths(self, paths, sb):
         sanitised_paths = []
@@ -696,9 +693,8 @@ class DeployPlugin(cliapp.Plugin):
     def setup_deploy(self, build_command, deploy_tempdir, root_repo_dir, ref,
                      artifact, deployment_type, location, env, components=[]):
         # Create a tempdir to extract the rootfs in
-        system_tree = tempfile.mkdtemp(dir=deploy_tempdir)
-
-        try:
+        with morphlib.util.temp_dir(dir=deploy_tempdir,
+                                    cleanup_on_success=False) as system_tree:
             # FIXME: This should be fixed in morphloader.
             morphlib.util.fix_chunk_build_mode(artifact)
             if self.app.settings['partial']:
@@ -716,18 +712,14 @@ class DeployPlugin(cliapp.Plugin):
                 json.dump(metadata, f, indent=4,
                           sort_keys=True, encoding='unicode-escape')
             return system_tree
-        except Exception:
-            shutil.rmtree(system_tree)
-            raise
 
     def run_deploy_commands(self, deploy_tempdir, env, artifact, root_repo_dir,
                             ref, deployment_type, system_tree, location):
         # Extensions get a private tempdir so we can more easily clean
         # up any files an extension left behind
-        deploy_private_tempdir = tempfile.mkdtemp(dir=deploy_tempdir)
-        env['TMPDIR'] = deploy_private_tempdir
-
-        try:
+        with morphlib.util.temp_dir(dir=deploy_tempdir) \
+             as deploy_private_tempdir:
+            env['TMPDIR'] = deploy_private_tempdir
             # Run configuration extensions.
             if not self.app.settings['partial']:
                 self.app.status(msg='Configure system')
@@ -752,11 +744,6 @@ class DeployPlugin(cliapp.Plugin):
                 [system_tree, location],
                 env)
 
-        finally:
-            # Cleanup.
-            self.app.status(msg='Cleaning up')
-            shutil.rmtree(deploy_private_tempdir)
-
     def _report_extension_stdout(self, line):
         self.app.status(msg=line.replace('%', '%%'))
     def _report_extension_stderr(self, error_list):
diff --git a/morphlib/plugins/system_manifests_plugin.py b/morphlib/plugins/system_manifests_plugin.py
index cc26e5ca..8e14d2eb 100644
--- a/morphlib/plugins/system_manifests_plugin.py
+++ b/morphlib/plugins/system_manifests_plugin.py
@@ -134,34 +134,33 @@ class SystemManifestsPlugin(cliapp.Plugin):
             trove_id = self.app.settings['trove-id'][0]
         except IndexError:
             trove_id = None
-        tempdir = tempfile.mkdtemp(dir=self.app.settings['tempdir'])
-        lorries = get_lorry_repos(tempdir, self.lrc, self.app.status,
-                                  trove_id, self.app.settings['trove-host'])
-        manifest = Manifest(system_artifact.name, tempdir, self.app.status,
-                            self.lrc)
-
-        old_prefix = self.app.status_prefix
-        sources = set(a.source for a in system_artifact.walk()
-                      if a.source.morphology['kind'] == 'chunk'
-                      and a.source.morphology['build-mode'] != 'bootstrap')
-        for i, source in enumerate(sources, start=1):
-            source.cache_key = ckc.compute_key(source)
-            source.cache_id = ckc.get_cache_id(source)
-            name = source.morphology['name']
-            ref = source.original_ref
-
-            # Ensure we have a cache of the repo
-            if not self.lrc.has_repo(source.repo_name):
-                self.lrc.cache_repo(source.repo_name)
-
-            cached = self.lrc.get_repo(source.repo_name)
-
-            new_prefix = '[%d/%d][%s] ' % (i, len(sources), name)
-            self.app.status_prefix = old_prefix + new_prefix
-            manifest.add_chunk(self.app, name, source.repo_name, ref, cached,
-                               resolver, lorries)
-        self.app.status_prefix = old_prefix
-        shutil.rmtree(tempdir)
+        with morphlib.util.temp_dir(dir=self.app.settings['tempdir']) as td:
+            lorries = get_lorry_repos(td, self.lrc, self.app.status, trove_id,
+                                      self.app.settings['trove-host'])
+            manifest = Manifest(system_artifact.name, td, self.app.status,
+                                self.lrc)
+
+            old_prefix = self.app.status_prefix
+            sources = set(a.source for a in system_artifact.walk()
+                          if a.source.morphology['kind'] == 'chunk'
+                          and a.source.morphology['build-mode'] != 'bootstrap')
+            for i, source in enumerate(sources, start=1):
+                source.cache_key = ckc.compute_key(source)
+                source.cache_id = ckc.get_cache_id(source)
+                name = source.morphology['name']
+                ref = source.original_ref
+
+                # Ensure we have a cache of the repo
+                if not self.lrc.has_repo(source.repo_name):
+                    self.lrc.cache_repo(source.repo_name)
+
+                cached = self.lrc.get_repo(source.repo_name)
+
+                new_prefix = '[%d/%d][%s] ' % (i, len(sources), name)
+                self.app.status_prefix = old_prefix + new_prefix
+                manifest.add_chunk(self.app, name, source.repo_name, ref,
+                                   cached, resolver, lorries)
+            self.app.status_prefix = old_prefix
 
 
 def run_licensecheck(filename):