Fixed #138: assertion error in obj.flush()

Author: abetkin

pony/orm/core.py

@@ -1593,7 +1593,7 @@ def flush(cache):
             cache.modified = False
 
             for obj, status in saved_objects:
-                if obj is not None: obj._after_save_(status)
+                obj._after_save_(status)
         else:
             if cache.modified: throw(TransactionError,
                 'Recursion depth limit reached in obj._after_save_() call')
@@ -4082,6 +4082,7 @@ def populate_criteria_list(criteria_list, columns, converters, params_count=0, t
 statuses = set(['created', 'cancelled', 'loaded', 'modified', 'inserted', 'updated', 'marked_to_delete', 'deleted'])
 del_statuses = set(['marked_to_delete', 'deleted', 'cancelled'])
 created_or_deleted_statuses = set(['created']) | del_statuses
+saved_statuses = set(['inserted', 'updated', 'deleted'])
 
 def throw_object_was_deleted(obj):
     assert obj._status_ in del_statuses
@@ -4723,18 +4724,37 @@ def _save_(obj, dependent_objects=None):
         cache = obj._session_cache_
         assert cache.is_alive
         status = obj._status_
-        if status in ('loaded', 'inserted', 'updated', 'cancelled'): return
 
         if status in ('created', 'modified'):
             obj._save_principal_objects_(dependent_objects)
 
-        obj._save_pos_ = None
         if status == 'created': obj._save_created_()
         elif status == 'modified': obj._save_updated_()
         elif status == 'marked_to_delete': obj._save_deleted_()
-        else: assert False  # pragma: no cover
+        else: assert False, "_save_() called for object %r with incorrect status %s" % (obj, status)  # pragma: no cover
+
+        assert obj._status_ in saved_statuses
+        obj._save_pos_ = None
     def flush(obj):
-        with obj._session_cache_.flush_disabled(): obj._save_()
+        status = obj._status_
+        if status not in ('created', 'modified', 'marked_to_delete'):
+            return
+
+        save_pos = obj._save_pos_
+        assert save_pos is not None, 'save_pos is None for %s object' % status
+
+        with obj._session_cache_.flush_disabled():
+            obj._before_save_() # should be inside flush_disabled to prevent infinite recursion
+                                # TODO: add to documentation that flush is disabled inside before_xxx hooks
+            obj._save_()
+
+        objects_to_save = obj._session_cache_.objects_to_save
+        if save_pos == len(objects_to_save) - 1:
+            objects_to_save.pop()
+        else:
+            objects_to_save[save_pos] = None
+
+        obj._after_save_(status)
     def _before_save_(obj):
         status = obj._status_
         if status == 'created': obj.before_insert()

pony/orm/tests/test_objects_to_save_cleanup.py

@@ -0,0 +1,69 @@
+
+
+import unittest
+
+
+from pony.orm import *
+
+db = Database()
+
+class TestPost(db.Entity):
+    category = Optional('TestCategory')
+    name = Optional(str, default='Noname')
+
+class TestCategory(db.Entity):
+    posts = Set(TestPost)
+
+db.bind('sqlite', ':memory:')
+db.generate_mapping(create_tables=True)
+
+
+class EntityStatusTestCase(object):
+
+    def make_flush(self, obj=None):
+        raise NotImplementedError
+
+    @db_session
+    def test_delete_updated(self):
+        p = TestPost()
+        self.make_flush(p)
+        p.name = 'Pony'
+        self.assertEqual(p._status_, 'modified')        
+        self.make_flush(p)
+        self.assertEqual(p._status_, 'updated')
+        p.delete()
+        self.assertEqual(p._status_, 'marked_to_delete')
+        self.make_flush(p)
+        self.assertEqual(p._status_, 'deleted')
+
+    @db_session
+    def test_delete_inserted(self):
+        p = TestPost()
+        self.assertEqual(p._status_, 'created')
+        self.make_flush(p)
+        self.assertEqual(p._status_, 'inserted')
+        p.delete()
+
+    @db_session
+    def test_cancelled(self):
+        p = TestPost()
+        self.assertEqual(p._status_, 'created')
+        p.delete()
+        self.assertEqual(p._status_, 'cancelled')
+        self.make_flush(p)
+        self.assertEqual(p._status_, 'cancelled')
+
+
+
+class EntityStatusTestCase_ObjectFlush(EntityStatusTestCase,
+                                       unittest.TestCase):
+
+    def make_flush(self, obj=None):
+        obj.flush()
+    
+
+class EntityStatusTestCase_FullFlush(EntityStatusTestCase,
+                                     unittest.TestCase):
+
+    def make_flush(self, obj=None):
+        flush() # full flush