Serialize the context dictionary in Load objects

Fixed bug where :class:`.Load` objects were not pickleable due to
mapper/relationship state in the internal context dictionary.  These
objects are now converted to picklable using similar techniques as that of
other elements within the loader option system that have long been
serializable.

Fixes: #4823
Change-Id: Id2a0d8b640ac475c86d6416ad540671e66d410e5

2 files changed, 36 insertions, 6 deletions

diff --git a/lib/sqlalchemy/orm/path_registry.py b/lib/sqlalchemy/orm/path_registry.py
index 4803dbecb..2f680a3a1 100644
--- a/lib/sqlalchemy/orm/path_registry.py
+++ b/lib/sqlalchemy/orm/path_registry.py
@@ -100,8 +100,8 @@ class PathRegistry(object):
     def __reduce__(self):
         return _unreduce_path, (self.serialize(),)
 
-    def serialize(self):
-        path = self.path
+    @classmethod
+    def _serialize_path(cls, path):
         return list(
             zip(
                 [m.class_ for m in [path[i] for i in range(0, len(path), 2)]],
@@ -110,10 +110,7 @@ class PathRegistry(object):
         )
 
     @classmethod
-    def deserialize(cls, path):
-        if path is None:
-            return None
-
+    def _deserialize_path(cls, path):
         p = tuple(
             chain(
                 *[
@@ -129,6 +126,35 @@ class PathRegistry(object):
         )
         if p and p[-1] is None:
             p = p[0:-1]
+        return p
+
+    @classmethod
+    def serialize_context_dict(cls, dict_, tokens):
+        return [
+            ((key, cls._serialize_path(path)), value)
+            for (key, path), value in [
+                (k, v)
+                for k, v in dict_.items()
+                if isinstance(k, tuple) and k[0] in tokens
+            ]
+        ]
+
+    @classmethod
+    def deserialize_context_dict(cls, serialized):
+        return util.OrderedDict(
+            ((key, tuple(cls._deserialize_path(path))), value)
+            for (key, path), value in serialized
+        )
+
+    def serialize(self):
+        path = self.path
+        return self._serialize_path(path)
+
+    @classmethod
+    def deserialize(cls, path):
+        if path is None:
+            return None
+        p = cls._deserialize_path(path)
         return cls.coerce(p)
 
     @classmethod
diff --git a/lib/sqlalchemy/orm/strategy_options.py b/lib/sqlalchemy/orm/strategy_options.py
index df7dd51a8..c50b7d041 100644
--- a/lib/sqlalchemy/orm/strategy_options.py
+++ b/lib/sqlalchemy/orm/strategy_options.py
@@ -465,12 +465,16 @@ class Load(Generative, MapperOption):
 
     def __getstate__(self):
         d = self.__dict__.copy()
+        d["context"] = PathRegistry.serialize_context_dict(
+            d["context"], ("loader",)
+        )
         d["path"] = self.path.serialize()
         return d
 
     def __setstate__(self, state):
         self.__dict__.update(state)
         self.path = PathRegistry.deserialize(self.path)
+        self.context = PathRegistry.deserialize_context_dict(self.context)
 
     def _chop_path(self, to_chop, path):
         i = -1