Small refactor: add keep parameter to stage, get rid of stage.destroy call.

- package.py uses context manager more effectively.

- Stage.__init__ has easier to understand method signature now.
  - keep can be used to override the default behavior either to keep
    the stage ALL the time or to delete the stage ALL the time.

4 files changed, 144 insertions, 93 deletions

diff --git a/lib/spack/llnl/util/filesystem.py b/lib/spack/llnl/util/filesystem.py
index 366237ef8f..a92cb0706d 100644
--- a/lib/spack/llnl/util/filesystem.py
+++ b/lib/spack/llnl/util/filesystem.py
@@ -362,8 +362,9 @@ def remove_dead_links(root):
 
 def remove_linked_tree(path):
     """
-    Removes a directory and its contents.  If the directory is a symlink, follows the link and removes the real
-    directory before removing the link.
+    Removes a directory and its contents.  If the directory is a
+    symlink, follows the link and removes the real directory before
+    removing the link.
 
     Args:
         path: directory to be removed
diff --git a/lib/spack/spack/cmd/clean.py b/lib/spack/spack/cmd/clean.py
index 0c8bd1d528..6e7179122c 100644
--- a/lib/spack/spack/cmd/clean.py
+++ b/lib/spack/spack/cmd/clean.py
@@ -43,4 +43,4 @@ def clean(parser, args):
     specs = spack.cmd.parse_specs(args.packages, concretize=True)
     for spec in specs:
         package = spack.repo.get(spec)
-        package.stage.destroy()
+        package.do_clean()
diff --git a/lib/spack/spack/package.py b/lib/spack/spack/package.py
index be45415b75..47d259968a 100644
--- a/lib/spack/spack/package.py
+++ b/lib/spack/spack/package.py
@@ -293,6 +293,7 @@ class Package(object):
 
     .. code-block:: python
 
+       p.do_clean()              # removes the stage directory entirely
        p.do_restage()            # removes the build directory and
                                  # re-expands the archive.
 
@@ -455,7 +456,7 @@ class Package(object):
         # Construct a composite stage on top of the composite FetchStrategy
         composite_fetcher = self.fetcher
         composite_stage = StageComposite()
-        resources = self._get_resources()
+        resources = self._get_needed_resources()
         for ii, fetcher in enumerate(composite_fetcher):
             if ii == 0:
                 # Construct root stage first
@@ -484,12 +485,14 @@ class Package(object):
 
 
     def _make_fetcher(self):
-        # Construct a composite fetcher that always contains at least one element (the root package). In case there
-        # are resources associated with the package, append their fetcher to the composite.
+        # Construct a composite fetcher that always contains at least
+        # one element (the root package). In case there are resources
+        # associated with the package, append their fetcher to the
+        # composite.
         root_fetcher = fs.for_package_version(self, self.version)
         fetcher = fs.FetchStrategyComposite()  # Composite fetcher
         fetcher.append(root_fetcher)  # Root fetcher is always present
-        resources = self._get_resources()
+        resources = self._get_needed_resources()
         for resource in resources:
             fetcher.append(resource.fetcher)
         return fetcher
@@ -706,6 +709,7 @@ class Package(object):
         self.stage.expand_archive()
         self.stage.chdir_to_source()
 
+
     def do_patch(self):
         """Calls do_stage(), then applied patches to the expanded tarball if they
            haven't been applied already."""
@@ -798,7 +802,7 @@ class Package(object):
         mkdirp(self.prefix.man1)
 
 
-    def _get_resources(self):
+    def _get_needed_resources(self):
         resources = []
         # Select the resources that are needed for this build
         for when_spec, resource_list in self.resources.items():
@@ -816,7 +820,7 @@ class Package(object):
 
 
     def do_install(self,
-                   keep_prefix=False,  keep_stage=False, ignore_deps=False,
+                   keep_prefix=False,  keep_stage=None, ignore_deps=False,
                    skip_patch=False, verbose=False, make_jobs=None, fake=False):
         """Called by commands to install a package and its dependencies.
 
@@ -825,7 +829,8 @@ class Package(object):
 
         Args:
         keep_prefix -- Keep install prefix on failure. By default, destroys it.
-        keep_stage  -- Keep stage on successful build. By default, destroys it.
+        keep_stage  -- Set to True or false to always keep or always delete stage.
+                       By default, stage is destroyed only if there are no exceptions.
         ignore_deps -- Do not install dependencies before installing this package.
         fake        -- Don't really build -- install fake stub files instead.
         skip_patch  -- Skip patch stage of build if True.
@@ -848,32 +853,33 @@ class Package(object):
                 make_jobs=make_jobs)
 
         start_time = time.time()
-        with self.stage:
-            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)
+        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)
 
+                self.stage.keep = keep_stage
+                with self.stage:
                     # Run the pre-install hook in the child process after
                     # the directory is created.
                     spack.hooks.pre_install(self)
@@ -888,7 +894,7 @@ class Package(object):
                         # Save the build environment in a file before building.
                         env_path = join_path(os.getcwd(), 'spack-build.env')
 
-                        # This redirects I/O to a build log (and optionally to the terminal)
+                        # Redirect 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):
@@ -908,29 +914,25 @@ class Package(object):
                     packages_dir = spack.install_layout.build_packages_path(self.spec)
                     dump_packages(self.spec, packages_dir)
 
-                    # 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
 
-                    # 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)
 
-                    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 as e:
+                # Annotate with location of build log.
+                e.build_log = log_path
+                cleanup()
+                raise e
 
-                except ProcessError as 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
+            except:
+                # other exceptions just clean up and raise.
+                cleanup()
+                raise
 
             # Set parallelism before starting build.
             self.make_jobs = make_jobs
@@ -1147,6 +1149,12 @@ class Package(object):
         """Reverts expanded/checked out source to a pristine state."""
         self.stage.restage()
 
+
+    def do_clean(self):
+        """Removes the package's build stage and source tarball."""
+        self.stage.destroy()
+
+
     def format_doc(self, **kwargs):
         """Wrap doc string at 72 characters and format nicely"""
         indent = kwargs.get('indent', 0)
diff --git a/lib/spack/spack/stage.py b/lib/spack/spack/stage.py
index a22982a6d4..b117c76aa1 100644
--- a/lib/spack/spack/stage.py
+++ b/lib/spack/spack/stage.py
@@ -42,29 +42,53 @@ STAGE_PREFIX = 'spack-stage-'
 
 
 class Stage(object):
-    """
-    A Stage object is a context manager that handles a directory where some source code is downloaded and built
-    before being installed.  It handles fetching the source code, either as an archive to be expanded or by checking
-    it out of a repository.  A stage's lifecycle looks like this:
+    """Manages a temporary stage directory for building.
+
+    A Stage object is a context manager that handles a directory where
+    some source code is downloaded and built before being installed.
+    It handles fetching the source code, either as an archive to be
+    expanded or by checking it out of a repository.  A stage's
+    lifecycle looks like this:
+
+    ```
+    with Stage() as stage:      # Context manager creates and destroys the stage directory
+        stage.fetch()           # Fetch a source archive into the stage.
+        stage.expand_archive()  # Expand the source archive.
+        <install>               # Build and install the archive. (handled by user of Stage)
+    ```
+
+    When used as a context manager, the stage is automatically
+    destroyed if no exception is raised by the context. If an
+    excpetion is raised, the stage is left in the filesystem and NOT
+    destroyed, for potential reuse later.
+
+    You can also use the stage's create/destroy functions manually,
+    like this:
 
     ```
-    with Stage() as stage:  # Context manager creates and destroys the stage directory
-        fetch()  # Fetch a source archive into the stage.
-        expand_archive()  # Expand the source archive.
-        <install>  # Build and install the archive.  This is handled by the Package class.
+    stage = Stage()
+    try:
+        stage.create()          # Explicitly create the stage directory.
+        stage.fetch()           # Fetch a source archive into the stage.
+        stage.expand_archive()  # Expand the source archive.
+        <install>               # Build and install the archive. (handled by user of Stage)
+    finally:
+        stage.destroy()         # Explicitly destroy the stage directory.
     ```
 
-    If spack.use_tmp_stage is True, spack will attempt to create stages in a tmp directory.
-    Otherwise, stages are created directly in spack.stage_path.
+    If spack.use_tmp_stage is True, spack will attempt to create
+    stages in a tmp directory.  Otherwise, stages are created directly
+    in spack.stage_path.
 
-    There are two kinds of stages: named and unnamed.  Named stages can persist between runs of spack, e.g. if you
-    fetched a tarball but didn't finish building it, you won't have to fetch it again.
+    There are two kinds of stages: named and unnamed.  Named stages
+    can persist between runs of spack, e.g. if you fetched a tarball
+    but didn't finish building it, you won't have to fetch it again.
 
-    Unnamed stages are created using standard mkdtemp mechanisms or similar, and are intended to persist for
-    only one run of spack.
+    Unnamed stages are created using standard mkdtemp mechanisms or
+    similar, and are intended to persist for only one run of spack.
     """
 
-    def __init__(self, url_or_fetch_strategy, **kwargs):
+    def __init__(self, url_or_fetch_strategy, name=None, mirror_path=None, keep=None):
         """Create a stage object.
            Parameters:
              url_or_fetch_strategy
@@ -76,6 +100,18 @@ class Stage(object):
                  and will persist between runs (or if you construct another
                  stage object later).  If name is not provided, then this
                  stage will be given a unique name automatically.
+
+             mirror_path
+                 If provided, Stage will search Spack's mirrors for
+                 this archive at the mirror_path, before using the
+                 default fetch strategy.
+
+             keep
+                 By default, when used as a context manager, the Stage
+                 is cleaned up when everything goes well, and it is
+                 kept intact when an exception is raised. You can
+                 override this behavior by setting keep to True
+                 (always keep) or False (always delete).
         """
         # TODO: fetch/stage coupling needs to be reworked -- the logic
         # TODO: here is convoluted and not modular enough.
@@ -91,15 +127,19 @@ class Stage(object):
 
         # 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.name = name
+        if name is None:
+            self.name = STAGE_PREFIX + next(tempfile._get_candidate_names())
+        self.mirror_path = mirror_path
         self.tmp_root = find_tmp_root()
 
         # 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)
+
         # Flag to decide whether to delete the stage folder on exit or not
-        self.delete_on_exit = True
+        self.keep = keep
+
 
     def __enter__(self):
         """
@@ -111,6 +151,7 @@ class Stage(object):
         self.create()
         return self
 
+
     def __exit__(self, exc_type, exc_val, exc_tb):
         """
         Exiting from a stage context will delete the stage directory unless:
@@ -125,11 +166,15 @@ class Stage(object):
         Returns:
             Boolean
         """
-        self.delete_on_exit = False if exc_type is not None else self.delete_on_exit
+        if self.keep is None:
+            # Default: delete when there are no exceptions.
+            if exc_type is None: self.destroy()
 
-        if self.delete_on_exit:
+        elif not self.keep:
+            # Overridden. Either always keep or always delete.
             self.destroy()
 
+
     def _need_to_create_path(self):
         """Makes sure nothing weird has happened since the last time we
            looked at path.  Returns True if path already exists and is ok.
@@ -201,7 +246,7 @@ class Stage(object):
         if os.path.isdir(self.path):
             os.chdir(self.path)
         else:
-            tty.die("Setup failed: no such directory: " + self.path)
+            raise ChdirError("Setup failed: no such directory: " + self.path)
 
     def fetch(self, mirror_only=False):
         """Downloads an archive or checks out code from a repository."""
@@ -302,11 +347,14 @@ class Stage(object):
         """
         Creates the stage directory
 
-        If self.tmp_root evaluates to False, the stage directory is created directly under spack.stage_path, otherwise
-        this will attempt to create a stage in a temporary directory and link it into spack.stage_path.
+        If self.tmp_root evaluates to False, the stage directory is
+        created directly under spack.stage_path, otherwise this will
+        attempt to create a stage in a temporary directory and link it
+        into spack.stage_path.
 
-        Spack will use the first writable location in spack.tmp_dirs to create a stage. If there is no valid location
-        in tmp_dirs, fall back to making the stage inside spack.stage_path.
+        Spack will use the first writable location in spack.tmp_dirs
+        to create a stage. If there is no valid location in tmp_dirs,
+        fall back to making the stage inside spack.stage_path.
         """
         # Create the top-level stage directory
         mkdirp(spack.stage_path)
@@ -323,9 +371,7 @@ class Stage(object):
         ensure_access(self.path)
 
     def destroy(self):
-        """
-        Removes this stage directory
-        """
+        """Removes this stage directory."""
         remove_linked_tree(self.path)
 
         # Make sure we don't end up in a removed directory
@@ -370,7 +416,7 @@ class ResourceStage(Stage):
                 shutil.move(source_path, destination_path)
 
 
-@pattern.composite(method_list=['fetch', 'check', 'expand_archive', 'restage', 'destroy'])
+@pattern.composite(method_list=['fetch', 'create', 'check', 'expand_archive', 'restage', 'destroy'])
 class StageComposite:
     """
     Composite for Stage type objects. The first item in this composite is considered to be the root package, and
@@ -410,7 +456,7 @@ class DIYStage(object):
         if os.path.isdir(self.path):
             os.chdir(self.path)
         else:
-            tty.die("Setup failed: no such directory: " + self.path)
+            raise ChdirError("Setup failed: no such directory: " + self.path)
 
     def chdir_to_source(self):
         self.chdir()
@@ -472,19 +518,15 @@ def find_tmp_root():
 
 
 class StageError(spack.error.SpackError):
-    def __init__(self, message, long_message=None):
-        super(self, StageError).__init__(message, long_message)
+    """"Superclass for all errors encountered during staging."""
 
 
 class RestageError(StageError):
-    def __init__(self, message, long_msg=None):
-        super(RestageError, self).__init__(message, long_msg)
+    """"Error encountered during restaging."""
 
 
 class ChdirError(StageError):
-    def __init__(self, message, long_msg=None):
-        super(ChdirError, self).__init__(message, long_msg)
-
+    """Raised when Spack can't change directories."""
 
 # Keep this in namespace for convenience
 FailedDownloadError = fs.FailedDownloadError