buildstream/_pickler.py

@@ -20,6 +20,7 @@
 import os
 import pickle
 import hashlib
+import io
 
 from contextlib import contextmanager
 from collections import namedtuple
@@ -53,8 +54,6 @@ class BstUnpickler(pickle.Unpickler):
     def persistent_load(self, pid):
         type_tag, key_id = pid
         if type_tag == "Project":
-            # XXX: This doesn't actually help, we need to load the junction
-            #      to create the project
             for project in self._context.get_projects():
                 if key_id == project.name:
                     return project
@@ -66,7 +65,27 @@ class BstUnpickler(pickle.Unpickler):
 
 
 CachedProject = namedtuple('CachedProject', ['path', 'project_sum', 'elements'])
-CachedYaml = namedtuple('CachedYaml', ['key', 'contents'])
+
+
+class CachedYaml():
+    def __init__(self, key, contents):
+        self.key = key
+        self.contents = contents
+        self.pickled_contents = None
+
+    # Intercept and return the pickled object
+    def __getstate__(self):
+        data = self.__dict__.copy()
+        if not self.pickled_contents:
+            # Pickle self.contents and store it as self.pickled_contents
+            picklestream = io.BytesIO()
+            BstPickler(picklestream).dump(self.contents)
+            provenance = _yaml.node_get_provenance(self.contents)
+            project = provenance.filename.project
+            data['pickled_contents'] = picklestream
+        data['contents'] = None
+
+        return data
 
 
 class YamlCache():
@@ -86,12 +105,15 @@ class YamlCache():
             if filepath in project_cache.elements:
                 cachedyaml = project_cache.elements[filepath]
                 if cachedyaml.key == key:
+                    if not cachedyaml.contents:
+                        cachedyaml.pickled_contents.seek(0)
+                        cachedyaml.contents = BstUnpickler(cachedyaml.pickled_contents, project._context).load()
                     return _yaml.node_copy(cachedyaml.contents)
         return None
 
     def put(self, project, filepath, key, value):
-        if project.name in self.project_caches:
-            # XXX: Needs a check that the project hasn't changed
+        if project.name in self.project_caches \
+                and project.shasum == self.project_caches[project.name].project_sum:
             project_cache = self.project_caches[project.name]
         else:
             project_cache = self.project_caches[project.name] = CachedProject(project.directory, project.shasum, {})