From 6e7d0ff1efdbfcc2e2c8b1a7f525f058a0076026 Mon Sep 17 00:00:00 2001
From: Mike Dirolf <mike@10gen.com>
Date: Mon, 2 Feb 2009 10:08:01 -0500
Subject: [PATCH] support encoding/decoding binary w/ different subtypes

---
 pymongo/bson.py   | 20 +++++++++++---------
 test/test_bson.py |  4 ++++
 2 files changed, 15 insertions(+), 9 deletions(-)

diff --git a/pymongo/bson.py b/pymongo/bson.py
index 04578ba15..caf38951e 100644
--- a/pymongo/bson.py
+++ b/pymongo/bson.py
@@ -198,14 +198,14 @@ def _get_array(data):
 
 def _get_binary(data):
     (length, data) = _get_int(data)
-    subtype = data[0]
+    subtype = ord(data[0])
     data = data[1:]
-    if subtype != "\x02":
-        raise InvalidBSON("binary subtype %r is unsupported" % subtype)
-    (length2, data) = _get_int(data)
-    if length2 != length - 4:
-        raise InvalidBSON("invalid binary - lengths don't match!")
-    return (Binary(data[:length2]), data[length2:])
+    if subtype == 2:
+        (length2, data) = _get_int(data)
+        if length2 != length - 4:
+            raise InvalidBSON("invalid binary (st 2) - lengths don't match!")
+        length = length2
+    return (Binary(data[:length], subtype), data[length:])
 
 def _get_oid(data):
     oid = _shuffle_oid(data[:12])
@@ -290,8 +290,10 @@ def _element_to_bson(key, value):
     if isinstance(value, float):
         return "\x01" + name + struct.pack("<d", value)
     if isinstance(value, Binary):
-        length = len(value)
-        return "\x05" + name + struct.pack("<i", length + 4) + "\x02" + struct.pack("<i", length) + value
+        subtype = value.subtype()
+        if subtype == 2:
+            value = struct.pack("<i", len(value)) + value
+        return "\x05" + name + struct.pack("<i", len(value)) + chr(subtype) + value
     if isinstance(value, Code):
         cstring = _make_c_string(value, False)
         length = struct.pack("<i", len(cstring))
diff --git a/test/test_bson.py b/test/test_bson.py
index 6107449d1..05ce654a2 100644
--- a/test/test_bson.py
+++ b/test/test_bson.py
@@ -83,6 +83,8 @@ class TestBSON(unittest.TestCase):
                          "\x10\x00\x00\x00\x03\x6E\x6F\x6E\x65\x00\x05\x00\x00\x00\x00\x00")
         self.assertEqual(BSON.from_dict({"test": Binary("test")}),
                          "\x18\x00\x00\x00\x05\x74\x65\x73\x74\x00\x08\x00\x00\x00\x02\x04\x00\x00\x00\x74\x65\x73\x74\x00")
+        self.assertEqual(BSON.from_dict({"test": Binary("test", 128)}),
+                         "\x14\x00\x00\x00\x05\x74\x65\x73\x74\x00\x04\x00\x00\x00\x80\x74\x65\x73\x74\x00")
         self.assertEqual(BSON.from_dict({"test": None}),
                          "\x0B\x00\x00\x00\x0A\x74\x65\x73\x74\x00\x00")
         self.assertEqual(BSON.from_dict({"date": datetime.datetime(2007, 1, 8, 0, 30, 11)}),
@@ -110,6 +112,8 @@ class TestBSON(unittest.TestCase):
         helper({"false": False})
         helper({"an array": [1, True, 3.8, u"world"]})
         helper({"an object": {"test": u"something"}})
+        helper({"a binary": Binary("test", 100)})
+        helper({"another binary": Binary("test")})
         helper(SON([(u'test dst', datetime.datetime(1993, 4, 4, 2))]))
 
         def from_then_to_dict(dict):