bson for UUID <=> Binary w/ subtype 3, only on Python >= 2.5

Author: Mike Dirolf

pymongo/_cbsonmodule.c

@@ -20,9 +20,14 @@
  * BSON encoding and decoding.
  */
 
+
+#define _GNU_SOURCE
+#include <stdio.h>
+#include <time.h>
+#undef _GNU_SOURCE // avoid multiple define from Python.h
+
 #include <Python.h>
 #include <datetime.h>
-#include <time.h>
 
 static PyObject* CBSONError;
 static PyObject* InvalidName;
@@ -34,6 +39,7 @@ static PyObject* Code;
 static PyObject* ObjectId;
 static PyObject* DBRef;
 static PyObject* RECompile;
+static PyObject* UUID;
 
 #if PY_VERSION_HEX < 0x02050000 && !defined(PY_SSIZE_T_MIN)
 typedef int Py_ssize_t;
@@ -473,6 +479,33 @@ static int write_element_to_buffer(bson_buffer* buffer, int type_byte, PyObject*
  }
  }
  return 1;
+ } else if (UUID && PyObject_IsInstance(value, UUID)) {
+ // Just a special case of Binary above, but simpler to do as a separate case
+
+ // UUID is always 16 bytes, subtype 3
+ int length = 16;
+ const char subtype = 3;
+
+ PyObject* bytes;
+
+ *(buffer->buffer + type_byte) = 0x05;
+ if (!buffer_write_bytes(buffer, (const char*)&length, 4)) {
+ return 0;
+ }
+ if (!buffer_write_bytes(buffer, &subtype, 1)) {
+ return 0;
+ }
+
+ bytes = PyObject_GetAttrString(value, "bytes");
+ if (!bytes) {
+ return 0;
+ }
+ if (!buffer_write_bytes(buffer, PyString_AsString(bytes), length)) {
+ Py_DECREF(bytes);
+ return 0;
+ }
+ Py_DECREF(bytes);
+ return 1;
  } else if (PyObject_IsInstance(value, Code)) {
  int start_position,
  length_location,
@@ -1072,6 +1105,7 @@ static PyObject* get_value(const char* buffer, int* position, int type) {
 
  memcpy(&length, buffer + *position, 4);
  subtype = (unsigned char)buffer[*position + 4];
+
  if (subtype == 2) {
  data = PyString_FromStringAndSize(buffer + *position + 9, length - 4);
  } else {
@@ -1080,6 +1114,35 @@ static PyObject* get_value(const char* buffer, int* position, int type) {
  if (!data) {
  return NULL;
  }
+
+ if (subtype == 3 && UUID) { // Encode as UUID, not Binary
+ PyObject* kwargs;
+ PyObject* args = PyTuple_New(0);
+ if (!args) {
+ return NULL;
+ }
+ kwargs = PyDict_New();
+ if (!kwargs) {
+ Py_DECREF(args);
+ return NULL;
+ }
+
+ assert(length == 16); // UUID should always be 16 bytes
+
+ PyDict_SetItemString(kwargs, "bytes", data);
+ value = PyObject_Call(UUID, args, kwargs);
+
+ Py_DECREF(args);
+ Py_DECREF(kwargs);
+ Py_DECREF(data);
+ if (!value) {
+ return NULL;
+ }
+
+ *position += length + 5;
+ break;
+ }
+
  st = PyInt_FromLong(subtype);
  if (!st) {
  Py_DECREF(data);
@@ -1433,4 +1496,13 @@ PyMODINIT_FUNC init_cbson(void) {
  }
  RECompile = PyObject_GetAttrString(module, "compile");
  Py_DECREF(module);
+
+ module = PyImport_ImportModule("uuid");
+ if (!module) {
+ UUID = NULL;
+ PyErr_Clear();
+ } else {
+ UUID = PyObject_GetAttrString(module, "UUID");
+ Py_DECREF(module);
+ }
 }

pymongo/bson.py

@@ -256,6 +256,12 @@ def _get_binary(data):
  if length2 != length - 4:
  raise InvalidBSON("invalid binary (st 2) - lengths don't match!")
  length = length2
+ if subtype == 3:
+ try:
+ import uuid
+ return (uuid.UUID(bytes=data[:length]), data[length:])
+ except ImportError:
+ pass
  return (Binary(data[:length], subtype), data[length:])
 
 
@@ -374,6 +380,16 @@ def _element_to_bson(key, value, check_keys):
  name = _make_c_string(key)
  if isinstance(value, float):
  return "\x01" + name + struct.pack("<d", value)
+
+ # Use Binary w/ subtype 3 for UUID instances
+ try:
+ import uuid
+
+ if isinstance(value, uuid.UUID):
+ value = Binary(value.bytes, subtype=3)
+ except ImportError:
+ pass
+
  if isinstance(value, Binary):
  subtype = value.subtype
  if subtype == 2:

test/test_bson.py

@@ -21,6 +21,14 @@
 import sys
 import types
 
+try:
+ import uuid
+ should_test_uuid = True
+except ImportError:
+ should_test_uuid = False
+
+from nose.plugins.skip import SkipTest
+
 sys.path[0:0] = [""]
 
 import qcheck
@@ -213,5 +221,17 @@ def test_tuple(self):
  self.assertEqual({"tuple": [1, 2]},
  BSON.from_dict({"tuple": (1, 2)}).to_dict())
 
+ def test_uuid(self):
+ if not should_test_uuid:
+ raise SkipTest()
+
+ id = uuid.uuid4()
+ transformed_id = (BSON.from_dict({"id": id})).to_dict()["id"]
+
+ self.assert_(isinstance(transformed_id, uuid.UUID))
+ self.assertEqual(id, transformed_id)
+ self.assertNotEqual(uuid.uuid4(), transformed_id)
+
+
 if __name__ == "__main__":
  unittest.main()