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

5 files changed, 177 insertions, 30 deletions

diff --git a/doc/build/changelog/unreleased_13/4823.rst b/doc/build/changelog/unreleased_13/4823.rst
new file mode 100644
index 000000000..7541330e6
--- /dev/null
+++ b/doc/build/changelog/unreleased_13/4823.rst
@@ -0,0 +1,9 @@
+.. change::
+    :tags: bug, orm
+    :tickets: 4823
+
+    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.
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
diff --git a/test/orm/test_pickled.py b/test/orm/test_pickled.py
index f2ef8e2e9..399c881ac 100644
--- a/test/orm/test_pickled.py
+++ b/test/orm/test_pickled.py
@@ -96,6 +96,32 @@ class PickleTest(fixtures.MappedTest):
             test_needs_fk=True,
         )
 
+    def _option_test_fixture(self):
+        users, addresses, dingalings = (
+            self.tables.users,
+            self.tables.addresses,
+            self.tables.dingalings,
+        )
+
+        mapper(
+            User,
+            users,
+            properties={"addresses": relationship(Address, backref="user")},
+        )
+        mapper(
+            Address,
+            addresses,
+            properties={"dingaling": relationship(Dingaling)},
+        )
+        mapper(Dingaling, dingalings)
+        sess = create_session()
+        u1 = User(name="ed")
+        u1.addresses.append(Address(email_address="ed@bar.com"))
+        sess.add(u1)
+        sess.flush()
+        sess.expunge_all()
+        return sess, User, Address, Dingaling
+
     def test_transient(self):
         users, addresses = (self.tables.users, self.tables.addresses)
 
@@ -418,40 +444,65 @@ class PickleTest(fixtures.MappedTest):
         eq_(sa.inspect(u2).info["some_key"], "value")
 
     @testing.requires.non_broken_pickle
-    def test_options_with_descriptors(self):
-        users, addresses, dingalings = (
-            self.tables.users,
-            self.tables.addresses,
-            self.tables.dingalings,
-        )
+    def test_unbound_options(self):
+        sess, User, Address, Dingaling = self._option_test_fixture()
 
-        mapper(
-            User,
-            users,
-            properties={"addresses": relationship(Address, backref="user")},
-        )
-        mapper(
-            Address,
-            addresses,
-            properties={"dingaling": relationship(Dingaling)},
-        )
-        mapper(Dingaling, dingalings)
-        sess = create_session()
-        u1 = User(name="ed")
-        u1.addresses.append(Address(email_address="ed@bar.com"))
-        sess.add(u1)
-        sess.flush()
-        sess.expunge_all()
+        for opt in [
+            sa.orm.joinedload(User.addresses),
+            sa.orm.joinedload("addresses"),
+            sa.orm.defer("name"),
+            sa.orm.defer(User.name),
+            sa.orm.joinedload("addresses").joinedload(Address.dingaling),
+        ]:
+            opt2 = pickle.loads(pickle.dumps(opt))
+            eq_(opt.path, opt2.path)
+
+        u1 = sess.query(User).options(opt).first()
+        pickle.loads(pickle.dumps(u1))
+
+    @testing.requires.non_broken_pickle
+    def test_bound_options(self):
+        sess, User, Address, Dingaling = self._option_test_fixture()
+
+        for opt in [
+            sa.orm.Load(User).joinedload(User.addresses),
+            sa.orm.Load(User).joinedload("addresses"),
+            sa.orm.Load(User).defer("name"),
+            sa.orm.Load(User).defer(User.name),
+            sa.orm.Load(User)
+            .joinedload("addresses")
+            .joinedload(Address.dingaling),
+            sa.orm.Load(User)
+            .joinedload("addresses", innerjoin=True)
+            .joinedload(Address.dingaling),
+        ]:
+            opt2 = pickle.loads(pickle.dumps(opt))
+            eq_(opt.path, opt2.path)
+            eq_(opt.context.keys(), opt2.context.keys())
+            eq_(opt.local_opts, opt2.local_opts)
+
+        u1 = sess.query(User).options(opt).first()
+        pickle.loads(pickle.dumps(u1))
+
+    @testing.requires.non_broken_pickle
+    def test_became_bound_options(self):
+        sess, User, Address, Dingaling = self._option_test_fixture()
 
         for opt in [
             sa.orm.joinedload(User.addresses),
             sa.orm.joinedload("addresses"),
             sa.orm.defer("name"),
             sa.orm.defer(User.name),
-            sa.orm.joinedload("addresses", Address.dingaling),
+            sa.orm.joinedload("addresses").joinedload(Address.dingaling),
         ]:
+            q = sess.query(User).options(opt)
+            opt = [
+                v for v in q._attributes.values() if isinstance(v, sa.orm.Load)
+            ][0]
+
             opt2 = pickle.loads(pickle.dumps(opt))
             eq_(opt.path, opt2.path)
+            eq_(opt.local_opts, opt2.local_opts)
 
         u1 = sess.query(User).options(opt).first()
         pickle.loads(pickle.dumps(u1))
diff --git a/test/orm/test_utils.py b/test/orm/test_utils.py
index 45dd2e38b..e47fc3f26 100644
--- a/test/orm/test_utils.py
+++ b/test/orm/test_utils.py
@@ -3,6 +3,7 @@ from sqlalchemy import inspect
 from sqlalchemy import Integer
 from sqlalchemy import MetaData
 from sqlalchemy import Table
+from sqlalchemy import util
 from sqlalchemy.ext.hybrid import hybrid_method
 from sqlalchemy.ext.hybrid import hybrid_property
 from sqlalchemy.orm import aliased
@@ -743,6 +744,62 @@ class PathRegistryTest(_fixtures.FixtureTest):
         eq_(p2.serialize(), [(User, "addresses"), (Address, None)])
         eq_(p3.serialize(), [(User, "addresses")])
 
+    def test_serialize_context_dict(self):
+        reg = util.OrderedDict()
+        umapper = inspect(self.classes.User)
+        amapper = inspect(self.classes.Address)
+
+        p1 = PathRegistry.coerce((umapper, umapper.attrs.addresses))
+        p2 = PathRegistry.coerce((umapper, umapper.attrs.addresses, amapper))
+        p3 = PathRegistry.coerce((amapper, amapper.attrs.email_address))
+
+        p1.set(reg, "p1key", "p1value")
+        p2.set(reg, "p2key", "p2value")
+        p3.set(reg, "p3key", "p3value")
+        eq_(
+            reg,
+            {
+                ("p1key", p1.path): "p1value",
+                ("p2key", p2.path): "p2value",
+                ("p3key", p3.path): "p3value",
+            },
+        )
+
+        serialized = PathRegistry.serialize_context_dict(
+            reg, ("p1key", "p2key")
+        )
+        eq_(
+            serialized,
+            [
+                (("p1key", p1.serialize()), "p1value"),
+                (("p2key", p2.serialize()), "p2value"),
+            ],
+        )
+
+    def test_deseralize_context_dict(self):
+        umapper = inspect(self.classes.User)
+        amapper = inspect(self.classes.Address)
+
+        p1 = PathRegistry.coerce((umapper, umapper.attrs.addresses))
+        p2 = PathRegistry.coerce((umapper, umapper.attrs.addresses, amapper))
+        p3 = PathRegistry.coerce((amapper, amapper.attrs.email_address))
+
+        serialized = [
+            (("p1key", p1.serialize()), "p1value"),
+            (("p2key", p2.serialize()), "p2value"),
+            (("p3key", p3.serialize()), "p3value"),
+        ]
+        deserialized = PathRegistry.deserialize_context_dict(serialized)
+
+        eq_(
+            deserialized,
+            {
+                ("p1key", p1.path): "p1value",
+                ("p2key", p2.path): "p2value",
+                ("p3key", p3.path): "p3value",
+            },
+        )
+
     def test_deseralize(self):
         User = self.classes.User
         Address = self.classes.Address