Merge branch 'develop' into path-33

Conflicts:
	sqlalchemy_mptt/events.py

Author: uralbash

sqlalchemy_mptt/events.py

@@ -12,6 +12,7 @@
 import weakref
 
 from sqlalchemy import and_, case, select, event, inspection
+from sqlalchemy.orm import object_session
 from sqlalchemy.orm.base import NO_VALUE
 from sqlalchemy.sql import func
 
@@ -325,14 +326,24 @@ def pop(self):
         return self.popitem()[0]
 
 
+class _WeakDefaultDict(weakref.WeakKeyDictionary, object):
+
+    def __getitem__(self, key):
+        try:
+            return super(_WeakDefaultDict, self).__getitem__(key)
+        except KeyError:
+            self[key] = value = _WeakDictBasedSet()
+            return value
+
+
 class TreesManager(object):
     """
     Manages events dispatching for all subclasses of a given class.
     """
     def __init__(self, base_class):
         self.base_class = base_class
         self.classes = set()
-        self.instances = _WeakDictBasedSet()
+        self.instances = _WeakDefaultDict()
 
     def register_mapper(self, mapper):
         for e, h in (
@@ -385,29 +396,37 @@ def register_factory(self, sessionmaker):
         return sessionmaker
 
     def before_insert(self, mapper, connection, instance):
-        self.instances.add(instance)
+        session = object_session(instance)
+        self.instances[session].add(instance)
         mptt_before_insert(mapper, connection, instance)
 
     def before_update(self, mapper, connection, instance):
-        self.instances.add(instance)
+        session = object_session(instance)
+        self.instances[session].add(instance)
         mptt_before_update(mapper, connection, instance)
 
     def before_delete(self, mapper, connection, instance):
-        self.instances.discard(instance)
+        session = object_session(instance)
+        self.instances[session].discard(instance)
         mptt_before_delete(mapper, connection, instance)
 
     def after_flush_postexec(self, session, context):
         """
         Event listener to recursively expire `left` and `right` attributes the
         parents of all modified instances part of this flush.
         """
-        while self.instances:
-            instance = self.instances.pop()
+        instances = self.instances[session]
+        while instances:
+            instance = instances.pop()
+            if instance not in session:
+                continue
             parent = self.get_parent_value(instance)
             while parent != NO_VALUE and parent is not None:
-                self.instances.discard(parent)
-                session.expire(parent, ['left', 'right', 'tree_id'])
+                instances.discard(parent)
+                session.expire(parent, ['left', 'right'])
                 parent = self.get_parent_value(parent)
+            else:
+                session.expire(instance, ['tree_id', 'level'])
 
     @staticmethod
     def get_parent_value(instance):

sqlalchemy_mptt/tests/tree_testing_base.py

@@ -8,6 +8,7 @@
 
 
 from sqlalchemy import create_engine
+from sqlalchemy.exc import IntegrityError
 from sqlalchemy.orm import sessionmaker
 
 from sqlalchemy_mptt import mptt_sessionmaker
@@ -169,6 +170,17 @@ def test_tree_orm_initialize(self):
         self.assertEqual(t5.left, 6)
         self.assertEqual(t5.right, 7)
 
+    def test_flush_with_transient_nodes_present(self):
+        transient_node = self.model(ppk=1, parent=None)
+        self.session.add(transient_node)
+        try:
+            self.session.flush()
+        except IntegrityError:
+            pass
+        self.session.rollback()
+        self.session.add(self.model(ppk=46, parent=None))
+        self.session.flush()
+
     def test_tree_initialize(self):
         """ Initial state of the trees
 
@@ -1726,6 +1738,7 @@ def test_session_expire_for_move_after_to_new_tree(self):
         self.session.flush()
 
         self.assertEqual(node.tree_id, 2)
+        self.assertEqual(node.level, 1)
         self.assertEqual(node.parent_id, None)
         self.assertEqual(children[0].tree_id, 2)
         self.assertEqual(children[1].tree_id, 2)