diff --git a/bson/__init__.py b/bson/__init__.py index e866a99c8d..fc6efe0d59 100644 --- 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 index a66071c285..d91c7e0536 100644 --- 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 index 272b81d6ae..4a7f5424b1 100644 --- 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 index b431f700dc..e550b538d3 100644 --- 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):