Merge "Improve support for with_polymorphic in mapper options"

3 files changed, 110 insertions, 48 deletions

diff --git a/lib/sqlalchemy/orm/path_registry.py b/lib/sqlalchemy/orm/path_registry.py
index 7d93da296..4803dbecb 100644
--- a/lib/sqlalchemy/orm/path_registry.py
+++ b/lib/sqlalchemy/orm/path_registry.py
@@ -62,14 +62,14 @@ class PathRegistry(object):
 
     def set(self, attributes, key, value):
         log.debug("set '%s' on path '%s' to '%s'", key, self, value)
-        attributes[(key, self.path)] = value
+        attributes[(key, self.natural_path)] = value
 
     def setdefault(self, attributes, key, value):
         log.debug("setdefault '%s' on path '%s' to '%s'", key, self, value)
-        attributes.setdefault((key, self.path), value)
+        attributes.setdefault((key, self.natural_path), value)
 
     def get(self, attributes, key, value=None):
-        key = (key, self.path)
+        key = (key, self.natural_path)
         if key in attributes:
             return attributes[key]
         else:
@@ -160,7 +160,7 @@ class RootRegistry(PathRegistry):
 
     """
 
-    path = ()
+    path = natural_path = ()
     has_entity = False
     is_aliased_class = False
     is_root = True
@@ -177,6 +177,7 @@ class TokenRegistry(PathRegistry):
         self.token = token
         self.parent = parent
         self.path = parent.path + (token,)
+        self.natural_path = parent.natural_path + (token,)
 
     has_entity = False
 
@@ -186,6 +187,13 @@ class TokenRegistry(PathRegistry):
         if not self.parent.is_aliased_class and not self.parent.is_root:
             for ent in self.parent.mapper.iterate_to_root():
                 yield TokenRegistry(self.parent.parent[ent], self.token)
+        elif (
+            self.parent.is_aliased_class
+            and self.parent.entity._is_with_polymorphic
+        ):
+            yield self
+            for ent in self.parent.entity._with_polymorphic_entities:
+                yield TokenRegistry(self.parent.parent[ent], self.token)
         else:
             yield self
 
@@ -211,6 +219,7 @@ class PropRegistry(PathRegistry):
         self.prop = prop
         self.parent = parent
         self.path = parent.path + (prop,)
+        self.natural_path = parent.natural_path + (prop,)
 
         self._wildcard_path_loader_key = (
             "loader",
@@ -255,6 +264,24 @@ class EntityRegistry(PathRegistry, dict):
         self.is_aliased_class = entity.is_aliased_class
         self.entity = entity
         self.path = parent.path + (entity,)
+
+        # the "natural path" is the path that we get when Query is traversing
+        # from the lead entities into the various relationships; it corresponds
+        # to the structure of mappers and relationships. when we are given a
+        # path that comes from loader options, as of 1.3 it can have ac-hoc
+        # with_polymorphic() and other AliasedInsp objects inside of it, which
+        # are usually not present in mappings.  So here we track both the
+        # "enhanced" path in self.path and the "natural" path that doesn't
+        # include those objects so these two traversals can be matched up.
+        if parent.path and self.is_aliased_class:
+            if entity.mapper.isa(parent.natural_path[-1].entity):
+                self.natural_path = parent.natural_path + (entity.mapper,)
+            else:
+                self.natural_path = parent.natural_path + (
+                    parent.natural_path[-1].entity,
+                )
+        else:
+            self.natural_path = self.path
         self.entity_path = self
 
     @property
diff --git a/lib/sqlalchemy/orm/strategy_options.py b/lib/sqlalchemy/orm/strategy_options.py
index ee3b83caa..8a34771c7 100644
--- a/lib/sqlalchemy/orm/strategy_options.py
+++ b/lib/sqlalchemy/orm/strategy_options.py
@@ -236,7 +236,9 @@ class Load(Generative, MapperOption):
 
             path = path[attr]
         elif _is_mapped_class(attr):
-            if not attr.common_parent(path.mapper):
+            if not orm_util._entity_corresponds_to_use_path_impl(
+                attr.parent, path[-1]
+            ):
                 if raiseerr:
                     raise sa_exc.ArgumentError(
                         "Attribute '%s' does not "
@@ -247,11 +249,13 @@ class Load(Generative, MapperOption):
         else:
             prop = found_property = attr.property
 
-            if not prop.parent.common_parent(path.mapper):
+            if not orm_util._entity_corresponds_to_use_path_impl(
+                attr.parent, path[-1]
+            ):
                 if raiseerr:
                     raise sa_exc.ArgumentError(
-                        "Attribute '%s' does not "
-                        "link from element '%s'" % (attr, path.entity)
+                        'Attribute "%s" does not '
+                        'link from element "%s"' % (attr, path.entity)
                     )
                 else:
                     return None
@@ -272,36 +276,15 @@ class Load(Generative, MapperOption):
                         _existing_alias=existing,
                     )
                     ext_info = inspect(ac)
-                elif not ext_info.with_polymorphic_mappers:
-                    ext_info = orm_util.AliasedInsp(
-                        ext_info.entity,
-                        ext_info.mapper.base_mapper,
-                        ext_info.selectable,
-                        ext_info.name,
-                        ext_info.with_polymorphic_mappers or [ext_info.mapper],
-                        ext_info.polymorphic_on,
-                        ext_info._base_alias,
-                        ext_info._use_mapper_path,
-                        ext_info._adapt_on_names,
-                        ext_info.represents_outer_join,
-                    )
 
                 path.entity_path[prop].set(
                     self.context, "path_with_polymorphic", ext_info
                 )
 
-                # the path here will go into the context dictionary and
-                # needs to match up to how the class graph is traversed.
-                # so we can't put an AliasedInsp in the path here, needs
-                # to be the base mapper.
-                path = path[prop][ext_info.mapper]
-
-                # but, we need to know what the original of_type()
-                # argument is for cache key purposes.  so....store that too.
-                # it might be better for "path" to really represent,
-                # "the path", but trying to keep the impact of the cache
-                # key feature localized for now
+                path = path[prop][ext_info]
+
                 self._of_type = of_type_info
+
             else:
                 path = path[prop]
 
diff --git a/lib/sqlalchemy/orm/util.py b/lib/sqlalchemy/orm/util.py
index 92dd4c4ec..f9258895d 100644
--- a/lib/sqlalchemy/orm/util.py
+++ b/lib/sqlalchemy/orm/util.py
@@ -589,12 +589,32 @@ class AliasedInsp(InspectionAttr):
             self.persist_selectable
         ) = self.local_table = selectable
         self.name = name
-        self.with_polymorphic_mappers = with_polymorphic_mappers
         self.polymorphic_on = polymorphic_on
         self._base_alias = _base_alias or self
         self._use_mapper_path = _use_mapper_path
         self.represents_outer_join = represents_outer_join
 
+        if with_polymorphic_mappers:
+            self._is_with_polymorphic = True
+            self.with_polymorphic_mappers = with_polymorphic_mappers
+            self._with_polymorphic_entities = []
+            for poly in self.with_polymorphic_mappers:
+                if poly is not mapper:
+                    ent = AliasedClass(
+                        poly.class_,
+                        selectable,
+                        base_alias=self,
+                        adapt_on_names=adapt_on_names,
+                        use_mapper_path=_use_mapper_path,
+                    )
+
+                    setattr(self.entity, poly.class_.__name__, ent)
+                    self._with_polymorphic_entities.append(ent._aliased_insp)
+
+        else:
+            self._is_with_polymorphic = False
+            self.with_polymorphic_mappers = [mapper]
+
         self._adapter = sql_util.ColumnAdapter(
             selectable,
             equivalents=mapper._equivalent_columns,
@@ -605,20 +625,6 @@ class AliasedInsp(InspectionAttr):
         self._adapt_on_names = adapt_on_names
         self._target = mapper.class_
 
-        for poly in self.with_polymorphic_mappers:
-            if poly is not mapper:
-                setattr(
-                    self.entity,
-                    poly.class_.__name__,
-                    AliasedClass(
-                        poly.class_,
-                        selectable,
-                        base_alias=self,
-                        adapt_on_names=adapt_on_names,
-                        use_mapper_path=_use_mapper_path,
-                    ),
-                )
-
     is_aliased_class = True
     "always returns True"
 
@@ -718,7 +724,17 @@ class AliasedInsp(InspectionAttr):
         )
 
     def __str__(self):
-        return "aliased(%s)" % (self._target.__name__,)
+        if self._is_with_polymorphic:
+            return "with_polymorphic(%s, [%s])" % (
+                self._target.__name__,
+                ", ".join(
+                    mp.class_.__name__
+                    for mp in self.with_polymorphic_mappers
+                    if mp is not self.mapper
+                ),
+            )
+        else:
+            return "aliased(%s)" % (self._target.__name__,)
 
 
 inspection._inspects(AliasedClass)(lambda target: target._aliased_insp)
@@ -1225,6 +1241,42 @@ def _entity_corresponds_to(given, entity):
     return entity.common_parent(given)
 
 
+def _entity_corresponds_to_use_path_impl(given, entity):
+    """determine if 'given' corresponds to 'entity', in terms
+    of a path of loader options where a mapped attribute is taken to
+    be a member of a parent entity.
+
+    e.g.::
+
+        someoption(A).someoption(A.b)  # -> fn(A, A) -> True
+        someoption(A).someoption(C.d)  # -> fn(A, C) -> False
+
+        a1 = aliased(A)
+        someoption(a1).someoption(A.b) # -> fn(a1, A) -> False
+        someoption(a1).someoption(a1.b) # -> fn(a1, a1) -> True
+
+        wp = with_polymorphic(A, [A1, A2])
+        someoption(wp).someoption(A1.foo)  # -> fn(wp, A1) -> False
+        someoption(wp).someoption(wp.A1.foo)  # -> fn(wp, wp.A1) -> True
+
+
+    """
+    if given.is_aliased_class:
+        return (
+            entity.is_aliased_class
+            and not entity._use_mapper_path
+            and given is entity
+            or given in entity._with_polymorphic_entities
+        )
+    elif not entity.is_aliased_class:
+        return given.common_parent(entity.mapper)
+    else:
+        return (
+            entity._use_mapper_path
+            and given in entity.with_polymorphic_mappers
+        )
+
+
 def _entity_isa(given, mapper):
     """determine if 'given' "is a" mapper, in terms of the given
     would load rows of type 'mapper'.