diff --git a/flask_pymongo/__init__.py b/flask_pymongo/__init__.py index 82c7f52..c4e5144 100644 --- a/flask_pymongo/__init__.py +++ b/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))) \ No newline at end of file diff --git a/tests/test_util_functions.py b/tests/test_util_functions.py new file mode 100644 index 0000000..1d1bdc0 --- /dev/null +++ b/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