Create and store data inside projects when opening workspaces

Changes to _context.py:
* Context has been extended to contain a WorkspaceProjectCache, as there
  are times when we want to use it before a Workspaces can be initialised
  (looking up a WorkspaceProject to find the directory that the project is
  in)

Changes to _stream.py:
* Removed staging the elements from workspace_open() and workspace_reset()

Changes in _workspaces.py:
* A new WorkspaceProject contains all the information needed to refer back
  to a project from its workspace (currently this is the project path and
  the element used to create this workspace)
* This is stored within a new WorkspaceProjectCache object, which keeps
  WorkspaceProjects around so they don't need to be loaded from disk
  repeatedly.
* Workspaces has been extended to contain the WorkspaceProjectCache, and
  will use it when opening and closing workspaces.
* Workspaces.create_workspace has been extended to handle the staging of
  the element into the workspace, in addition to creating the equivalent
  WorkspaceProject file.

This is a part of #222

3 files changed, 249 insertions, 23 deletions

diff --git a/buildstream/_context.py b/buildstream/_context.py
index 7ca60e7aa..55d0fd489 100644
--- a/buildstream/_context.py
+++ b/buildstream/_context.py
@@ -32,7 +32,7 @@ from ._message import Message, MessageType
 from ._profile import Topics, profile_start, profile_end
 from ._artifactcache import ArtifactCache
 from ._artifactcache.cascache import CASCache
-from ._workspaces import Workspaces
+from ._workspaces import Workspaces, WorkspaceProjectCache
 from .plugin import _plugin_lookup
 
 
@@ -140,6 +140,7 @@ class Context():
         self._projects = []
         self._project_overrides = {}
         self._workspaces = None
+        self._workspace_project_cache = WorkspaceProjectCache()
         self._log_handle = None
         self._log_filename = None
         self._cascache = None
@@ -285,7 +286,7 @@ class Context():
     #
     def add_project(self, project):
         if not self._projects:
-            self._workspaces = Workspaces(project)
+            self._workspaces = Workspaces(project, self._workspace_project_cache)
         self._projects.append(project)
 
     # get_projects():
@@ -312,6 +313,16 @@ class Context():
     def get_workspaces(self):
         return self._workspaces
 
+    # get_workspace_project_cache():
+    #
+    # Return the WorkspaceProjectCache object used for this BuildStream invocation
+    #
+    # Returns:
+    #    (WorkspaceProjectCache): The WorkspaceProjectCache object
+    #
+    def get_workspace_project_cache(self):
+        return self._workspace_project_cache
+
     # get_overrides():
     #
     # Fetch the override dictionary for the active project. This returns
diff --git a/buildstream/_stream.py b/buildstream/_stream.py
index 6f298c259..3256003c8 100644
--- a/buildstream/_stream.py
+++ b/buildstream/_stream.py
@@ -581,15 +581,7 @@ class Stream():
                     todo_elements = "\nDid not try to create workspaces for " + todo_elements
                 raise StreamError("Failed to create workspace directory: {}".format(e) + todo_elements) from e
 
-            workspaces.create_workspace(target._get_full_name(), directory)
-
-            if not no_checkout:
-                with target.timed_activity("Staging sources to {}".format(directory)):
-                    target._open_workspace()
-
-            # Saving the workspace once it is set up means that if the next workspace fails to be created before
-            # the configuration gets saved. The successfully created workspace still gets saved.
-            workspaces.save_config()
+            workspaces.create_workspace(target, directory, checkout=not no_checkout)
             self._message(MessageType.INFO, "Created a workspace for element: {}"
                           .format(target._get_full_name()))
 
@@ -672,10 +664,7 @@ class Stream():
                                       .format(workspace_path, e)) from e
 
             workspaces.delete_workspace(element._get_full_name())
-            workspaces.create_workspace(element._get_full_name(), workspace_path)
-
-            with element.timed_activity("Staging sources to {}".format(workspace_path)):
-                element._open_workspace()
+            workspaces.create_workspace(element, workspace_path, checkout=True)
 
             self._message(MessageType.INFO,
                           "Reset workspace for {} at: {}".format(element.name,
diff --git a/buildstream/_workspaces.py b/buildstream/_workspaces.py
index 468073f05..466c55ce7 100644
--- a/buildstream/_workspaces.py
+++ b/buildstream/_workspaces.py
@@ -25,6 +25,202 @@ from ._exceptions import LoadError, LoadErrorReason
 
 
 BST_WORKSPACE_FORMAT_VERSION = 3
+BST_WORKSPACE_PROJECT_FORMAT_VERSION = 1
+WORKSPACE_PROJECT_FILE = ".bstproject.yaml"
+
+
+# WorkspaceProject()
+#
+# An object to contain various helper functions and data required for
+# referring from a workspace back to buildstream.
+#
+# Args:
+#    directory (str): The directory that the workspace exists in.
+#
+class WorkspaceProject():
+    def __init__(self, directory):
+        self._projects = []
+        self._directory = directory
+
+    # get_default_project_path()
+    #
+    # Retrieves the default path to a project.
+    #
+    # Returns:
+    #    (str): The path to a project
+    #
+    def get_default_project_path(self):
+        return self._projects[0]['project-path']
+
+    # get_default_element()
+    #
+    # Retrieves the name of the element that owns this workspace.
+    #
+    # Returns:
+    #    (str): The name of an element
+    #
+    def get_default_element(self):
+        return self._projects[0]['element-name']
+
+    # to_dict()
+    #
+    # Turn the members data into a dict for serialization purposes
+    #
+    # Returns:
+    #    (dict): A dict representation of the WorkspaceProject
+    #
+    def to_dict(self):
+        ret = {
+            'projects': self._projects,
+            'format-version': BST_WORKSPACE_PROJECT_FORMAT_VERSION,
+        }
+        return ret
+
+    # from_dict()
+    #
+    # Loads a new WorkspaceProject from a simple dictionary
+    #
+    # Args:
+    #    directory (str): The directory that the workspace exists in
+    #    dictionary (dict): The dict to generate a WorkspaceProject from
+    #
+    # Returns:
+    #   (WorkspaceProject): A newly instantiated WorkspaceProject
+    #
+    @classmethod
+    def from_dict(cls, directory, dictionary):
+        # Only know how to handle one format-version at the moment.
+        format_version = int(dictionary['format-version'])
+        assert format_version == BST_WORKSPACE_PROJECT_FORMAT_VERSION, \
+            "Format version {} not found in {}".format(BST_WORKSPACE_PROJECT_FORMAT_VERSION, dictionary)
+
+        workspace_project = cls(directory)
+        for item in dictionary['projects']:
+            workspace_project.add_project(item['project-path'], item['element-name'])
+
+        return workspace_project
+
+    # load()
+    #
+    # Loads the WorkspaceProject for a given directory.
+    #
+    # Args:
+    #    directory (str): The directory
+    # Returns:
+    #    (WorkspaceProject): The created WorkspaceProject, if in a workspace, or
+    #    (NoneType): None, if the directory is not inside a workspace.
+    #
+    @classmethod
+    def load(cls, directory):
+        workspace_file = os.path.join(directory, WORKSPACE_PROJECT_FILE)
+        if os.path.exists(workspace_file):
+            data_dict = _yaml.load(workspace_file)
+            return cls.from_dict(directory, data_dict)
+        else:
+            return None
+
+    # write()
+    #
+    # Writes the WorkspaceProject to disk
+    #
+    def write(self):
+        os.makedirs(self._directory, exist_ok=True)
+        _yaml.dump(self.to_dict(), self.get_filename())
+
+    # get_filename()
+    #
+    # Returns the full path to the workspace local project file
+    #
+    def get_filename(self):
+        return os.path.join(self._directory, WORKSPACE_PROJECT_FILE)
+
+    # add_project()
+    #
+    # Adds an entry containing the project's path and element's name.
+    #
+    # Args:
+    #    project_path (str): The path to the project that opened the workspace.
+    #    element_name (str): The name of the element that the workspace belongs to.
+    #
+    def add_project(self, project_path, element_name):
+        assert (project_path and element_name)
+        self._projects.append({'project-path': project_path, 'element-name': element_name})
+
+
+# WorkspaceProjectCache()
+#
+# A class to manage workspace project data for multiple workspaces.
+#
+class WorkspaceProjectCache():
+    def __init__(self):
+        self._projects = {}  # Mapping of a workspace directory to its WorkspaceProject
+
+    # get()
+    #
+    # Returns a WorkspaceProject for a given directory, retrieving from the cache if
+    # present.
+    #
+    # Args:
+    #    directory (str): The directory to search for a WorkspaceProject.
+    #
+    # Returns:
+    #    (WorkspaceProject): The WorkspaceProject that was found for that directory.
+    #    or      (NoneType): None, if no WorkspaceProject can be found.
+    #
+    def get(self, directory):
+        try:
+            workspace_project = self._projects[directory]
+        except KeyError:
+            workspace_project = WorkspaceProject.load(directory)
+            if workspace_project:
+                self._projects[directory] = workspace_project
+
+        return workspace_project
+
+    # add()
+    #
+    # Adds the project path and element name to the WorkspaceProject that exists
+    # for that directory
+    #
+    # Args:
+    #    directory (str): The directory to search for a WorkspaceProject.
+    #    project_path (str): The path to the project that refers to this workspace
+    #    element_name (str): The element in the project that was refers to this workspace
+    #
+    # Returns:
+    #    (WorkspaceProject): The WorkspaceProject that was found for that directory.
+    #
+    def add(self, directory, project_path, element_name):
+        workspace_project = self.get(directory)
+        if not workspace_project:
+            workspace_project = WorkspaceProject(directory)
+            self._projects[directory] = workspace_project
+
+        workspace_project.add_project(project_path, element_name)
+        return workspace_project
+
+    # remove()
+    #
+    # Removes the project path and element name from the WorkspaceProject that exists
+    # for that directory.
+    #
+    # NOTE: This currently just deletes the file, but with support for multiple
+    # projects opening the same workspace, this will involve decreasing the count
+    # and deleting the file if there are no more projects.
+    #
+    # Args:
+    #    directory (str): The directory to search for a WorkspaceProject.
+    #
+    def remove(self, directory):
+        workspace_project = self.get(directory)
+        if not workspace_project:
+            raise LoadError(LoadErrorReason.MISSING_FILE,
+                            "Failed to find a {} file to remove".format(WORKSPACE_PROJECT_FILE))
+        path = workspace_project.get_filename()
+        try:
+            os.unlink(path)
+        except FileNotFoundError:
+            pass
 
 
 # Workspace()
@@ -199,12 +395,14 @@ class Workspace():
 #
 # Args:
 #    toplevel_project (Project): Top project used to resolve paths.
+#    workspace_project_cache (WorkspaceProjectCache): The cache of WorkspaceProjects
 #
 class Workspaces():
-    def __init__(self, toplevel_project):
+    def __init__(self, toplevel_project, workspace_project_cache):
         self._toplevel_project = toplevel_project
         self._bst_directory = os.path.join(toplevel_project.directory, ".bst")
         self._workspaces = self._load_config()
+        self._workspace_project_cache = workspace_project_cache
 
     # list()
     #
@@ -219,19 +417,36 @@ class Workspaces():
 
     # create_workspace()
     #
-    # Create a workspace in the given path for the given element.
+    # Create a workspace in the given path for the given element, and potentially
+    # checks-out the target into it.
     #
     # Args:
-    #    element_name (str) - The element name to create a workspace for
+    #    target (Element) - The element to create a workspace for
     #    path (str) - The path in which the workspace should be kept
+    #    checkout (bool): Whether to check-out the element's sources into the directory
     #
-    def create_workspace(self, element_name, path):
-        if path.startswith(self._toplevel_project.directory):
-            path = os.path.relpath(path, self._toplevel_project.directory)
+    def create_workspace(self, target, path, *, checkout):
+        element_name = target._get_full_name()
+        project_dir = self._toplevel_project.directory
+        if path.startswith(project_dir):
+            workspace_path = os.path.relpath(path, project_dir)
+        else:
+            workspace_path = path
 
-        self._workspaces[element_name] = Workspace(self._toplevel_project, path=path)
+        self._workspaces[element_name] = Workspace(self._toplevel_project, path=workspace_path)
 
-        return self._workspaces[element_name]
+        if checkout:
+            with target.timed_activity("Staging sources to {}".format(path)):
+                target._open_workspace()
+
+        workspace_project = self._workspace_project_cache.add(path, project_dir, element_name)
+        project_file_path = workspace_project.get_filename()
+
+        if os.path.exists(project_file_path):
+            target.warn("{} was staged from this element's sources".format(WORKSPACE_PROJECT_FILE))
+        workspace_project.write()
+
+        self.save_config()
 
     # get_workspace()
     #
@@ -280,8 +495,19 @@ class Workspaces():
     #    element_name (str) - The element name whose workspace to delete
     #
     def delete_workspace(self, element_name):
+        workspace = self.get_workspace(element_name)
         del self._workspaces[element_name]
 
+        # Remove from the cache if it exists
+        try:
+            self._workspace_project_cache.remove(workspace.get_absolute_path())
+        except LoadError as e:
+            # We might be closing a workspace with a deleted directory
+            if e.reason == LoadErrorReason.MISSING_FILE:
+                pass
+            else:
+                raise
+
     # save_config()
     #
     # Dump the current workspace element to the project configuration