From bd9960b71067a51d17dadc17274f88de601681cb Mon Sep 17 00:00:00 2001
From: Christoph Heer <Christoph.Heer@gmail.com>
Date: Sun, 9 Oct 2011 09:53:47 +0200
Subject: [PATCH] Add method get_or_404 method to the document class

---
 flaskext/mongokit.py | 16 +++++++++++++++-
 tests.py             | 23 +++++++++++++++++++++--
 2 files changed, 36 insertions(+), 3 deletions(-)

diff --git a/flaskext/mongokit.py b/flaskext/mongokit.py
index d2e9789..403e3f0 100644
--- a/flaskext/mongokit.py
+++ b/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.
 
diff --git a/tests.py b/tests.py
index 84fa253..d1ab0cb 100644
--- a/tests.py
+++ b/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()