stage : on-going refactoring

3 files changed, 180 insertions, 185 deletions

diff --git a/lib/spack/spack/package.py b/lib/spack/spack/package.py
index 8019b29cba..4fdc582479 100644
--- a/lib/spack/spack/package.py
+++ b/lib/spack/spack/package.py
@@ -847,93 +847,94 @@ class Package(object):
                 make_jobs=make_jobs)
 
         start_time = time.time()
-        if not fake:
-            if not skip_patch:
-                self.do_patch()
-            else:
-                self.do_stage()
-
-        # create the install directory.  The install layout
-        # handles this in case so that it can use whatever
-        # package naming scheme it likes.
-        spack.install_layout.create_install_directory(self.spec)
-
-        def cleanup():
-            if not keep_prefix:
-                # If anything goes wrong, remove the install prefix
-                self.remove_prefix()
-            else:
-                tty.warn("Keeping install prefix in place despite error.",
-                         "Spack will think this package is installed." +
-                         "Manually remove this directory to fix:",
-                         self.prefix, wrap=True)
-
-
-        def real_work():
-            try:
-                tty.msg("Building %s." % self.name)
-
-                # Run the pre-install hook in the child process after
-                # the directory is created.
-                spack.hooks.pre_install(self)
-
-                # Set up process's build environment before running install.
-                if fake:
-                    self.do_fake_install()
+        with self.stage:
+            if not fake:
+                if not skip_patch:
+                    self.do_patch()
                 else:
-                    # Do the real install in the source directory.
-                    self.stage.chdir_to_source()
-
-                    # This redirects I/O to a build log (and optionally to the terminal)
-                    log_path = join_path(os.getcwd(), 'spack-build.out')
-                    log_file = open(log_path, 'w')
-                    with log_output(log_file, verbose, sys.stdout.isatty(), True):
-                        self.install(self.spec, self.prefix)
-
-                # Ensure that something was actually installed.
-                self._sanity_check_install()
-
-                # Move build log into install directory on success
-                if not fake:
-                    log_install_path = spack.install_layout.build_log_path(self.spec)
-                    install(log_path, log_install_path)
-
-                # On successful install, remove the stage.
-                if not keep_stage:
-                    self.stage.destroy()
-
-                # Stop timer.
-                self._total_time = time.time() - start_time
-                build_time = self._total_time - self._fetch_time
-
-                tty.msg("Successfully installed %s." % self.name,
-                        "Fetch: %s.  Build: %s.  Total: %s."
-                        % (_hms(self._fetch_time), _hms(build_time), _hms(self._total_time)))
-                print_pkg(self.prefix)
-
-            except ProcessError, e:
-                # Annotate with location of build log.
-                e.build_log = log_path
-                cleanup()
-                raise e
-
-            except:
-                # other exceptions just clean up and raise.
-                cleanup()
-                raise
+                    self.do_stage()
 
-        # Set parallelism before starting build.
-        self.make_jobs = make_jobs
+            # create the install directory.  The install layout
+            # handles this in case so that it can use whatever
+            # package naming scheme it likes.
+            spack.install_layout.create_install_directory(self.spec)
 
-        # Do the build.
-        spack.build_environment.fork(self, real_work)
-
-        # note: PARENT of the build process adds the new package to
-        # the database, so that we don't need to re-read from file.
-        spack.installed_db.add(self.spec, self.prefix)
-
-        # Once everything else is done, run post install hooks
-        spack.hooks.post_install(self)
+            def cleanup():
+                if not keep_prefix:
+                    # If anything goes wrong, remove the install prefix
+                    self.remove_prefix()
+                else:
+                    tty.warn("Keeping install prefix in place despite error.",
+                             "Spack will think this package is installed." +
+                             "Manually remove this directory to fix:",
+                             self.prefix, wrap=True)
+
+
+            def real_work():
+                try:
+                    tty.msg("Building %s." % self.name)
+
+                    # Run the pre-install hook in the child process after
+                    # the directory is created.
+                    spack.hooks.pre_install(self)
+
+                    # Set up process's build environment before running install.
+                    if fake:
+                        self.do_fake_install()
+                    else:
+                        # Do the real install in the source directory.
+                        self.stage.chdir_to_source()
+
+                        # This redirects I/O to a build log (and optionally to the terminal)
+                        log_path = join_path(os.getcwd(), 'spack-build.out')
+                        log_file = open(log_path, 'w')
+                        with log_output(log_file, verbose, sys.stdout.isatty(), True):
+                            self.install(self.spec, self.prefix)
+
+                    # Ensure that something was actually installed.
+                    self._sanity_check_install()
+
+                    # Move build log into install directory on success
+                    if not fake:
+                        log_install_path = spack.install_layout.build_log_path(self.spec)
+                        install(log_path, log_install_path)
+
+                    # On successful install, remove the stage.
+                    if not keep_stage:
+                        self.stage.destroy()
+
+                    # Stop timer.
+                    self._total_time = time.time() - start_time
+                    build_time = self._total_time - self._fetch_time
+
+                    tty.msg("Successfully installed %s." % self.name,
+                            "Fetch: %s.  Build: %s.  Total: %s."
+                            % (_hms(self._fetch_time), _hms(build_time), _hms(self._total_time)))
+                    print_pkg(self.prefix)
+
+                except ProcessError, e:
+                    # Annotate with location of build log.
+                    e.build_log = log_path
+                    cleanup()
+                    raise e
+
+                except:
+                    # other exceptions just clean up and raise.
+                    cleanup()
+                    raise
+
+            # Set parallelism before starting build.
+            self.make_jobs = make_jobs
+
+            # Do the build.
+            spack.build_environment.fork(self, real_work)
+
+            # note: PARENT of the build process adds the new package to
+            # the database, so that we don't need to re-read from file.
+            spack.installed_db.add(self.spec, self.prefix)
+
+            # Once everything else is done, run post install hooks
+            spack.hooks.post_install(self)
 
 
     def _sanity_check_install(self):
diff --git a/lib/spack/spack/stage.py b/lib/spack/spack/stage.py
index f217450d42..96b1eaf3f2 100644
--- a/lib/spack/spack/stage.py
+++ b/lib/spack/spack/stage.py
@@ -96,12 +96,44 @@ class Stage(object):
         self.default_fetcher = self.fetcher  # self.fetcher can change with mirrors.
         self.skip_checksum_for_mirror = True  # used for mirrored archives of repositories.
 
-        self.name = kwargs.get('name')
+        # TODO : this uses a protected member of tempfile, but seemed the only way to get a temporary name
+        # TODO : besides, the temporary link name won't be the same as the temporary stage area in tmp_root
+        self.name = kwargs.get('name') if 'name' in kwargs else STAGE_PREFIX + next(tempfile._get_candidate_names())
         self.mirror_path = kwargs.get('mirror_path')
         self.tmp_root = find_tmp_root()
 
-        self.path = None
+        # Try to construct here a temporary name for the stage directory
+        # If this is a named stage, then construct a named path.
+        self.path = join_path(spack.stage_path, self.name)
+
+        self.delete_on_exit = True
+
+    def __enter__(self):
+        """
+        Entering a stage context will create the stage directory
+        """
+        # FIXME : if _setup is used only here, then it makes no sense to retain the function
         self._setup()
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """
+        Exiting from a stage context will delete the stage directory unless:
+        - it was explicitly requested not to do so
+        - an exception has been raised
+
+        Args:
+            exc_type: exception type
+            exc_val: exception value
+            exc_tb: exception traceback
+
+        Returns:
+            Boolean
+        """
+        self.delete_on_exit = False if exc_type is not None else self.delete_on_exit
+
+        if self.delete_on_exit:
+            self.destroy()
 
     def _cleanup_dead_links(self):
         """Remove any dead links in the stage directory."""
@@ -163,35 +195,17 @@ class Stage(object):
         mkdirp(spack.stage_path)
         self._cleanup_dead_links()
 
-        # If this is a named stage, then construct a named path.
-        if self.name is not None:
-            self.path = join_path(spack.stage_path, self.name)
-
         # If this is a temporary stage, them make the temp directory
-        tmp_dir = None
         if self.tmp_root:
-            if self.name is None:
-                # Unnamed tmp root.  Link the path in
+            if self._need_to_create_path():
                 tmp_dir = tempfile.mkdtemp('', STAGE_PREFIX, self.tmp_root)
-                self.name = os.path.basename(tmp_dir)
-                self.path = join_path(spack.stage_path, self.name)
-                if self._need_to_create_path():
-                    os.symlink(tmp_dir, self.path)
-
-            else:
-                if self._need_to_create_path():
-                    tmp_dir = tempfile.mkdtemp('', STAGE_PREFIX, self.tmp_root)
-                    os.symlink(tmp_dir, self.path)
+                os.symlink(tmp_dir, self.path)
 
         # if we're not using a tmp dir, create the stage directly in the
         # stage dir, rather than linking to it.
         else:
-            if self.name is None:
-                self.path = tempfile.mkdtemp('', STAGE_PREFIX, spack.stage_path)
-                self.name = os.path.basename(self.path)
-            else:
-                if self._need_to_create_path():
-                    mkdirp(self.path)
+            if self._need_to_create_path():
+                mkdirp(self.path)
 
         # Make sure we can actually do something with the stage we made.
         ensure_access(self.path)
@@ -389,6 +403,12 @@ class StageComposite:
     def path(self):
         return self[0].path
 
+    def __enter__(self):
+        return self[0].__enter__()
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        return self[0].__exit__(exc_type, exc_val, exc_tb)
+
     def chdir_to_source(self):
         return self[0].chdir_to_source()
 
diff --git a/lib/spack/spack/test/stage.py b/lib/spack/spack/test/stage.py
index c1b2a2a573..a7314eba4c 100644
--- a/lib/spack/spack/test/stage.py
+++ b/lib/spack/spack/test/stage.py
@@ -192,116 +192,90 @@ class StageTest(unittest.TestCase):
 
     def test_setup_and_destroy_name_with_tmp(self):
         with use_tmp(True):
-            stage = Stage(archive_url, name=stage_name)
-            self.check_setup(stage, stage_name)
-
-            stage.destroy()
+            with Stage(archive_url, name=stage_name) as stage:
+                self.check_setup(stage, stage_name)
             self.check_destroy(stage, stage_name)
 
 
     def test_setup_and_destroy_name_without_tmp(self):
         with use_tmp(False):
-            stage = Stage(archive_url, name=stage_name)
-            self.check_setup(stage, stage_name)
-
-            stage.destroy()
+            with Stage(archive_url, name=stage_name) as stage:
+                self.check_setup(stage, stage_name)
             self.check_destroy(stage, stage_name)
 
 
     def test_setup_and_destroy_no_name_with_tmp(self):
         with use_tmp(True):
-            stage = Stage(archive_url)
-            self.check_setup(stage, None)
-
-            stage.destroy()
+            with Stage(archive_url) as stage:
+                self.check_setup(stage, None)
             self.check_destroy(stage, None)
 
 
     def test_setup_and_destroy_no_name_without_tmp(self):
         with use_tmp(False):
-            stage = Stage(archive_url)
-            self.check_setup(stage, None)
-
-            stage.destroy()
+            with Stage(archive_url) as stage:
+                self.check_setup(stage, None)
             self.check_destroy(stage, None)
 
 
     def test_chdir(self):
-        stage = Stage(archive_url, name=stage_name)
-
-        stage.chdir()
-        self.check_setup(stage, stage_name)
-        self.check_chdir(stage, stage_name)
-
-        stage.destroy()
+        with Stage(archive_url, name=stage_name) as stage:
+            stage.chdir()
+            self.check_setup(stage, stage_name)
+            self.check_chdir(stage, stage_name)
         self.check_destroy(stage, stage_name)
 
 
     def test_fetch(self):
-        stage = Stage(archive_url, name=stage_name)
-
-        stage.fetch()
-        self.check_setup(stage, stage_name)
-        self.check_chdir(stage, stage_name)
-        self.check_fetch(stage, stage_name)
-
-        stage.destroy()
+        with Stage(archive_url, name=stage_name) as stage:
+            stage.fetch()
+            self.check_setup(stage, stage_name)
+            self.check_chdir(stage, stage_name)
+            self.check_fetch(stage, stage_name)
         self.check_destroy(stage, stage_name)
 
 
     def test_expand_archive(self):
-        stage = Stage(archive_url, name=stage_name)
-
-        stage.fetch()
-        self.check_setup(stage, stage_name)
-        self.check_fetch(stage, stage_name)
-
-        stage.expand_archive()
-        self.check_expand_archive(stage, stage_name)
-
-        stage.destroy()
+        with Stage(archive_url, name=stage_name) as stage:
+            stage.fetch()
+            self.check_setup(stage, stage_name)
+            self.check_fetch(stage, stage_name)
+            stage.expand_archive()
+            self.check_expand_archive(stage, stage_name)
         self.check_destroy(stage, stage_name)
 
 
     def test_expand_archive(self):
-        stage = Stage(archive_url, name=stage_name)
-
-        stage.fetch()
-        self.check_setup(stage, stage_name)
-        self.check_fetch(stage, stage_name)
-
-        stage.expand_archive()
-        stage.chdir_to_source()
-        self.check_expand_archive(stage, stage_name)
-        self.check_chdir_to_source(stage, stage_name)
-
-        stage.destroy()
+        with Stage(archive_url, name=stage_name) as stage:
+            stage.fetch()
+            self.check_setup(stage, stage_name)
+            self.check_fetch(stage, stage_name)
+            stage.expand_archive()
+            stage.chdir_to_source()
+            self.check_expand_archive(stage, stage_name)
+            self.check_chdir_to_source(stage, stage_name)
         self.check_destroy(stage, stage_name)
 
 
     def test_restage(self):
-        stage = Stage(archive_url, name=stage_name)
-
-        stage.fetch()
-        stage.expand_archive()
-        stage.chdir_to_source()
-        self.check_expand_archive(stage, stage_name)
-        self.check_chdir_to_source(stage, stage_name)
-
-        # Try to make a file in the old archive dir
-        with open('foobar', 'w') as file:
-            file.write("this file is to be destroyed.")
-
-        self.assertTrue('foobar' in os.listdir(stage.source_path))
-
-        # Make sure the file is not there after restage.
-        stage.restage()
-        self.check_chdir(stage, stage_name)
-        self.check_fetch(stage, stage_name)
-
-        stage.chdir_to_source()
-        self.check_chdir_to_source(stage, stage_name)
-        self.assertFalse('foobar' in os.listdir(stage.source_path))
-
-        stage.destroy()
+        with Stage(archive_url, name=stage_name) as stage:
+            stage.fetch()
+            stage.expand_archive()
+            stage.chdir_to_source()
+            self.check_expand_archive(stage, stage_name)
+            self.check_chdir_to_source(stage, stage_name)
+
+            # Try to make a file in the old archive dir
+            with open('foobar', 'w') as file:
+                file.write("this file is to be destroyed.")
+
+            self.assertTrue('foobar' in os.listdir(stage.source_path))
+
+            # Make sure the file is not there after restage.
+            stage.restage()
+            self.check_chdir(stage, stage_name)
+            self.check_fetch(stage, stage_name)
+            stage.chdir_to_source()
+            self.check_chdir_to_source(stage, stage_name)
+            self.assertFalse('foobar' in os.listdir(stage.source_path))
         self.check_destroy(stage, stage_name)