Added json helpers for Binary and Code types PYTHON-362

Author: rozza

bson/__init__.py

@@ -565,3 +565,11 @@ def has_c():
     .. versionadded:: 1.9
     """
     return _use_c
+
+
+def has_uuid():
+    """Is the uuid module available?
+
+    .. versionadded:: 2.2.1+
+    """
+    return _use_uuid

bson/json_util.py

@@ -14,25 +14,43 @@
 
 """Tools for using Python's :mod:`json` module with BSON documents.
 
-This module provides two methods: `object_hook` and `default`. These
-names are pretty terrible, but match the names used in Python's `json
-library <http://docs.python.org/library/json.html>`_. They allow for
-specialized encoding and decoding of BSON documents into `Mongo
-Extended JSON
+This module provides two helper methods `dumps` and `loads` that wrap the
+native :mod:`json` methods and provide explicit BSON conversion to and from
+json.  This allows for specialized encoding and decoding of BSON documents
+into `Mongo Extended JSON
 <http://www.mongodb.org/display/DOCS/Mongo+Extended+JSON>`_'s *Strict*
 mode.  This lets you encode / decode BSON documents to JSON even when
 they use special BSON types.
 
 Example usage (serialization)::
 
->>> json.dumps(..., default=json_util.default)
+.. doctest::
+
+   >>> from bson import Binary, Code
+   >>> from bson.json_util import dumps
+   >>> dumps([{'foo': [1, 2]},
+   ...        {'bar': {'hello': 'world'}},
+   ...        {'code': Code("function x() { return 1; }")},
+   ...        {'bin': Binary("\x00\x01\x02\x03\x04")}])
+   '[{"foo": [1, 2]}, {"bar": {"hello": "world"}}, {"code": {"$scope": {}, "$code": "function x() { return 1; }"}}, {"bin": {"$type": 0, "$binary": "AAECAwQ=\\n"}}]'
 
 Example usage (deserialization)::
 
->>> json.loads(..., object_hook=json_util.object_hook)
+.. doctest::
+
+   >>> from bson.json_util import loads
+   >>> loads('[{"foo": [1, 2]}, {"bar": {"hello": "world"}}, {"code": {"$scope": {}, "$code": "function x() { return 1; }"}}, {"bin": {"$type": 0, "$binary": "AAECAwQ=\\n"}}]')
+   [{u'foo': [1, 2]}, {u'bar': {u'hello': u'world'}}, {u'code': Code('function x() { return 1; }', {})}, {u'bin': Binary('\x00\x01\x02\x03\x04', 0)}]
 
-Currently this does not handle special encoding and decoding for
-:class:`~bson.binary.Binary` and :class:`~bson.code.Code` instances.
+Alternatively, you can manually pass the `default` to :func:`json.dumps`.
+It won't handle :class:`~bson.binary.Binary` and :class:`~bson.code.Code`
+instances (as they are extended strings you can't provide custom defaults),
+but it will be faster as there is less recursion.
+
+.. versionchanged:: 2.2.1+
+   Added dumps and loads helpers to automatically handle conversion to and
+   from json and supports :class:`~bson.binary.Binary` and
+   :class:`~bson.code.Code`
 
 .. versionchanged:: 1.9
    Handle :class:`uuid.UUID` instances, whenever possible.
@@ -50,34 +68,72 @@
    Added support for encoding/decoding datetimes and regular expressions.
 """
 
+import base64
 import calendar
 import datetime
 import re
+
+json_lib = True
 try:
-    import uuid
-    _use_uuid = True
+    import json
 except ImportError:
-    _use_uuid = False
+    try:
+        import simplejson as json
+    except ImportError:
+        json_lib = False
 
+import bson
 from bson import EPOCH_AWARE
+from bson.binary import Binary
+from bson.code import Code
 from bson.dbref import DBRef
 from bson.max_key import MaxKey
 from bson.min_key import MinKey
 from bson.objectid import ObjectId
 from bson.timestamp import Timestamp
-from bson.tz_util import utc
 
-# TODO support Binary and Code
-# Binary and Code are tricky because they subclass str so json thinks it can
-# handle them. Not sure what the proper way to get around this is...
-#
-# One option is to just add some other method that users need to call _before_
-# calling json.dumps or json.loads. That is pretty terrible though...
+from bson.py3compat import PY3, binary_type, string_types
 
 # TODO share this with bson.py?
 _RE_TYPE = type(re.compile("foo"))
 
 
+def dumps(obj, *args, **kwargs):
+    """Helper function that wraps :class:`json.dumps`.
+
+    Recursive function that handles all BSON types incuding
+    :class:`~bson.binary.Binary` and :class:`~bson.code.Code`.
+    """
+    if not json_lib:
+        raise Exception("No json library available")
+    return json.dumps(_json_convert(obj), *args, **kwargs)
+
+
+def loads(s, *args, **kwargs):
+    """Helper function that wraps :class:`json.loads`.
+
+    Automatically passes the object_hook for BSON type conversion.
+    """
+    if not json_lib:
+        raise Exception("No json library available")
+    kwargs['object_hook'] = object_hook
+    return json.loads(s, *args, **kwargs)
+
+
+def _json_convert(obj):
+    """Recursive helper method that converts BSON types so they can be
+    converted into json.
+    """
+    if hasattr(obj, 'iteritems') or hasattr(obj, 'items'):  # PY3 support
+        return dict(((k, _json_convert(v)) for k, v in obj.iteritems()))
+    elif hasattr(obj, '__iter__') and not isinstance(obj, string_types):
+        return list((_json_convert(v) for v in obj))
+    try:
+        return default(obj)
+    except TypeError:
+        return obj
+
+
 def object_hook(dct):
     if "$oid" in dct:
         return ObjectId(str(dct["$oid"]))
@@ -97,8 +153,12 @@ def object_hook(dct):
         return MinKey()
     if "$maxKey" in dct:
         return MaxKey()
-    if _use_uuid and "$uuid" in dct:
-        return uuid.UUID(dct["$uuid"])
+    if "$binary" in dct:
+        return Binary(base64.b64decode(dct["$binary"].encode()), dct["$type"])
+    if "$code" in dct:
+        return Code(dct["$code"], dct.get("$scope"))
+    if bson.has_uuid() and "$uuid" in dct:
+        return bson.uuid.UUID(dct["$uuid"])
     return dct
 
 
@@ -128,6 +188,14 @@ def default(obj):
         return {"$maxKey": 1}
     if isinstance(obj, Timestamp):
         return {"t": obj.time, "i": obj.inc}
-    if _use_uuid and isinstance(obj, uuid.UUID):
+    if isinstance(obj, Code):
+        return {'$code': "%s" % obj, '$scope': obj.scope}
+    if isinstance(obj, Binary):
+        return {'$binary': base64.b64encode(obj).decode(),
+                '$type': obj.subtype}
+    if PY3 and isinstance(obj, binary_type):
+        return {'$binary': base64.b64encode(obj).decode(),
+                '$type': 0}
+    if bson.has_uuid() and isinstance(obj, bson.uuid.UUID):
         return {"$uuid": obj.hex}
     raise TypeError("%r is not JSON serializable" % obj)

test/test_json_util.py

@@ -18,44 +18,38 @@
 import datetime
 import re
 import sys
-json_lib = True
-try:
-    import json
-except ImportError:
-    try:
-        import simplejson as json
-    except ImportError:
-        json_lib = False
-try:
-    import uuid
-    should_test_uuid = True
-except ImportError:
-    should_test_uuid = False
 
 from nose.plugins.skip import SkipTest
 
 sys.path[0:0] = [""]
 
+import bson
+from bson.py3compat import b
+from bson import json_util
+from bson.binary import Binary, MD5_SUBTYPE
+from bson.code import Code
 from bson.dbref import DBRef
-from bson.json_util import default, object_hook
-from bson.min_key import MinKey
 from bson.max_key import MaxKey
+from bson.min_key import MinKey
 from bson.objectid import ObjectId
 from bson.timestamp import Timestamp
 from bson.tz_util import utc
 
+from test.test_connection import get_connection
+
 PY3 = sys.version_info[0] == 3
 
 
 class TestJsonUtil(unittest.TestCase):
 
     def setUp(self):
-        if not json_lib:
+        if not json_util.json_lib:
             raise SkipTest()
 
+        self.db = get_connection().pymongo_test
+
     def round_tripped(self, doc):
-        return json.loads(json.dumps(doc, default=default),
-                          object_hook=object_hook)
+        return json_util.loads(json_util.dumps(doc))
 
     def round_trip(self, doc):
         self.assertEqual(doc, self.round_tripped(doc))
@@ -76,11 +70,11 @@ def test_dbref(self):
         #
         # self.assertEqual("{\"ref\": {\"$ref\": \"foo\", \"$id\": 5}}",
         #                  json.dumps({"ref": DBRef("foo", 5)},
-        #                  default=default))
+        #                  default=json_util.default))
         # self.assertEqual("{\"ref\": {\"$ref\": \"foo\",
         #                              \"$id\": 5, \"$db\": \"bar\"}}",
         #                  json.dumps({"ref": DBRef("foo", 5, "bar")},
-        #                  default=default))
+        #                  default=json_util.default))
 
     def test_datetime(self):
         # only millis, not micros
@@ -92,27 +86,55 @@ def test_regex(self):
         self.assertEqual("a*b", res.pattern)
         if PY3:
             # re.UNICODE is a default in python 3.
-            self.assertEqual(re.IGNORECASE|re.UNICODE, res.flags)
+            self.assertEqual(re.IGNORECASE | re.UNICODE, res.flags)
         else:
             self.assertEqual(re.IGNORECASE, res.flags)
 
     def test_minkey(self):
         self.round_trip({"m": MinKey()})
 
     def test_maxkey(self):
-        self.round_trip({"m": MinKey()})
+        self.round_trip({"m": MaxKey()})
 
     def test_timestamp(self):
-        res = json.dumps({"ts": Timestamp(4, 13)}, default=default)
-        dct = json.loads(res)
+        res = json_util.json.dumps({"ts": Timestamp(4, 13)},
+                                     default=json_util.default)
+        dct = json_util.json.loads(res)
         self.assertEqual(dct['ts']['t'], 4)
         self.assertEqual(dct['ts']['i'], 13)
 
     def test_uuid(self):
-        if not should_test_uuid:
+        if not bson.has_uuid():
             raise SkipTest()
         self.round_trip(
-                {'uuid': uuid.UUID('f47ac10b-58cc-4372-a567-0e02b2c3d479')})
+                {'uuid': bson.uuid.UUID(
+                            'f47ac10b-58cc-4372-a567-0e02b2c3d479')})
+
+    def test_binary(self):
+        self.round_trip({"bin": Binary(b("\x00\x01\x02\x03\x04"))})
+        self.round_trip({
+            "md5": Binary(b(' n7\x18\xaf\t/\xd1\xd1/\x80\xca\xe7q\xcc\xac'),
+                MD5_SUBTYPE)})
+
+    def test_code(self):
+        self.round_trip({"code": Code("function x() { return 1; }")})
+        self.round_trip({"code": Code("function y() { return z; }", z=2)})
+
+    def test_cursor(self):
+        db = self.db
+
+        db.drop_collection("test")
+        docs = [
+            {'foo': [1, 2]},
+            {'bar': {'hello': 'world'}},
+            {'code': Code("function x() { return 1; }")},
+            {'bin': Binary(b("\x00\x01\x02\x03\x04"))}
+        ]
+
+        db.test.insert(docs)
+        reloaded_docs = json_util.loads(json_util.dumps(db.test.find()))
+        for doc in docs:
+            self.assertTrue(doc in reloaded_docs)
 
 if __name__ == "__main__":
     unittest.main()