_stream.py: Absorb workspace functionality from App.

3 files changed, 214 insertions, 157 deletions

diff --git a/buildstream/_frontend/app.py b/buildstream/_frontend/app.py
index feac0ab9c..2e53b54de 100644
--- a/buildstream/_frontend/app.py
+++ b/buildstream/_frontend/app.py
@@ -20,7 +20,6 @@
 
 import os
 import sys
-import shutil
 import resource
 import traceback
 import datetime
@@ -32,7 +31,7 @@ import click
 from click import UsageError
 
 # Import buildstream public symbols
-from .. import Scope, Consistency
+from .. import Scope
 
 # Import various buildstream internals
 from .._context import Context
@@ -212,6 +211,9 @@ class App():
         except BstError as e:
             self._error_exit(e, "Error loading project")
 
+        # Create the stream now.
+        self.stream = Stream(self.context)
+
         # Run the body of the session here, once everything is loaded
         try:
             yield
@@ -296,7 +298,11 @@ class App():
             except BstError as e:
                 self._error_exit(e, "Error initializing pipeline")
 
-            self.stream = Stream(self.context, self.scheduler, self.pipeline)
+            # XXX This is going to change soon !
+            #
+            self.stream._scheduler = self.scheduler
+            self.stream._pipeline = self.pipeline
+            self.stream.total_elements = len(list(self.pipeline.dependencies(Scope.ALL)))
 
             # Pipeline is loaded, now we can tell the logger about it
             self.logger.size_request(self.pipeline)
@@ -448,116 +454,6 @@ class App():
             self.pipeline.cleanup()
 
     ############################################################
-    #                   Workspace Commands                     #
-    ############################################################
-
-    # open_workspace
-    #
-    # Open a project workspace - this requires full initialization
-    #
-    # Args:
-    #    target (Element): The element to open the workspace for
-    #    directory (str): The directory to stage the source in
-    #    no_checkout (bool): Whether to skip checking out the source
-    #    track_first (bool): Whether to track and fetch first
-    #    force (bool): Whether to ignore contents in an existing directory
-    #
-    def open_workspace(self, target, directory, no_checkout, track_first, force):
-
-        workdir = os.path.abspath(directory)
-
-        if not list(target.sources()):
-            build_depends = [x.name for x in target.dependencies(Scope.BUILD, recurse=False)]
-            if not build_depends:
-                raise AppError("The given element has no sources")
-            detail = "Try opening a workspace on one of its dependencies instead:\n"
-            detail += "  \n".join(build_depends)
-            raise AppError("The given element has no sources", detail=detail)
-
-        # Check for workspace config
-        if self.project.workspaces.get_workspace(target):
-            raise AppError("Workspace '{}' is already defined.".format(target.name))
-
-        # If we're going to checkout, we need at least a fetch,
-        # if we were asked to track first, we're going to fetch anyway.
-        if not no_checkout or track_first:
-            self.stream.fetch(self.scheduler, [target])
-
-        if not no_checkout and target._get_consistency() != Consistency.CACHED:
-            raise AppError("Could not stage uncached source. " +
-                           "Use `--track` to track and " +
-                           "fetch the latest version of the " +
-                           "source.")
-
-        try:
-            os.makedirs(directory, exist_ok=True)
-        except OSError as e:
-            raise AppError("Failed to create workspace directory: {}".format(e)) from e
-
-        self.project.workspaces.create_workspace(target.name, workdir)
-
-        if not no_checkout:
-            with target.timed_activity("Staging sources to {}".format(directory)):
-                target._open_workspace()
-
-        self.project.workspaces.save_config()
-        self._message(MessageType.INFO, "Saved workspace configuration")
-
-    # close_workspace
-    #
-    # Close a project workspace - this requires only partial initialization
-    #
-    # Args:
-    #    element_name (str): The element name to close the workspace for
-    #    remove_dir (bool): Whether to remove the associated directory
-    #
-    def close_workspace(self, element_name, remove_dir):
-
-        workspace = self.project.workspaces.get_workspace(element_name)
-
-        if workspace is None:
-            raise AppError("Workspace '{}' does not exist".format(element_name))
-
-        if self.interactive and remove_dir:
-            if not click.confirm('This will remove all your changes, are you sure?'):
-                click.echo('Aborting', err=True)
-                sys.exit(-1)
-
-        # Remove workspace directory if prompted
-        if remove_dir:
-            with self.context.timed_activity("Removing workspace directory {}"
-                                             .format(workspace.path)):
-                try:
-                    shutil.rmtree(workspace.path)
-                except OSError as e:
-                    raise AppError("Could not remove  '{}': {}"
-                                   .format(workspace.path, e)) from e
-
-        # Delete the workspace and save the configuration
-        self.project.workspaces.delete_workspace(element_name)
-        self.project.workspaces.save_config()
-        self._message(MessageType.INFO, "Saved workspace configuration")
-
-    # reset_workspace
-    #
-    # Reset a workspace to its original state, discarding any user
-    # changes.
-    #
-    # Args:
-    #    target (Element): The element to reset the workspace for
-    #    track (bool): Whether to also track the source
-    #
-    def reset_workspace(self, target, track):
-        workspace = self.project.workspaces.get_workspace(target.name)
-
-        if workspace is None:
-            raise AppError("Workspace '{}' is currently not defined"
-                           .format(target.name))
-
-        self.close_workspace(target.name, True)
-        self.open_workspace(target, workspace.path, False, track, False)
-
-    ############################################################
     #                      Local Functions                     #
     ############################################################
 
diff --git a/buildstream/_frontend/cli.py b/buildstream/_frontend/cli.py
index 11a0ca2cc..1393db35a 100644
--- a/buildstream/_frontend/cli.py
+++ b/buildstream/_frontend/cli.py
@@ -3,7 +3,7 @@ import sys
 
 import click
 from .. import _yaml
-from .._exceptions import BstError, LoadError
+from .._exceptions import BstError, LoadError, AppError
 from .._versions import BST_FORMAT_VERSION
 from .complete import main_bashcomplete, complete_path, CompleteUnhandled
 
@@ -635,7 +635,7 @@ def workspace_open(app, no_checkout, force, track_, element, directory):
     with app.initialized((element,), rewritable=track_, track_elements=[element] if track_ else None):
         # This command supports only one target
         target = app.pipeline.targets[0]
-        app.open_workspace(target, directory, no_checkout, track_, force)
+        app.stream.workspace_open(target, directory, no_checkout, track_, force)
 
 
 ##################################################################
@@ -657,10 +657,29 @@ def workspace_close(app, remove_dir, all_, elements):
         sys.exit(-1)
 
     with app.partially_initialized():
+
+        # Early exit if we specified `all` and there are no workspaces
+        if all_ and not app.stream.workspace_exists():
+            click.echo('No open workspaces to close', err=True)
+            sys.exit(0)
+
+        # Check that the workspaces in question exist
+        nonexisting = []
+        for element_name in elements:
+            if not app.stream.workspace_exists(element_name):
+                nonexisting.append(element_name)
+        if nonexisting:
+            raise AppError("Workspace does not exist", detail="\n".join(nonexisting))
+
+        if app.interactive and remove_dir:
+            if not click.confirm('This will remove all your changes, are you sure?'):
+                click.echo('Aborting', err=True)
+                sys.exit(-1)
+
         if all_:
             elements = [element_name for element_name, _ in app.project.workspaces.list()]
-        for element in elements:
-            app.close_workspace(element, remove_dir)
+        for element_name in elements:
+            app.stream.workspace_close(element_name, remove_dir)
 
 
 ##################################################################
@@ -692,7 +711,7 @@ def workspace_reset(app, track_, all_, elements):
 
     with app.initialized(elements):
         for target in app.pipeline.targets:
-            app.reset_workspace(target, track_)
+            app.stream.workspace_reset(target, track_)
 
 
 ##################################################################
@@ -704,14 +723,4 @@ def workspace_list(app):
     """List open workspaces"""
 
     with app.partially_initialized():
-        workspaces = []
-        for element_name, workspace_ in app.project.workspaces.list():
-            workspace_detail = {
-                'element': element_name,
-                'directory': workspace_.path,
-            }
-            workspaces.append(workspace_detail)
-
-        _yaml.dump({
-            'workspaces': workspaces
-        })
+        app.stream.workspace_list()
diff --git a/buildstream/_stream.py b/buildstream/_stream.py
index 725c3836f..93a12f630 100644
--- a/buildstream/_stream.py
+++ b/buildstream/_stream.py
@@ -20,14 +20,15 @@
 import os
 import stat
 import shlex
+import shutil
 import tarfile
 from tempfile import TemporaryDirectory
 
 from ._exceptions import StreamError, ImplError, BstError
-from . import _site
-from . import utils
-from . import Scope, Consistency
+from ._message import Message, MessageType
 from ._scheduler import SchedStatus, TrackQueue, FetchQueue, BuildQueue, PullQueue, PushQueue
+from . import utils, _yaml, _site
+from . import Scope, Consistency
 
 
 # Stream()
@@ -39,15 +40,13 @@ from ._scheduler import SchedStatus, TrackQueue, FetchQueue, BuildQueue, PullQue
 #
 class Stream():
 
-    def __init__(self, context, scheduler, pipeline):
+    def __init__(self, context):
         self.session_elements = 0  # Number of elements to process in this session
         self.total_elements = 0    # Number of total potential elements for this pipeline
 
         self._context = context
-        self._scheduler = scheduler
-        self._pipeline = pipeline
-
-        self.total_elements = len(list(self._pipeline.dependencies(Scope.ALL)))
+        self._scheduler = None
+        self._pipeline = None
 
     # track()
     #
@@ -208,7 +207,7 @@ class Stream():
                 with target.timed_activity("Checking out files in {}".format(directory)):
                     try:
                         if hardlinks:
-                            self.checkout_hardlinks(sandbox_root, directory)
+                            self._checkout_hardlinks(sandbox_root, directory)
                         else:
                             utils.copy_files(sandbox_root, directory)
                     except OSError as e:
@@ -217,25 +216,6 @@ class Stream():
             raise StreamError("Error while staging dependencies into a sandbox: {}".format(e),
                               reason=e.reason) from e
 
-    # Helper function for checkout()
-    #
-    def checkout_hardlinks(self, sandbox_root, directory):
-        try:
-            removed = utils.safe_remove(directory)
-        except OSError as e:
-            raise StreamError("Failed to remove checkout directory: {}".format(e)) from e
-
-        if removed:
-            # Try a simple rename of the sandbox root; if that
-            # doesnt cut it, then do the regular link files code path
-            try:
-                os.rename(sandbox_root, directory)
-            except OSError:
-                os.makedirs(directory, exist_ok=True)
-                utils.link_files(sandbox_root, directory)
-        else:
-            utils.link_files(sandbox_root, directory)
-
     # pull()
     #
     # Pulls elements from the pipeline
@@ -290,6 +270,150 @@ class Stream():
         elif status == SchedStatus.TERMINATED:
             raise StreamError(terminated=True)
 
+    # workspace_open
+    #
+    # Open a project workspace
+    #
+    # Args:
+    #    target (Element): The element to open the workspace for
+    #    directory (str): The directory to stage the source in
+    #    no_checkout (bool): Whether to skip checking out the source
+    #    track_first (bool): Whether to track and fetch first
+    #    force (bool): Whether to ignore contents in an existing directory
+    #
+    def workspace_open(self, target, directory, no_checkout, track_first, force):
+        project = self._context.get_toplevel_project()
+        workdir = os.path.abspath(directory)
+
+        if not list(target.sources()):
+            build_depends = [x.name for x in target.dependencies(Scope.BUILD, recurse=False)]
+            if not build_depends:
+                raise StreamError("The given element has no sources")
+            detail = "Try opening a workspace on one of its dependencies instead:\n"
+            detail += "  \n".join(build_depends)
+            raise StreamError("The given element has no sources", detail=detail)
+
+        # Check for workspace config
+        workspace = project.workspaces.get_workspace(target.name)
+        if workspace:
+            raise StreamError("Workspace '{}' is already defined at: {}"
+                              .format(target.name, workspace.path))
+
+        # If we're going to checkout, we need at least a fetch,
+        # if we were asked to track first, we're going to fetch anyway.
+        if not no_checkout or track_first:
+            self.fetch(self._scheduler, [target])
+
+        if not no_checkout and target._get_consistency() != Consistency.CACHED:
+            raise StreamError("Could not stage uncached source. " +
+                              "Use `--track` to track and " +
+                              "fetch the latest version of the " +
+                              "source.")
+
+        try:
+            os.makedirs(directory, exist_ok=True)
+        except OSError as e:
+            raise StreamError("Failed to create workspace directory: {}".format(e)) from e
+
+        project.workspaces.create_workspace(target.name, workdir)
+
+        if not no_checkout:
+            with target.timed_activity("Staging sources to {}".format(directory)):
+                target._open_workspace()
+
+        project.workspaces.save_config()
+        self._message(MessageType.INFO, "Saved workspace configuration")
+
+    # workspace_close
+    #
+    # Close a project workspace
+    #
+    # Args:
+    #    element_name (str): The element name to close the workspace for
+    #    remove_dir (bool): Whether to remove the associated directory
+    #
+    def workspace_close(self, element_name, remove_dir):
+        project = self._context.get_toplevel_project()
+        workspace = project.workspaces.get_workspace(element_name)
+
+        # Remove workspace directory if prompted
+        if remove_dir:
+            with self._context.timed_activity("Removing workspace directory {}"
+                                              .format(workspace.path)):
+                try:
+                    shutil.rmtree(workspace.path)
+                except OSError as e:
+                    raise StreamError("Could not remove  '{}': {}"
+                                      .format(workspace.path, e)) from e
+
+        # Delete the workspace and save the configuration
+        project.workspaces.delete_workspace(element_name)
+        project.workspaces.save_config()
+        self._message(MessageType.INFO, "Closed workspace for {}".format(element_name))
+
+    # workspace_reset
+    #
+    # Reset a workspace to its original state, discarding any user
+    # changes.
+    #
+    # Args:
+    #    target (Element): The element to reset the workspace for
+    #    track (bool): Whether to also track the source
+    #
+    def workspace_reset(self, target, track):
+        project = self._context.get_toplevel_project()
+        workspace = project.workspaces.get_workspace(target.name)
+
+        if workspace is None:
+            raise StreamError("Workspace '{}' is currently not defined"
+                              .format(target.name))
+
+        self.workspace_close(target.name, True)
+        self.workspace_open(target, workspace.path, False, track, False)
+
+    # workspace_exists
+    #
+    # Check if a workspace exists
+    #
+    # Args:
+    #    element_name (str): The element name to close the workspace for, or None
+    #
+    # Returns:
+    #    (bool): True if the workspace exists
+    #
+    # If None is specified for `element_name`, then this will return
+    # True if there are any existing workspaces.
+    #
+    def workspace_exists(self, element_name=None):
+        project = self._context.get_toplevel_project()
+
+        if element_name:
+            workspace = project.workspaces.get_workspace(element_name)
+            if workspace:
+                return True
+        elif any(project.workspaces.list()):
+            return True
+
+        return False
+
+    # workspace_list
+    #
+    # Serializes the workspaces and dumps them in YAML to stdout.
+    #
+    def workspace_list(self):
+        project = self._context.get_toplevel_project()
+        workspaces = []
+        for element_name, workspace_ in project.workspaces.list():
+            workspace_detail = {
+                'element': element_name,
+                'directory': workspace_.path,
+            }
+            workspaces.append(workspace_detail)
+
+        _yaml.dump({
+            'workspaces': workspaces
+        })
+
     # source_bundle()
     #
     # Create a build bundle for the given artifact.
@@ -351,6 +475,34 @@ class Stream():
     #                     Private Methods                       #
     #############################################################
 
+    # _message()
+    #
+    # Local message propagator
+    #
+    def _message(self, message_type, message, **kwargs):
+        args = dict(kwargs)
+        self._context.message(
+            Message(None, message_type, message, **args))
+
+    # Helper function for checkout()
+    #
+    def _checkout_hardlinks(self, sandbox_root, directory):
+        try:
+            removed = utils.safe_remove(directory)
+        except OSError as e:
+            raise StreamError("Failed to remove checkout directory: {}".format(e)) from e
+
+        if removed:
+            # Try a simple rename of the sandbox root; if that
+            # doesnt cut it, then do the regular link files code path
+            try:
+                os.rename(sandbox_root, directory)
+            except OSError:
+                os.makedirs(directory, exist_ok=True)
+                utils.link_files(sandbox_root, directory)
+        else:
+            utils.link_files(sandbox_root, directory)
+
     # Write the element build script to the given directory
     def _write_element_script(self, directory, element):
         try: