From a31cac08896138abe6bce06647c94c8f56f480fc Mon Sep 17 00:00:00 2001
From: Ross Lawley <ross.lawley@gmail.com>
Date: Wed, 5 Dec 2012 16:23:33 +0000
Subject: [PATCH] json_util fix format for binary type PYTHON-443

---
 bson/json_util.py      | 15 ++++++++++-----
 test/test_json_util.py | 40 ++++++++++++++++++++++++++++++++++++----
 2 files changed, 46 insertions(+), 9 deletions(-)

diff --git a/bson/json_util.py b/bson/json_util.py
index da4d06554..55e002d0f 100644
--- a/bson/json_util.py
+++ b/bson/json_util.py
@@ -32,14 +32,14 @@ Example usage (serialization)::
    ...        {'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"}}]'
+   '[{"foo": [1, 2]}, {"bar": {"hello": "world"}}, {"code": {"$scope": {}, "$code": "function x() { return 1; }"}}, {"bin": {"$type": "00", "$binary": "AAECAwQ="}}]'
 
 Example usage (deserialization)::
 
 .. 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"}}]')
+   >>> loads('[{"foo": [1, 2]}, {"bar": {"hello": "world"}}, {"code": {"$scope": {}, "$code": "function x() { return 1; }"}}, {"bin": {"$type": "00", "$binary": "AAECAwQ="}}]')
    [{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)}]
 
 Alternatively, you can manually pass the `default` to :func:`json.dumps`.
@@ -151,7 +151,12 @@ def object_hook(dct):
     if "$maxKey" in dct:
         return MaxKey()
     if "$binary" in dct:
-        return Binary(base64.b64decode(dct["$binary"].encode()), dct["$type"])
+        if isinstance(dct["$type"], int):
+            dct["$type"] = "%02x" % dct["$type"]
+        subtype = int(dct["$type"], 16)
+        if subtype >= 0xffffff80:  # Handle mongoexport values
+            subtype = int(dct["$type"][6:], 16)
+        return Binary(base64.b64decode(dct["$binary"].encode()), subtype)
     if "$code" in dct:
         return Code(dct["$code"], dct.get("$scope"))
     if bson.has_uuid() and "$uuid" in dct:
@@ -189,10 +194,10 @@ def default(obj):
         return {'$code': "%s" % obj, '$scope': obj.scope}
     if isinstance(obj, Binary):
         return {'$binary': base64.b64encode(obj).decode(),
-                '$type': obj.subtype}
+                '$type': "%02x" % obj.subtype}
     if PY3 and isinstance(obj, binary_type):
         return {'$binary': base64.b64encode(obj).decode(),
-                '$type': 0}
+                '$type': "00"}
     if bson.has_uuid() and isinstance(obj, bson.uuid.UUID):
         return {"$uuid": obj.hex}
     raise TypeError("%r is not JSON serializable" % obj)
diff --git a/test/test_json_util.py b/test/test_json_util.py
index df1cca448..21ed9900f 100644
--- a/test/test_json_util.py
+++ b/test/test_json_util.py
@@ -26,7 +26,7 @@ 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.binary import Binary, MD5_SUBTYPE, USER_DEFINED_SUBTYPE
 from bson.code import Code
 from bson.dbref import DBRef
 from bson.max_key import MaxKey
@@ -101,10 +101,42 @@ class TestJsonUtil(unittest.TestCase):
                             'f47ac10b-58cc-4372-a567-0e02b2c3d479')})
 
     def test_binary(self):
-        self.round_trip({"bin": Binary(b("\x00\x01\x02\x03\x04"))})
-        self.round_trip({
+        bin_type_dict = {"bin": Binary(b("\x00\x01\x02\x03\x04"))}
+        md5_type_dict = {
             "md5": Binary(b(' n7\x18\xaf\t/\xd1\xd1/\x80\xca\xe7q\xcc\xac'),
-                MD5_SUBTYPE)})
+                MD5_SUBTYPE)}
+        custom_type_dict = {"custom": Binary(b("hello"), USER_DEFINED_SUBTYPE)}
+
+        self.round_trip(bin_type_dict)
+        self.round_trip(md5_type_dict)
+        self.round_trip(custom_type_dict)
+
+        # PYTHON-443 ensure old type formats are supported
+        json_bin_dump = json_util.dumps(bin_type_dict)
+        self.assertTrue('"$type": "00"' in json_bin_dump)
+        self.assertEqual(bin_type_dict,
+            json_util.loads('{"bin": {"$type": 0, "$binary": "AAECAwQ="}}'))
+
+        json_bin_dump = json_util.dumps(md5_type_dict)
+        self.assertTrue('"$type": "05"' in json_bin_dump)
+        self.assertEqual(md5_type_dict,
+            json_util.loads('{"md5": {"$type": 5, "$binary":'
+                            ' "IG43GK8JL9HRL4DK53HMrA=="}}'))
+
+        json_bin_dump = json_util.dumps(custom_type_dict)
+        self.assertTrue('"$type": "80"' in json_bin_dump)
+        self.assertEqual(custom_type_dict,
+            json_util.loads('{"custom": {"$type": 128, "$binary":'
+                            ' "aGVsbG8="}}'))
+
+        # Handle mongoexport where subtype >= 128
+        self.assertEqual(128,
+            json_util.loads('{"custom": {"$type": "ffffff80", "$binary":'
+                            ' "aGVsbG8="}}')['custom'].subtype)
+
+        self.assertEqual(255,
+            json_util.loads('{"custom": {"$type": "ffffffff", "$binary":'
+                            ' "aGVsbG8="}}')['custom'].subtype)
 
     def test_code(self):
         self.round_trip({"code": Code("function x() { return 1; }")})