Redirect elements when performing tracking and workspace commands

i.e. when an element is specified, it may be replaced with its source element.

3 files changed, 54 insertions, 12 deletions

diff --git a/buildstream/_frontend/cli.py b/buildstream/_frontend/cli.py
index c321fa99f..41e97cb0e 100644
--- a/buildstream/_frontend/cli.py
+++ b/buildstream/_frontend/cli.py
@@ -625,6 +625,11 @@ def workspace_close(app, remove_dir, all_, elements):
             click.echo('No open workspaces to close', err=True)
             sys.exit(0)
 
+        if all_:
+            elements = [element_name for element_name, _ in app.project.workspaces.list()]
+
+        elements = app.stream.redirect_element_names(elements)
+
         # Check that the workspaces in question exist
         nonexisting = []
         for element_name in elements:
@@ -638,8 +643,6 @@ def workspace_close(app, remove_dir, all_, elements):
                 click.echo('Aborting', err=True)
                 sys.exit(-1)
 
-        if all_:
-            elements = [element_name for element_name, _ in app.project.workspaces.list()]
         for element_name in elements:
             app.stream.workspace_close(element_name, remove_dir=remove_dir)
 
@@ -669,13 +672,6 @@ def workspace_reset(app, soft, track_, all_, elements):
         if all_ and not app.stream.workspace_exists():
             raise AppError("No open workspaces to reset")
 
-        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 not soft:
             if not click.confirm('This will remove all your changes, are you sure?'):
                 click.echo('Aborting', err=True)
diff --git a/buildstream/_pipeline.py b/buildstream/_pipeline.py
index 04979bc7c..ba27ca6b6 100644
--- a/buildstream/_pipeline.py
+++ b/buildstream/_pipeline.py
@@ -212,11 +212,19 @@ class Pipeline():
     # use in the result, this function reports a list that is appropriate for
     # the selected option.
     #
-    def get_selection(self, targets, mode):
+    def get_selection(self, targets, mode, *, silent=True):
 
         elements = None
         if mode == PipelineSelection.NONE:
-            elements = targets
+            # Redirect and log if permitted
+            elements = []
+            for t in targets:
+                new_elm = t._get_source_element()
+                if new_elm != t and not silent:
+                    self._message(MessageType.INFO, "Element '{}' redirected to '{}'"
+                                  .format(t.name, new_elm.name))
+                if new_elm not in elements:
+                    elements.append(new_elm)
         elif mode == PipelineSelection.PLAN:
             elements = self.plan(targets)
         else:
diff --git a/buildstream/_stream.py b/buildstream/_stream.py
index cdef69094..f2806b4c8 100644
--- a/buildstream/_stream.py
+++ b/buildstream/_stream.py
@@ -514,6 +514,13 @@ class Stream():
 
         elements, track_elements = self._load(targets, track_targets)
 
+        nonexisting = []
+        for element in elements:
+            if not self.workspace_exists(element.name):
+                nonexisting.append(element.name)
+        if nonexisting:
+            raise StreamError("Workspace does not exist", detail="\n".join(nonexisting))
+
         # Do the tracking first
         if track_first:
             self._fetch(elements, track_elements=track_elements)
@@ -663,6 +670,37 @@ class Stream():
             self._collect_sources(tempdir, tar_location,
                                   target.normal_name, compression)
 
+    # redirect_element_names()
+    #
+    # Takes a list of element names and returns a list where elements have been
+    # redirected to their source elements if the element file exists, and just
+    # the name, if not.
+    #
+    # Args:
+    #    elements (list of str): The element names to redirect
+    #
+    # Returns:
+    #    (list of str): The element names after redirecting
+    #
+    def redirect_element_names(self, elements):
+        element_dir = self._project.element_path
+        load_elements = []
+        output_elements = set()
+
+        for e in elements:
+            element_path = os.path.join(element_dir, e)
+            if os.path.exists(element_path):
+                load_elements.append(e)
+            else:
+                output_elements.add(e)
+        if load_elements:
+            loaded_elements, _ = self._load(load_elements, ())
+
+            for e in loaded_elements:
+                output_elements.add(e.name)
+
+        return list(output_elements)
+
     #############################################################
     #                 Scheduler API forwarding                  #
     #############################################################
@@ -803,7 +841,7 @@ class Stream():
         # Now move on to loading primary selection.
         #
         self._pipeline.resolve_elements(elements)
-        selected = self._pipeline.get_selection(elements, selection)
+        selected = self._pipeline.get_selection(elements, selection, silent=False)
         selected = self._pipeline.except_elements(elements,
                                                   selected,
                                                   except_elements)