Update sqlalchemy

libs/sqlalchemy/orm/collections.py

@@ -112,12 +112,32 @@ from sqlalchemy.sql import expression
 from sqlalchemy import schema, util, exc as sa_exc
 
 
+
 __all__ = ['collection', 'collection_adapter',
            'mapped_collection', 'column_mapped_collection',
            'attribute_mapped_collection']
 
 __instrumentation_mutex = util.threading.Lock()
 
+class _SerializableColumnGetter(object):
+    def __init__(self, colkeys):
+        self.colkeys = colkeys
+        self.composite = len(colkeys) > 1
+
+    def __reduce__(self):
+        return _SerializableColumnGetter, (self.colkeys,)
+
+    def __call__(self, value):
+        state = instance_state(value)
+        m = _state_mapper(state)
+        key = [m._get_state_attr_by_column(
+                        state, state.dict, 
+                        m.mapped_table.columns[k])
+                     for k in self.colkeys]
+        if self.composite:
+            return tuple(key)
+        else:
+            return key[0]
 
 def column_mapped_collection(mapping_spec):
     """A dictionary-based collection type with column-based keying.
@@ -131,25 +151,27 @@ def column_mapped_collection(mapping_spec):
     after a session flush.
 
     """
+    global _state_mapper, instance_state
     from sqlalchemy.orm.util import _state_mapper
     from sqlalchemy.orm.attributes import instance_state
 
-    cols = [expression._only_column_elements(q, "mapping_spec") 
-                for q in util.to_list(mapping_spec)]
-    if len(cols) == 1:
-        def keyfunc(value):
-            state = instance_state(value)
-            m = _state_mapper(state)
-            return m._get_state_attr_by_column(state, state.dict, cols[0])
-    else:
-        mapping_spec = tuple(cols)
-        def keyfunc(value):
-            state = instance_state(value)
-            m = _state_mapper(state)
-            return tuple(m._get_state_attr_by_column(state, state.dict, c)
-                         for c in mapping_spec)
+    cols = [c.key for c in [
+                expression._only_column_elements(q, "mapping_spec") 
+                for q in util.to_list(mapping_spec)]]
+    keyfunc = _SerializableColumnGetter(cols)
     return lambda: MappedCollection(keyfunc)
 
+class _SerializableAttrGetter(object):
+    def __init__(self, name):
+        self.name = name
+        self.getter = operator.attrgetter(name)
+
+    def __call__(self, target):
+        return self.getter(target)
+
+    def __reduce__(self):
+        return _SerializableAttrGetter, (self.name, )
+
 def attribute_mapped_collection(attr_name):
     """A dictionary-based collection type with attribute-based keying.
 
@@ -163,7 +185,8 @@ def attribute_mapped_collection(attr_name):
     after a session flush.
 
     """
-    return lambda: MappedCollection(operator.attrgetter(attr_name))
+    getter = _SerializableAttrGetter(attr_name)
+    return lambda: MappedCollection(getter)
 
 
 def mapped_collection(keyfunc):
@@ -814,6 +837,7 @@ def _instrument_class(cls):
             methods[name] = None, None, after
 
     # apply ABC auto-decoration to methods that need it
+
     for method, decorator in decorators.items():
         fn = getattr(cls, method, None)
         if (fn and method not in methods and
@@ -1465,3 +1489,13 @@ class MappedCollection(dict):
                     incoming_key, value, new_key))
             yield value
     _convert = collection.converter(_convert)
+
+# ensure instrumentation is associated with
+# these built-in classes; if a user-defined class
+# subclasses these and uses @internally_instrumented,
+# the superclass is otherwise not instrumented.
+# see [ticket:2406].
+_instrument_class(MappedCollection)
+_instrument_class(InstrumentedList)
+_instrument_class(InstrumentedSet)
+