- the column assigned to polymorphic_on now behaves like any other

mapped attribute, in that it can be assigned to, mapped to multiple
columns.  It is also populated immediately upon object construction
with its class-based value, so the attribute can be read before
any flush occurs.  [ticket:1895]

1 files changed, 60 insertions, 40 deletions

diff --git a/lib/sqlalchemy/orm/mapper.py b/lib/sqlalchemy/orm/mapper.py
index 8fc0b4132..32eb8f643 100644
--- a/lib/sqlalchemy/orm/mapper.py
+++ b/lib/sqlalchemy/orm/mapper.py
@@ -208,6 +208,7 @@ class Mapper(object):
             self._configure_class_instrumentation()
             self._configure_listeners()
             self._configure_properties()
+            self._configure_polymorphic_setter()
             self._configure_pks()
             global _new_mappers
             _new_mappers = True
@@ -310,7 +311,7 @@ class Mapper(object):
             if self.polymorphic_identity is not None:
                 self.polymorphic_map[self.polymorphic_identity] = self
             self._identity_class = self.class_
-
+        
         if self.mapped_table is None:
             raise sa_exc.ArgumentError(
                     "Mapper '%s' does not have a mapped_table specified." 
@@ -560,33 +561,65 @@ class Mapper(object):
                                     init=False, 
                                     setparent=True)
 
+    def _configure_polymorphic_setter(self):
+        """Configure an attribute on the mapper representing the 
+        'polymorphic_on' column, if applicable, and not 
+        already generated by _configure_properties (which is typical).
+        
+        Also create a setter function which will assign this
+        attribute to the value of the 'polymorphic_identity'
+        upon instance construction, also if applicable.  This 
+        routine will run when an instance is created.
+        
+        """
         # do a special check for the "discriminiator" column, as it 
         # may only be present in the 'with_polymorphic' selectable 
         # but we need it for the base mapper
-        if self.polymorphic_on is not None and \
-                    self.polymorphic_on not in self._columntoproperty:
-            col = self.mapped_table.corresponding_column(self.polymorphic_on)
-            if col is None:
-                instrument = False
-                col = self.polymorphic_on
-                if self.with_polymorphic is None \
-                    or self.with_polymorphic[1].corresponding_column(col) \
-                    is None:
-                    util.warn("Could not map polymorphic_on column "
-                              "'%s' to the mapped table - polymorphic "
-                              "loads will not function properly"
-                              % col.description)
-            else:
-                instrument = True
-            if self._should_exclude(col.key, col.key, local=False, column=col):
-                raise sa_exc.InvalidRequestError(
-                    "Cannot exclude or override the discriminator column %r" %
-                    col.key)
+        setter = False
+        
+        if self.polymorphic_on is not None:
+            setter = True
+
+            if self.polymorphic_on not in self._columntoproperty:
+                col = self.mapped_table.corresponding_column(self.polymorphic_on)
+                if col is None:
+                    setter = False
+                    instrument = False
+                    col = self.polymorphic_on
+                    if self.with_polymorphic is None \
+                        or self.with_polymorphic[1].corresponding_column(col) \
+                        is None:
+                        util.warn("Could not map polymorphic_on column "
+                                  "'%s' to the mapped table - polymorphic "
+                                  "loads will not function properly"
+                                  % col.description)
+                else:
+                    instrument = True
+
+                if self._should_exclude(col.key, col.key, False, col):
+                    raise sa_exc.InvalidRequestError(
+                        "Cannot exclude or override the discriminator column %r" %
+                        col.key)
                     
-            self._configure_property(
-                            col.key, 
-                            properties.ColumnProperty(col, _instrument=instrument),
-                            init=False, setparent=True)
+                self._configure_property(
+                                col.key, 
+                                properties.ColumnProperty(col, _instrument=instrument),
+                                init=False, setparent=True)
+                polymorphic_key = col.key
+            else:
+                polymorphic_key = self._columntoproperty[self.polymorphic_on].key
+
+        if setter:
+            def _set_polymorphic_identity(state):
+                dict_ = state.dict
+                state.get_impl(polymorphic_key).set(state, dict_,
+                        self.polymorphic_identity, None)
+
+            self._set_polymorphic_identity = _set_polymorphic_identity
+        else:
+            self._set_polymorphic_identity = None
+        
+
 
     def _adapt_inherited_property(self, key, prop, init):
         if not self.concrete:
@@ -1656,13 +1689,6 @@ class Mapper(object):
                         if col is mapper.version_id_col:
                             params[col.key] = \
                               mapper.version_id_generator(None)
-                        elif mapper.polymorphic_on is not None and \
-                                mapper.polymorphic_on.shares_lineage(col):
-                            value = mapper.polymorphic_identity
-                            if ((col.default is None and
-                                 col.server_default is None) or
-                                value is not None):
-                                params[col.key] = value
                         elif col in pks:
                             value = \
                              mapper._get_state_attr_by_column(
@@ -1715,10 +1741,6 @@ class Mapper(object):
                                                     passive=True)
                                     if history.added:
                                         hasdata = True
-                        elif mapper.polymorphic_on is not None and \
-                            mapper.polymorphic_on.shares_lineage(col) and \
-                                col not in pks:
-                            pass
                         else:
                             prop = mapper._columntoproperty[col]
                             history = attributes.get_state_history(
@@ -1890,11 +1912,6 @@ class Mapper(object):
         postfetch_cols = resultproxy.postfetch_cols()
         generated_cols = list(resultproxy.prefetch_cols())
 
-        if self.polymorphic_on is not None:
-            po = table.corresponding_column(self.polymorphic_on)
-            if po is not None:
-                generated_cols.append(po)
-
         if self.version_id_col is not None:
             generated_cols.append(self.version_id_col)
 
@@ -2390,6 +2407,9 @@ def _event_on_init(state, args, kwargs):
     if instrumenting_mapper:
         # compile() always compiles all mappers
         instrumenting_mapper.compile()
+        
+        if instrumenting_mapper._set_polymorphic_identity:
+            instrumenting_mapper._set_polymorphic_identity(state)
 
 def _event_on_resurrect(state):
     # re-populate the primary key elements