raise InvalidDocument when trying to encode a document with non-string keys

Mike Dirolf · Mike Dirolf · commit 9bbeea47825c · 2009-11-10T11:15:24.000-05:00
diff --git a/pymongo/_cbsonmodule.c b/pymongo/_cbsonmodule.c
@@ -846,7 +846,7 @@ static int write_pair(bson_buffer* buffer, const char* name, Py_ssize_t name_len
     int type_byte;
 
     /* Don't write any _id elements unless we're explicitly told to -
-     * _id has to be written first so we write do so, but don't bother
+     * _id has to be written first so we do so, but don't bother
      * deleting it from the dictionary being written. */
     if (!allow_id && strcmp(name, "_id") == 0) {
         return 1;
@@ -868,6 +868,37 @@ static int write_pair(bson_buffer* buffer, const char* name, Py_ssize_t name_len
     return 1;
 }
 
+static int decode_and_write_pair(bson_buffer* buffer, PyObject* key,
+                                 PyObject* value, unsigned char check_keys) {
+    PyObject* encoded;
+    if (PyUnicode_Check(key)) {
+        encoded = PyUnicode_AsUTF8String(key);
+        if (!encoded) {
+            return 0;
+        }
+    } else if (PyString_Check(key)) {
+        encoded = key;
+        Py_INCREF(encoded);
+    } else {
+        PyObject* errmsg = PyString_FromString("documents must have only string keys, key was ");
+        PyObject* repr = PyObject_Repr(key);
+        PyString_ConcatAndDel(&errmsg, repr);
+        PyErr_SetString(InvalidDocument, PyString_AsString(errmsg));
+        Py_DECREF(errmsg);
+        return 0;
+    }
+
+    /* Don't allow writing _id here - it was already written. */
+    if (!write_pair(buffer, PyString_AsString(encoded),
+                    PyString_Size(encoded), value, check_keys, 0)) {
+        Py_DECREF(encoded);
+        return 0;
+    }
+
+    Py_DECREF(encoded);
+    return 1;
+}
+
 static int write_son(bson_buffer* buffer, PyObject* dict, int start_position,
                      int length_location, unsigned char check_keys) {
     PyObject* keys = PyObject_CallMethod(dict, "keys", NULL);
@@ -880,33 +911,16 @@ static int write_son(bson_buffer* buffer, PyObject* dict, int start_position,
     for(i = 0; i < items; i++) {
         PyObject* key;
         PyObject* value;
-        PyObject* encoded;
 
         key = PyList_GetItem(keys, i);
         if (!key) {
             Py_DECREF(keys);
             return 0;
         }
         value = PyDict_GetItem(dict, key);
-        if (!value) {
-            Py_DECREF(keys);
-            return 0;
-        }
-        if (PyUnicode_CheckExact(key)) {
-            encoded = PyUnicode_AsUTF8String(key);
-            if (!encoded) {
-                Py_DECREF(keys);
-                return 0;
-            }
-        } else {
-            encoded = key;
-            Py_INCREF(encoded);
-        }
-        /* Don't allow writing _id here - it was written above. */
-        if (!write_pair(buffer, PyString_AsString(encoded),
-                        PyString_Size(encoded), value, check_keys, 0)) {
+        if (!value ||
+            !decode_and_write_pair(buffer, key, value, check_keys)) {
             Py_DECREF(keys);
-            Py_DECREF(encoded);
             return 0;
         }
     }
@@ -950,20 +964,7 @@ static int write_dict(bson_buffer* buffer, PyObject* dict, unsigned char check_k
         Py_ssize_t pos = 0;
 
         while (PyDict_Next(dict, &pos, &key, &value)) {
-            PyObject* encoded;
-            if (PyUnicode_CheckExact(key)) {
-                encoded = PyUnicode_AsUTF8String(key);
-                if (!encoded) {
-                    return 0;
-                }
-            } else {
-                encoded = key;
-                Py_INCREF(encoded);
-            }
-            /* Don't allow writing _id here - it was written above. */
-            if (!write_pair(buffer, PyString_AsString(encoded),
-                            PyString_Size(encoded), value, check_keys, 0)) {
-                Py_DECREF(encoded);
+            if (!decode_and_write_pair(buffer, key, value, check_keys)) {
                 return 0;
             }
         }
diff --git a/pymongo/bson.py b/pymongo/bson.py
@@ -373,6 +373,9 @@ def _bson_to_dict(data):
 
 
 def _element_to_bson(key, value, check_keys):
+    if not isinstance(key, (str, unicode)):
+        raise InvalidDocument("documents must have only string keys, key was %r" % key)
+
     if check_keys:
         if key.startswith("$"):
             raise InvalidName("key %r must not start with '$'" % key)
diff --git a/test/test_bson.py b/test/test_bson.py
@@ -238,6 +238,9 @@ def test_unicode_regex(self):
         regex = re.compile(u'revisi\xf3n')
         BSON.from_dict({"regex": regex}).to_dict()
 
+    def test_non_string_keys(self):
+        self.assertRaises(InvalidDocument, BSON.from_dict, {8.9: "test"})
+
 # TODO this test doesn't pass w/ C extension
 #
 # timegm doesn't handle years < 1900 (negative), at least on OS X

Original file line number	Diff line number	Diff line change
`@@ -238,6 +238,9 @@ def test_unicode_regex(self):`
`238`	`238`	`regex = re.compile(u'revisi\xf3n')`
`239`	`239`	`BSON.from_dict({"regex": regex}).to_dict()`
`240`	`240`
	`241`	`+ def test_non_string_keys(self):`
	`242`	`+ self.assertRaises(InvalidDocument, BSON.from_dict, {8.9: "test"})`
	`243`	`+`
`241`	`244`	`# TODO this test doesn't pass w/ C extension`
`242`	`245`	`#`
`243`	`246`	`# timegm doesn't handle years < 1900 (negative), at least on OS X`