Add method get_or_404 method to the document class

flaskext/mongokit.py

@@ -36,12 +36,26 @@ def show_task(task_id):
     def to_python(self, value):
         try:
             return bson.ObjectId(value)
-        except bson.errors.InvalidId, e:
+        except bson.errors.InvalidId:
             raise abort(400)
 
     def to_url(self, value):
         return str(value)
 
+class Document(Document):
+    def get_or_404(self, id):
+        """This method get one document over the _id. If there no document
+        with this id then it will raised a 404 error.
+        
+        :param id: The id from the document. The most time there will be 
+                   a ObjectId.
+        """
+        doc = self.get_from_id(id)
+        if doc is None:
+            abort(404)
+        else:
+            return doc
+
 class MongoKit(object):
     """This class is used to integrate `MongoKit`_ into a Flask application.
 

tests.py

@@ -6,8 +6,8 @@
 
 from flask import Flask
 from flaskext.mongokit import MongoKit, BSONObjectIdConverter, \
-                              Document, Database, Collection
-from werkzeug.exceptions import BadRequest
+                              Document, Collection
+from werkzeug.exceptions import BadRequest, NotFound
 from bson import ObjectId
 
 class BlogPost(Document):
@@ -84,6 +84,25 @@ def test_save_and_find_document(self):
         assert rec_post.body == rec_post.body
         assert rec_post.author == rec_post.author
 
+    def test_get_or_404(self):
+        self.db.register([BlogPost])
+
+        assert len(self.db.registered_documents) > 0
+        assert self.db.registered_documents[0] == BlogPost
+
+        post = self.db.BlogPost()
+        post.title = u"Flask-MongoKit"
+        post.body = u"Flask-MongoKit is a layer between Flask and MongoKit"
+        post.author = u"Christoph Heer"
+        post.save()
+
+        assert self.db.BlogPost.find().count() > 0
+        assert "get_or_404" in dir(self.db.BlogPost)
+        try:
+            self.db.BlogPost.get_or_404(post['_id'])
+        except NotFound:
+            self.fail("There should be a document with this id")
+        self.assertRaises(NotFound, self.db.BlogPost.get_or_404, ObjectId())
 
 def suite():
     suite = unittest.TestSuite()