From b4388759260a107a67c7c74b52742e9569c3df84 Mon Sep 17 00:00:00 2001
From: LJ <linghua@cocoindex.io>
Date: Mon, 21 Apr 2025 15:21:28 -0700
Subject: [PATCH] feat(py-err): attach Python exception backtrace when
 propagating to Rust

---
 src/ops/py_factory.rs | 47 ++++++++++++++++++++++++++-----------------
 src/py/convert.rs     |  4 ++--
 src/py/mod.rs         | 19 +++++++++++++++++
 3 files changed, 50 insertions(+), 20 deletions(-)

diff --git a/src/ops/py_factory.rs b/src/ops/py_factory.rs
index 3eec3ac8..adabc6ba 100644
--- a/src/ops/py_factory.rs
+++ b/src/ops/py_factory.rs
@@ -12,7 +12,7 @@ use pythonize::pythonize;
 use crate::{
     base::{schema, value},
     builder::plan,
-    py,
+    py::{self, FromPyResult},
 };
 use anyhow::{anyhow, Result};
 
@@ -74,14 +74,17 @@ impl PyFunctionExecutor {
             Some(kwargs)
         };
 
-        let result = self.py_function_executor.call(
-            py,
-            PyTuple::new(py, args.into_iter())?,
-            kwargs
-                .map(|kwargs| -> Result<_> { Ok(kwargs.into_py_dict(py)?) })
-                .transpose()?
-                .as_ref(),
-        )?;
+        let result = self
+            .py_function_executor
+            .call(
+                py,
+                PyTuple::new(py, args.into_iter())?,
+                kwargs
+                    .map(|kwargs| -> Result<_> { Ok(kwargs.into_py_dict(py)?) })
+                    .transpose()?
+                    .as_ref(),
+            )
+            .from_py_result(py)?;
         Ok(result.into_bound(py))
     }
 }
@@ -99,8 +102,9 @@ impl SimpleFunctionExecutor for Arc<PyFunctionExecutor> {
                 result_coro,
             )?)
         })?;
-        let result = result_fut.await?;
+        let result = result_fut.await;
         Python::with_gil(|py| -> Result<_> {
+            let result = result.from_py_result(py)?;
             Ok(py::value_from_py_object(
                 &self.result_type.typ,
                 &result.into_bound(py),
@@ -156,11 +160,14 @@ impl SimpleFunctionFactory for PyFunctionFactory {
                     .iter()
                     .map(|(name, _)| PyString::new(py, name).unbind())
                     .collect::<Vec<_>>();
-                let result = self.py_function_factory.call(
-                    py,
-                    PyTuple::new(py, args.into_iter())?,
-                    Some(&kwargs.into_py_dict(py)?),
-                )?;
+                let result = self
+                    .py_function_factory
+                    .call(
+                        py,
+                        PyTuple::new(py, args.into_iter())?,
+                        Some(&kwargs.into_py_dict(py)?),
+                    )
+                    .from_py_result(py)?;
                 let (result_type, executor) = result
                     .extract::<(crate::py::Pythonized<schema::EnrichedValueType>, Py<PyAny>)>(py)?;
                 Ok((
@@ -181,7 +188,9 @@ impl SimpleFunctionFactory for PyFunctionFactory {
                     .clone();
                 let (prepare_fut, enable_cache, behavior_version) =
                     Python::with_gil(|py| -> anyhow::Result<_> {
-                        let prepare_coro = executor.call_method(py, "prepare", (), None)?;
+                        let prepare_coro = executor
+                            .call_method(py, "prepare", (), None)
+                            .from_py_result(py)?;
                         let prepare_fut = pyo3_async_runtimes::into_future_with_locals(
                             &pyo3_async_runtimes::TaskLocals::new(
                                 py_exec_ctx.event_loop.bind(py).clone(),
@@ -189,10 +198,12 @@ impl SimpleFunctionFactory for PyFunctionFactory {
                             prepare_coro.into_bound(py),
                         )?;
                         let enable_cache = executor
-                            .call_method(py, "enable_cache", (), None)?
+                            .call_method(py, "enable_cache", (), None)
+                            .from_py_result(py)?
                             .extract::<bool>(py)?;
                         let behavior_version = executor
-                            .call_method(py, "behavior_version", (), None)?
+                            .call_method(py, "behavior_version", (), None)
+                            .from_py_result(py)?
                             .extract::<Option<u32>>(py)?;
                         Ok((prepare_fut, enable_cache, behavior_version))
                     })?;
diff --git a/src/py/convert.rs b/src/py/convert.rs
index afd01245..6410f1a1 100644
--- a/src/py/convert.rs
+++ b/src/py/convert.rs
@@ -1,3 +1,4 @@
+use bytes::Bytes;
 use pyo3::types::{PyList, PyTuple};
 use pyo3::IntoPyObjectExt;
 use pyo3::{exceptions::PyException, prelude::*};
@@ -6,8 +7,7 @@ use serde::de::DeserializeOwned;
 use serde::Serialize;
 use std::collections::BTreeMap;
 use std::ops::Deref;
-use std::sync::Arc; 
-use bytes::Bytes;
+use std::sync::Arc;
 
 use super::IntoPyResult;
 use crate::base::{schema, value};
diff --git a/src/py/mod.rs b/src/py/mod.rs
index 35eda644..e42fb999 100644
--- a/src/py/mod.rs
+++ b/src/py/mod.rs
@@ -12,6 +12,7 @@ use crate::setup;
 use pyo3::{exceptions::PyException, prelude::*};
 use pyo3_async_runtimes::tokio::future_into_py;
 use std::collections::btree_map;
+use std::fmt::Write;
 
 mod convert;
 pub use convert::*;
@@ -26,6 +27,24 @@ impl PythonExecutionContext {
     }
 }
 
+pub trait FromPyResult<T> {
+    fn from_py_result(self, py: Python<'_>) -> anyhow::Result<T>;
+}
+
+impl<T> FromPyResult<T> for Result<T, PyErr> {
+    fn from_py_result(self, py: Python<'_>) -> anyhow::Result<T> {
+        match self {
+            Ok(value) => Ok(value),
+            Err(err) => {
+                let mut err_str = format!("Error calling Python function: {}", err);
+                if let Some(tb) = err.traceback(py) {
+                    write!(&mut err_str, "\n{}", tb.format()?)?;
+                }
+                Err(anyhow::anyhow!(err_str))
+            }
+        }
+    }
+}
 pub trait IntoPyResult<T> {
     fn into_py_result(self) -> PyResult<T>;
 }