address review

Author: blink1073

bson/_cbsonmodule.c

@@ -1644,6 +1644,68 @@ static int write_raw_doc(buffer_t buffer, PyObject* raw, PyObject* _raw_str) {
     return bytes_written;
 }
 
+
+/* Update Invalid Document error message to include doc.
+ */
+void handle_invalid_doc_error() {
+    PyObject *etype = NULL, *evalue = NULL, *etrace = NULL;
+    PyErr_Fetch(&etype, &evalue, &etrace);
+    PyObject *InvalidDocument = _error("InvalidDocument");
+    if (InvalidDocument == NULL) {
+        PyErr_Restore(etype, evalue, etrace);
+        return 0;
+    }
+
+    if (PyErr_GivenExceptionMatches(etype, InvalidDocument)) {
+
+        Py_DECREF(etype);
+        etype = InvalidDocument;
+
+        if (evalue) {
+            PyObject *msg = PyObject_Str(evalue);
+            Py_DECREF(evalue);
+
+            if (msg) {
+                // Prepend doc to the existing message
+                PyObject *dict_str = PyObject_Str(dict);
+                if (dict_str == NULL) {
+                    Py_DECREF(msg);
+                    return 0;
+                }
+                PyObject *dict_str_utf8 = PyUnicode_AsUTF8(dict);
+                Py_DECREF(dict_str);
+                if (dict_str_utf8 == NULL) {
+                    Py_DECREF(msg);
+                    return 0
+                }
+                PyObject *msg_utf8 = PyUnicode_AsUTF8(msg);
+                if (msg_utf8 == NULL) {
+                    Py_DECREF(msg);
+                    Py_DECREF(dict_str_utf8)
+                    return 0
+                }
+                PyObject *new_msg = PyUnicode_FromFormat("Invalid document %s | %s", dict_str_utf8, msg_utf8);
+                Py_DECREF(msg_utf8);
+
+                if (new_msg) {
+                    Py_DECREF(msg);
+                    evalue = new_msg;
+                }
+                else {
+                    Py_DECREF(new_msg);
+                    evalue = msg;
+                }
+            }
+        }
+        PyErr_NormalizeException(&etype, &evalue, &etrace);
+    }
+    else {
+        Py_DECREF(InvalidDocument);
+    }
+    PyErr_Restore(etype, evalue, etrace);
+}
+
+
 /* returns the number of bytes written or 0 on failure */
 int write_dict(PyObject* self, buffer_t buffer,
                PyObject* dict, unsigned char check_keys,
@@ -1743,46 +1805,10 @@ int write_dict(PyObject* self, buffer_t buffer,
         while (PyDict_Next(dict, &pos, &key, &value)) {
             if (!decode_and_write_pair(self, buffer, key, value,
                                     check_keys, options, top_level)) {
-                if (PyErr_Occurred()) {
-                    PyObject *etype = NULL, *evalue = NULL, *etrace = NULL;
-                    PyErr_Fetch(&etype, &evalue, &etrace);
-                    PyObject *InvalidDocument = _error("InvalidDocument");
-                    if (InvalidDocument == NULL) {
-                        return 0;
-                    }
-
-                    if (top_level && InvalidDocument && PyErr_GivenExceptionMatches(etype, InvalidDocument)) {
-
-                        Py_DECREF(etype);
-                        etype = InvalidDocument;
-
-                        if (evalue) {
-                            PyObject *msg = PyObject_Str(evalue);
-                            Py_DECREF(evalue);
-
-                            if (msg) {
-                                // Prepend doc to the existing message
-                                PyObject *dict_str = PyObject_Str(dict);
-                                if (dict_str ==  NULL) {
-                                    return 0;
-                                }
-                                PyObject *new_msg = PyUnicode_FromFormat("Invalid document %s | %s", PyUnicode_AsUTF8(dict_str), PyUnicode_AsUTF8(msg));
-                                Py_DECREF(dict_str);
-
-                                if (new_msg) {
-                                    evalue = new_msg;
-                                }
-                                else {
-                                evalue = msg;
-                                }
-                            }
-                        }
-                        PyErr_NormalizeException(&etype, &evalue, &etrace);
-                    }
-                    else {
-                        Py_DECREF(InvalidDocument);
-                    }
-                    PyErr_Restore(etype, evalue, etrace);
+                Py_DECREF(key);
+                Py_DECREF(value);
+                if (PyErr_Occurred() && top_level) {
+                    handle_invalid_doc_error();
                 }
                 return 0;
             }
@@ -1805,6 +1831,9 @@ int write_dict(PyObject* self, buffer_t buffer,
                 Py_DECREF(key);
                 Py_DECREF(value);
                 Py_DECREF(iter);
+                if (PyErr_Occurred() && top_level) {
+                    handle_invalid_doc_error();
+                }
                 return 0;
             }
             Py_DECREF(key);