From 3fbdc863cb8d96a5ce17add13f2dcd70b69a12bf Mon Sep 17 00:00:00 2001
From: Rico Hageman <info@ricohageman.nl>
Date: Tue, 22 Feb 2022 09:48:41 +0100
Subject: [PATCH] Include the causes when throwing a PyTypeError when argument
 parsing failed

---
 CHANGELOG.md                  |  2 +-
 src/impl_/extract_argument.rs | 19 ++++----------
 tests/test_pyfunction.rs      | 49 +++++++++++++++++++++++++----------
 3 files changed, 42 insertions(+), 28 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 56901c03..0e3a9e4e 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -72,7 +72,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   - `PyDate_FromTimestamp`
 - Deprecate the `gc` option for `pyclass` (e.g. `#[pyclass(gc)]`). Just implement a `__traverse__` `#[pymethod]`. [#2159](https://github.com/PyO3/pyo3/pull/2159)
 - The `ml_meth` field of `PyMethodDef` is now represented by the `PyMethodDefPointer` union [2166](https://github.com/PyO3/pyo3/pull/2166)
-- The error messages of `PyTypeError` contain now detailed information when argument parsing failed [2177](https://github.com/PyO3/pyo3/pull/2178)
+- The `PyTypeError` thrown when argument parsing failed now contains the nested causes [2177](https://github.com/PyO3/pyo3/pull/2178)
 
 ### Removed
 
diff --git a/src/impl_/extract_argument.rs b/src/impl_/extract_argument.rs
index 680e802a..9d5333a7 100644
--- a/src/impl_/extract_argument.rs
+++ b/src/impl_/extract_argument.rs
@@ -101,21 +101,12 @@ pub fn from_py_with_with_default<'py, T>(
 /// single string.)
 #[doc(hidden)]
 #[cold]
-pub fn argument_extraction_error(py: Python, arg_name: &str, mut error: PyErr) -> PyErr {
+pub fn argument_extraction_error(py: Python, arg_name: &str, error: PyErr) -> PyErr {
     if error.get_type(py) == PyTypeError::type_object(py) {
-        let mut reason = error
-            .value(py)
-            .str()
-            .unwrap_or_else(|_| PyString::new(py, ""))
-            .to_string();
-
-        while let Some(cause) = error.cause(py) {
-            reason.push_str(":\n\t");
-            reason.push_str(&cause.to_string());
-            error = cause
-        }
-
-        PyTypeError::new_err(format!("argument '{}': {}", arg_name, reason))
+        let remapped_error =
+            PyTypeError::new_err(format!("argument '{}': {}", arg_name, error.value(py)));
+        remapped_error.set_cause(py, error.cause(py));
+        remapped_error
     } else {
         error
     }
diff --git a/tests/test_pyfunction.rs b/tests/test_pyfunction.rs
index f54110ef..14920dc1 100644
--- a/tests/test_pyfunction.rs
+++ b/tests/test_pyfunction.rs
@@ -165,17 +165,20 @@ fn test_function_with_custom_conversion_error() {
     );
 }
 
+#[pyclass]
 #[derive(Debug, FromPyObject)]
 struct ValueClass {
+    #[pyo3(get)]
     value: usize,
 }
 
 #[pyfunction]
-fn conversion_error(str_arg: &str,
-                    int_arg: i64,
-                    tuple_arg: (&str, f64),
-                    option_arg: Option<i64>,
-                    struct_arg: Option<ValueClass>
+fn conversion_error(
+    str_arg: &str,
+    int_arg: i64,
+    tuple_arg: (&str, f64),
+    option_arg: Option<i64>,
+    struct_arg: Option<ValueClass>,
 ) {
     println!(
         "{:?} {:?} {:?} {:?} {:?}",
@@ -224,7 +227,7 @@ fn test_conversion_error() {
         PyTypeError,
         "argument 'option_arg': 'str' object cannot be interpreted as an integer"
     );
-    py_expect_exception!(
+    let exception = py_expect_exception!(
         py,
         conversion_error,
         "
@@ -232,11 +235,15 @@ class ValueClass:
     def __init__(self, value):
         self.value = value
 conversion_error('string1', -100, ('string2', 10.), None, ValueClass(\"no_expected_type\"))",
-        PyTypeError,
-        "argument 'struct_arg': failed to extract field ValueClass.value:\n\tTypeError: 'str' \
-        object cannot be interpreted as an integer"
+        PyTypeError
     );
-    py_expect_exception!(
+    assert_eq!(
+        extract_traceback(py, exception),
+        "TypeError: argument 'struct_arg': failed to \
+    extract field ValueClass.value: TypeError: 'str' object cannot be interpreted as an integer"
+    );
+
+    let exception = py_expect_exception!(
         py,
         conversion_error,
         "
@@ -244,10 +251,26 @@ class ValueClass:
     def __init__(self, value):
         self.value = value
 conversion_error('string1', -100, ('string2', 10.), None, ValueClass(-5))",
-        PyTypeError,
-        "argument 'struct_arg': failed to extract field ValueClass.value:\n\tOverflowError: can't \
-        convert negative int to unsigned"
+        PyTypeError
     );
+    assert_eq!(
+        extract_traceback(py, exception),
+        "TypeError: argument 'struct_arg': failed to \
+    extract field ValueClass.value: OverflowError: can't convert negative int to unsigned"
+    );
+}
+
+/// Helper function that concatenates the error message from
+/// each error in the traceback into a single string that can
+/// be tested.
+fn extract_traceback(py: Python, mut error: PyErr) -> String {
+    let mut error_msg = error.to_string();
+    while let Some(cause) = error.cause(py) {
+        error_msg.push_str(": ");
+        error_msg.push_str(&cause.to_string());
+        error = cause
+    }
+    error_msg
 }
 
 #[test]