Add signal_kwargs arg for save/delete/bulk insert

docs/changelog.rst

@@ -6,6 +6,7 @@ Changes in 0.10.7 - DEV
 =======================
 - Fixed not being able to specify `use_db_field=False` on `ListField(EmbeddedDocumentField)` instances
 - Fixed cascade delete mixing among collections #1224
+- Add `signal_kwargs` argument to `Document.save`, `Document.delete` and `BaseQuerySet.insert` to be passed to signals calls #1206
 
 Changes in 0.10.6
 =================

mongoengine/document.py

@@ -250,7 +250,7 @@ class Document(BaseDocument):
 
     def save(self, force_insert=False, validate=True, clean=True,
              write_concern=None, cascade=None, cascade_kwargs=None,
-             _refs=None, save_condition=None, **kwargs):
+             _refs=None, save_condition=None, signal_kwargs=None, **kwargs):
         """Save the :class:`~mongoengine.Document` to the database. If the
         document already exists, it will be updated, otherwise it will be
         created.
@@ -276,6 +276,8 @@ class Document(BaseDocument):
         :param save_condition: only perform save if matching record in db
             satisfies condition(s) (e.g. version number).
             Raises :class:`OperationError` if the conditions are not satisfied
+        :parm signal_kwargs: (optional) kwargs dictionary to be passed to
+            the signal calls.
 
         .. versionchanged:: 0.5
             In existing documents it only saves changed fields using
@@ -297,8 +299,11 @@ class Document(BaseDocument):
             :class:`OperationError` exception raised if save_condition fails.
         .. versionchanged:: 0.10.1
             :class: save_condition failure now raises a `SaveConditionError`
+        .. versionchanged:: 0.10.7
+            Add signal_kwargs argument
         """
-        signals.pre_save.send(self.__class__, document=self)
+        signal_kwargs = signal_kwargs or {}
+        signals.pre_save.send(self.__class__, document=self, **signal_kwargs)
 
         if validate:
             self.validate(clean=clean)
@@ -311,7 +316,7 @@ class Document(BaseDocument):
         created = ('_id' not in doc or self._created or force_insert)
 
         signals.pre_save_post_validation.send(self.__class__, document=self,
-                                              created=created)
+                                              created=created, **signal_kwargs)
 
         try:
             collection = self._get_collection()
@@ -400,7 +405,8 @@ class Document(BaseDocument):
         if created or id_field not in self._meta.get('shard_key', []):
             self[id_field] = self._fields[id_field].to_python(object_id)
 
-        signals.post_save.send(self.__class__, document=self, created=created)
+        signals.post_save.send(self.__class__, document=self,
+                               created=created, **signal_kwargs)
         self._clear_changed_fields()
         self._created = False
         return self
@@ -476,18 +482,24 @@ class Document(BaseDocument):
         # Need to add shard key to query, or you get an error
         return self._qs.filter(**self._object_key).update_one(**kwargs)
 
-    def delete(self, **write_concern):
+    def delete(self, signal_kwargs=None, **write_concern):
         """Delete the :class:`~mongoengine.Document` from the database. This
         will only take effect if the document has been previously saved.
 
+        :parm signal_kwargs: (optional) kwargs dictionary to be passed to
+            the signal calls.
         :param write_concern: Extra keyword arguments are passed down which
             will be used as options for the resultant
             ``getLastError`` command.  For example,
             ``save(..., write_concern={w: 2, fsync: True}, ...)`` will
             wait until at least two servers have recorded the write and
             will force an fsync on the primary server.
+
+        .. versionchanged:: 0.10.7
+            Add signal_kwargs argument
         """
-        signals.pre_delete.send(self.__class__, document=self)
+        signal_kwargs = signal_kwargs or {}
+        signals.pre_delete.send(self.__class__, document=self, **signal_kwargs)
 
         # Delete FileFields separately 
         FileField = _import_class('FileField')
@@ -501,7 +513,7 @@ class Document(BaseDocument):
         except pymongo.errors.OperationFailure, err:
             message = u'Could not delete document (%s)' % err.message
             raise OperationError(message)
-        signals.post_delete.send(self.__class__, document=self)
+        signals.post_delete.send(self.__class__, document=self, **signal_kwargs)
 
     def switch_db(self, db_alias, keep_created=True):
         """

mongoengine/queryset/base.py

@@ -266,7 +266,8 @@ class BaseQuerySet(object):
             result = None
         return result
 
-    def insert(self, doc_or_docs, load_bulk=True, write_concern=None):
+    def insert(self, doc_or_docs, load_bulk=True,
+               write_concern=None, signal_kwargs=None):
         """bulk insert documents
 
         :param doc_or_docs: a document or list of documents to be inserted
@@ -279,11 +280,15 @@ class BaseQuerySet(object):
                 ``insert(..., {w: 2, fsync: True})`` will wait until at least
                 two servers have recorded the write and will force an fsync on
                 each server being written to.
+        :parm signal_kwargs: (optional) kwargs dictionary to be passed to
+            the signal calls.
 
         By default returns document instances, set ``load_bulk`` to False to
         return just ``ObjectIds``
 
         .. versionadded:: 0.5
+        .. versionchanged:: 0.10.7
+            Add signal_kwargs argument
         """
         Document = _import_class('Document')
 
@@ -305,7 +310,9 @@ class BaseQuerySet(object):
                 msg = "Some documents have ObjectIds use doc.update() instead"
                 raise OperationError(msg)
 
-        signals.pre_bulk_insert.send(self._document, documents=docs)
+        signal_kwargs = signal_kwargs or {}
+        signals.pre_bulk_insert.send(self._document,
+                                     documents=docs, **signal_kwargs)
 
         raw = [doc.to_mongo() for doc in docs]
         try:
@@ -324,7 +331,7 @@ class BaseQuerySet(object):
 
         if not load_bulk:
             signals.post_bulk_insert.send(
-                self._document, documents=docs, loaded=False)
+                self._document, documents=docs, loaded=False, **signal_kwargs)
             return return_one and ids[0] or ids
 
         documents = self.in_bulk(ids)
@@ -332,7 +339,7 @@ class BaseQuerySet(object):
         for obj_id in ids:
             results.append(documents.get(obj_id))
         signals.post_bulk_insert.send(
-            self._document, documents=results, loaded=True)
+            self._document, documents=results, loaded=True, **signal_kwargs)
         return return_one and results[0] or results
 
     def count(self, with_limit_and_skip=False):

tests/test_signals.py

@@ -25,6 +25,8 @@ class SignalTests(unittest.TestCase):
         connect(db='mongoenginetest')
 
         class Author(Document):
+            # Make the id deterministic for easier testing
+            id = SequenceField(primary_key=True)
             name = StringField()
 
             def __unicode__(self):
@@ -33,7 +35,7 @@ class SignalTests(unittest.TestCase):
             @classmethod
             def pre_init(cls, sender, document, *args, **kwargs):
                 signal_output.append('pre_init signal, %s' % cls.__name__)
-                signal_output.append(str(kwargs['values']))
+                signal_output.append(kwargs['values'])
 
             @classmethod
             def post_init(cls, sender, document, **kwargs):
@@ -43,48 +45,55 @@ class SignalTests(unittest.TestCase):
             @classmethod
             def pre_save(cls, sender, document, **kwargs):
                 signal_output.append('pre_save signal, %s' % document)
+                signal_output.append(kwargs)
 
             @classmethod
             def pre_save_post_validation(cls, sender, document, **kwargs):
                 signal_output.append('pre_save_post_validation signal, %s' % document)
-                if 'created' in kwargs:
+                if kwargs.pop('created', False):
-                    if kwargs['created']:
                     signal_output.append('Is created')
                 else:
                     signal_output.append('Is updated')
+                signal_output.append(kwargs)
 
             @classmethod
             def post_save(cls, sender, document, **kwargs):
                 dirty_keys = document._delta()[0].keys() + document._delta()[1].keys()
                 signal_output.append('post_save signal, %s' % document)
                 signal_output.append('post_save dirty keys, %s' % dirty_keys)
-                if 'created' in kwargs:
+                if kwargs.pop('created', False):
-                    if kwargs['created']:
                     signal_output.append('Is created')
                 else:
                     signal_output.append('Is updated')
+                signal_output.append(kwargs)
 
             @classmethod
             def pre_delete(cls, sender, document, **kwargs):
                 signal_output.append('pre_delete signal, %s' % document)
+                signal_output.append(kwargs)
 
             @classmethod
             def post_delete(cls, sender, document, **kwargs):
                 signal_output.append('post_delete signal, %s' % document)
+                signal_output.append(kwargs)
 
             @classmethod
             def pre_bulk_insert(cls, sender, documents, **kwargs):
                 signal_output.append('pre_bulk_insert signal, %s' % documents)
+                signal_output.append(kwargs)
 
             @classmethod
             def post_bulk_insert(cls, sender, documents, **kwargs):
                 signal_output.append('post_bulk_insert signal, %s' % documents)
-                if kwargs.get('loaded', False):
+                if kwargs.pop('loaded', False):
                     signal_output.append('Is loaded')
                 else:
                     signal_output.append('Not loaded')
+                signal_output.append(kwargs)
+
         self.Author = Author
         Author.drop_collection()
+        Author.id.set_next_value(0)
 
         class Another(Document):
 
@@ -96,10 +105,12 @@ class SignalTests(unittest.TestCase):
             @classmethod
             def pre_delete(cls, sender, document, **kwargs):
                 signal_output.append('pre_delete signal, %s' % document)
+                signal_output.append(kwargs)
 
             @classmethod
             def post_delete(cls, sender, document, **kwargs):
                 signal_output.append('post_delete signal, %s' % document)
+                signal_output.append(kwargs)
 
         self.Another = Another
         Another.drop_collection()
@@ -137,12 +148,18 @@ class SignalTests(unittest.TestCase):
                 for document in documents:
                     if not document.active:
                         document.active = True
+                signal_output.append(kwargs)
 
             @classmethod
             def post_bulk_insert(cls, sender, documents, **kwargs):
                 signal_output.append('post_bulk_insert signal, %s' %
                                      [(doc, {'active': documents[n].active})
                                       for n, doc in enumerate(documents)])
+                if kwargs.pop('loaded', False):
+                    signal_output.append('Is loaded')
+                else:
+                    signal_output.append('Not loaded')
+                signal_output.append(kwargs)
 
         self.Post = Post
         Post.drop_collection()
@@ -237,63 +254,118 @@ class SignalTests(unittest.TestCase):
 
         self.assertEqual(self.get_signal_output(create_author), [
             "pre_init signal, Author",
-            "{'name': 'Bill Shakespeare'}",
+            {'name': 'Bill Shakespeare'},
             "post_init signal, Bill Shakespeare, document._created = True",
         ])
 
         a1 = self.Author(name='Bill Shakespeare')
         self.assertEqual(self.get_signal_output(a1.save), [
             "pre_save signal, Bill Shakespeare",
+            {},
             "pre_save_post_validation signal, Bill Shakespeare",
             "Is created",
+            {},
             "post_save signal, Bill Shakespeare",
             "post_save dirty keys, ['name']",
-            "Is created"
+            "Is created",
+            {}
         ])
 
         a1.reload()
         a1.name = 'William Shakespeare'
         self.assertEqual(self.get_signal_output(a1.save), [
             "pre_save signal, William Shakespeare",
+            {},
             "pre_save_post_validation signal, William Shakespeare",
             "Is updated",
+            {},
             "post_save signal, William Shakespeare",
             "post_save dirty keys, ['name']",
-            "Is updated"
+            "Is updated",
+            {}
         ])
 
         self.assertEqual(self.get_signal_output(a1.delete), [
             'pre_delete signal, William Shakespeare',
+            {},
             'post_delete signal, William Shakespeare',
+            {}
         ])
 
-        signal_output = self.get_signal_output(load_existing_author)
+        self.assertEqual(self.get_signal_output(load_existing_author), [
-        # test signal_output lines separately, because of random ObjectID after object load
-        self.assertEqual(signal_output[0],
             "pre_init signal, Author",
-        )
+            {'id': 2, 'name': 'Bill Shakespeare'},
-        self.assertEqual(signal_output[2],
+            "post_init signal, Bill Shakespeare, document._created = False"
-            "post_init signal, Bill Shakespeare, document._created = False",
+        ])
-        )
 
-
+        self.assertEqual(self.get_signal_output(bulk_create_author_with_load), [
-        signal_output = self.get_signal_output(bulk_create_author_with_load)
+            'pre_init signal, Author',
-
+            {'name': 'Bill Shakespeare'},
-        # The output of this signal is not entirely deterministic. The reloaded
+            'post_init signal, Bill Shakespeare, document._created = True',
-        # object will have an object ID. Hence, we only check part of the output
+            'pre_bulk_insert signal, [<Author: Bill Shakespeare>]',
-        self.assertEqual(signal_output[3], "pre_bulk_insert signal, [<Author: Bill Shakespeare>]"
+            {},
-        )
+            'pre_init signal, Author',
-        self.assertEqual(signal_output[-2:],
+            {'id': 3, 'name': 'Bill Shakespeare'},
-            ["post_bulk_insert signal, [<Author: Bill Shakespeare>]",
+            'post_init signal, Bill Shakespeare, document._created = False',
-             "Is loaded",])
+            'post_bulk_insert signal, [<Author: Bill Shakespeare>]',
+            'Is loaded',
+            {}
+        ])
 
         self.assertEqual(self.get_signal_output(bulk_create_author_without_load), [
             "pre_init signal, Author",
-            "{'name': 'Bill Shakespeare'}",
+            {'name': 'Bill Shakespeare'},
             "post_init signal, Bill Shakespeare, document._created = True",
             "pre_bulk_insert signal, [<Author: Bill Shakespeare>]",
+            {},
             "post_bulk_insert signal, [<Author: Bill Shakespeare>]",
             "Not loaded",
+            {}
+        ])
+
+    def test_signal_kwargs(self):
+        """ Make sure signal_kwargs is passed to signals calls. """
+
+        def live_and_let_die():
+            a = self.Author(name='Bill Shakespeare')
+            a.save(signal_kwargs={'live': True, 'die': False})
+            a.delete(signal_kwargs={'live': False, 'die': True})
+
+        self.assertEqual(self.get_signal_output(live_and_let_die), [
+            "pre_init signal, Author",
+            {'name': 'Bill Shakespeare'},
+            "post_init signal, Bill Shakespeare, document._created = True",
+            "pre_save signal, Bill Shakespeare",
+            {'die': False, 'live': True},
+            "pre_save_post_validation signal, Bill Shakespeare",
+            "Is created",
+            {'die': False, 'live': True},
+            "post_save signal, Bill Shakespeare",
+            "post_save dirty keys, ['name']",
+            "Is created",
+            {'die': False, 'live': True},
+            'pre_delete signal, Bill Shakespeare',
+            {'die': True, 'live': False},
+            'post_delete signal, Bill Shakespeare',
+            {'die': True, 'live': False}
+        ])
+
+        def bulk_create_author():
+            a1 = self.Author(name='Bill Shakespeare')
+            self.Author.objects.insert([a1], signal_kwargs={'key': True})
+
+        self.assertEqual(self.get_signal_output(bulk_create_author), [
+            'pre_init signal, Author',
+            {'name': 'Bill Shakespeare'},
+            'post_init signal, Bill Shakespeare, document._created = True',
+            'pre_bulk_insert signal, [<Author: Bill Shakespeare>]',
+            {'key': True},
+            'pre_init signal, Author',
+            {'id': 2, 'name': 'Bill Shakespeare'},
+            'post_init signal, Bill Shakespeare, document._created = False',
+            'post_bulk_insert signal, [<Author: Bill Shakespeare>]',
+            'Is loaded',
+            {'key': True}
         ])
 
     def test_queryset_delete_signals(self):
@@ -302,7 +374,9 @@ class SignalTests(unittest.TestCase):
         self.Another(name='Bill Shakespeare').save()
         self.assertEqual(self.get_signal_output(self.Another.objects.delete), [
             'pre_delete signal, Bill Shakespeare',
+            {},
             'post_delete signal, Bill Shakespeare',
+            {}
         ])
 
     def test_signals_with_explicit_doc_ids(self):
@@ -353,7 +427,10 @@ class SignalTests(unittest.TestCase):
         results = self.get_signal_output(bulk_set_active_post)
         self.assertEqual(results, [
             "pre_bulk_insert signal, [(<Post: Post 1>, {'active': False}), (<Post: Post 2>, {'active': False}), (<Post: Post 3>, {'active': False})]",
-            "post_bulk_insert signal, [(<Post: Post 1>, {'active': True}), (<Post: Post 2>, {'active': True}), (<Post: Post 3>, {'active': True})]"
+            {},
+            "post_bulk_insert signal, [(<Post: Post 1>, {'active': True}), (<Post: Post 2>, {'active': True}), (<Post: Post 3>, {'active': True})]",
+            'Is loaded',
+            {}
         ])
 
 if __name__ == '__main__':