PYTHON-1982 Update Invalid Document error message to include doc (#1854)

Co-authored-by: Navjot Singh <[email protected]> Co-authored-by: Navjot Singh <[email protected]> Co-authored-by: Steven Silvester <[email protected]>
mongodb · Dec 3, 2024 · fdcbe2e · fdcbe2e
1 parent f45b35a
commit fdcbe2e
Show file tree

Hide file tree

Showing 4 changed files with 53 additions and 1 deletion.
diff --git a/bson/__init__.py b/bson/__init__.py
@@ -1006,7 +1006,10 @@ def _dict_to_bson(
             elements.append(_name_value_to_bson(b"_id\x00", doc["_id"], check_keys, opts))
         for key, value in doc.items():
             if not top_level or key != "_id":
-                elements.append(_element_to_bson(key, value, check_keys, opts))
+                try:
+                    elements.append(_element_to_bson(key, value, check_keys, opts))
+                except InvalidDocument as err:
+                    raise InvalidDocument(f"Invalid document {doc} | {err}") from err
     except AttributeError:
         raise TypeError(f"encoder expected a mapping type but got: {doc!r}") from None
 

diff --git a/bson/_cbsonmodule.c b/bson/_cbsonmodule.c
@@ -1743,6 +1743,41 @@ 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 (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);
+                                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);
+                }
                 return 0;
             }
         }

diff --git a/doc/contributors.rst b/doc/contributors.rst
@@ -102,3 +102,4 @@ The following is a list of people who have contributed to
 - Ivan Lukyanchikov (ilukyanchikov)
 - Terry Patterson
 - Romain Morotti
+- Navjot Singh (navjots18)
diff --git a/test/test_bson.py b/test/test_bson.py
@@ -1099,6 +1099,19 @@ def __repr__(self):
         ):
             encode({"t": Wrapper(1)})
 
+    def test_doc_in_invalid_document_error_message(self):
+        class Wrapper:
+            def __init__(self, val):
+                self.val = val
+
+            def __repr__(self):
+                return repr(self.val)
+
+        self.assertEqual("1", repr(Wrapper(1)))
+        doc = {"t": Wrapper(1)}
+        with self.assertRaisesRegex(InvalidDocument, f"Invalid document {doc}"):
+            encode(doc)
+
 
 class TestCodecOptions(unittest.TestCase):
     def test_document_class(self):