Function to jsonify BSON objects

flask_pymongo/__init__.py

@@ -24,11 +24,17 @@
 # POSSIBILITY OF SUCH DAMAGE.
 
 
-__all__ = ('PyMongo', 'ASCENDING', 'DESCENDING')
+__all__ = ('PyMongo', 'ASCENDING', 'DESCENDING', 'jsonify')
 
+try:
+    import simplejson as json
+except ImportError:
+    import json
+
+from bson import ObjectId
 from bson.errors import InvalidId
-from bson.objectid import ObjectId
-from flask import abort, current_app, request
+from bson.json_util import dumps as bson_dumps
+from flask import abort, current_app, request, jsonify as flask_jsonify
 from gridfs import GridFS, NoFile
 from mimetypes import guess_type
 from pymongo import uri_parser
@@ -339,3 +345,22 @@ def save_upload(filename):
 
         storage = GridFS(self.db, base)
         storage.put(fileobj, filename=filename, content_type=content_type)
+
+
+def jsonify(obj, *args, **kwargs):
+    """Same call signature as `flask.jsonify`.  Works with MongoDB
+    BSON objects.
+
+    .. code-block:: python
+
+        @app.route('/get_bson_doc')
+        def get_bson_doc():
+            data = bson.BSON.encode({'a': 1, 'id_': ObjectId()})
+            return jsonify(data)
+
+    :param object obj: The object to be jsonified.
+
+    :return Response response: The response object to be returned.
+    """
+
+    return flask_jsonify(**json.loads(bson_dumps(obj, *args, **kwargs)))

tests/test_util_functions.py

@@ -0,0 +1,78 @@
+import sys
+
+from tests import util
+
+from flask import Flask, jsonify as flask_jsonify
+from flask.ext.pymongo import jsonify
+from bson import ObjectId, Binary, Code, Regex, DBRef
+
+
+if sys.version_info[0] == 2:
+    str = unicode
+
+
+class BsonJsonifyTest(util.FlaskRequestTest):
+
+    def setUp(self):
+        self.app = Flask('test')
+        self.context = self.app.test_request_context('/')
+        self.context.push()
+
+    def tearDown(self):
+        self.context.pop()
+
+    def test_jsonify_ObjectId(self):
+        objectid = ObjectId(b'foo-bar-quuz')
+        json = {'a': 1, 'id_': objectid}
+        safe_json = {'a': 1, 'id_': {'$oid': str(objectid)}}
+
+        jsonified_bson = jsonify(json).response
+        jsonified = flask_jsonify(safe_json).response
+
+        assert jsonified_bson == jsonified
+
+    def test_jsonify_Binary(self):
+        binary = Binary(b"hello")
+        json = {'a': 1, 'bin': binary}
+        safe_json = {'a': 1, 'bin': {'$binary': "aGVsbG8=", "$type": "00"}}
+
+        jsonified_bson = jsonify(json).response
+        jsonified = flask_jsonify(safe_json).response
+
+        assert jsonified_bson == jsonified
+
+    def test_jsonify_Code(self):
+        code = Code("function () { console.log('Hello, world!'); }();")
+        json = {'a': 1, 'code': code}
+        safe_json = {'a': 1, 'code': {'$code': str(code), '$scope': {}}}
+
+        jsonified_bson = jsonify(json).response
+        jsonified = flask_jsonify(safe_json).response
+
+        assert jsonified_bson == jsonified
+
+    def test_jsonify_Regex(self):
+        regex = Regex("bb|[^b]{2}")
+        json = {'a': 1, 'regex': regex}
+        safe_json = {'a': 1, 'regex': {'$regex': "bb|[^b]{2}", "$options": ""}}
+
+        jsonified_bson = jsonify(json).response
+        jsonified = flask_jsonify(safe_json).response
+
+        assert jsonified_bson == jsonified
+
+    def test_jsonify_DBRef(self):
+        dbref = DBRef("fake_document", "helloworld")
+        json = {'a': 1, 'dbref': dbref}
+        safe_json = {
+            'a': 1,
+            'dbref': {
+                '$id': 'helloworld',
+                '$ref': 'fake_document'
+            }
+        }
+
+        jsonified_bson = jsonify(json).response
+        jsonified = flask_jsonify(safe_json).response
+
+        assert jsonified_bson == jsonified