Encode DBRef and ObjectId

Author: idocyabra

flask_mongoengine/json.py

@@ -1,4 +1,4 @@
-from bson import json_util
+from bson import json_util, DBRef, ObjectId
 from flask.json import JSONEncoder
 from mongoengine.base import BaseDocument
 from mongoengine.queryset import QuerySet
@@ -16,6 +16,10 @@ def default(self, obj):
                 return json_util._json_convert(obj.to_mongo())
             elif isinstance(obj, QuerySet):
                 return json_util._json_convert(obj.as_pymongo())
+            elif isinstance(obj, DBRef):
+                return obj.id
+            elif isinstance(obj, ObjectId):
+                return obj.__str__()
             return superclass.default(self, obj)
 
     return MongoEngineJSONEncoder

tests/test_json_app.py

@@ -1,6 +1,6 @@
 import flask
 import pytest
-from bson import ObjectId
+from bson import ObjectId, DBRef
 
 
 @pytest.fixture(autouse=True)
@@ -22,6 +22,14 @@ def add():
     def show(id):
         return flask.jsonify(result=Todo.objects.get_or_404(id=id))
 
+    @app.route("/object_id")
+    def object_id():
+        return flask.jsonify(result=ObjectId())
+
+    @app.route("/dbref")
+    def dbref():
+        return flask.jsonify(result=DBRef('Todo', ObjectId()))
+
 
 def test_with_id(app, todo):
     Todo = todo
@@ -32,6 +40,12 @@ def test_with_id(app, todo):
     response = client.post("/add", data={"title": "First Item", "text": "The text"})
     assert response.status_code == 200
 
+    response = client.get("/dbref")
+    assert response.status_code == 200
+
+    response = client.get("/object_id")
+    assert response.status_code == 200
+
     response = client.get("/show/%s/" % Todo.objects.first().id)
     assert response.status_code == 200