1use pyo3::{Bound, PyAny, PyResult, prelude::PyAnyMethods, types::PyTypeMethods};
4use crate::{Node, dump};
5
6pub fn extract_operator_type<T>(
8 ob: &Bound<PyAny>,
9 attr_name: &str,
10 context: &str,
11) -> PyResult<String>
12where
13 T: std::fmt::Debug,
14{
15 let op = ob.getattr(attr_name).map_err(|_| {
16 pyo3::exceptions::PyAttributeError::new_err(
17 ob.error_message("<unknown>", format!("error getting {}", context))
18 )
19 })?;
20
21 let op_type = op.get_type().name().map_err(|_| {
22 pyo3::exceptions::PyTypeError::new_err(
23 ob.error_message(
24 "<unknown>",
25 format!("extracting type name for {}", context),
26 )
27 )
28 })?;
29
30 op_type.extract()
31}
32
33pub fn extract_binary_operands<L, R>(
35 ob: &Bound<PyAny>,
36 left_attr: &str,
37 right_attr: &str,
38 context: &str,
39) -> PyResult<(L, R)>
40where
41 L: for<'any, 'py> pyo3::FromPyObject<'any, 'py>,
42 R: for<'any, 'py> pyo3::FromPyObject<'any, 'py>,
43{
44 let left = ob.getattr(left_attr).map_err(|_| {
45 pyo3::exceptions::PyAttributeError::new_err(
46 ob.error_message("<unknown>", format!("error getting {} left operand", context))
47 )
48 })?;
49
50 let right = ob.getattr(right_attr).map_err(|_| {
51 pyo3::exceptions::PyAttributeError::new_err(
52 ob.error_message("<unknown>", format!("error getting {} right operand", context))
53 )
54 })?;
55
56 let left = left.extract().map_err(Into::into).map_err(|e: pyo3::PyErr| {
57 let e: pyo3::PyErr = e.into();
58 pyo3::exceptions::PyValueError::new_err(
59 format!("Failed to extract {} left operand: {}", context, e)
60 )
61 })?;
62
63 let right = right.extract().map_err(Into::into).map_err(|e: pyo3::PyErr| {
64 let e: pyo3::PyErr = e.into();
65 pyo3::exceptions::PyValueError::new_err(
66 format!("Failed to extract {} right operand: {}", context, e)
67 )
68 })?;
69
70 Ok((left, right))
71}
72
73pub fn extract_list<T>(
75 ob: &Bound<PyAny>,
76 attr_name: &str,
77 context: &str,
78) -> PyResult<Vec<T>>
79where
80 T: for<'any, 'py> pyo3::FromPyObject<'any, 'py>,
81{
82 let list_obj = ob.getattr(attr_name).map_err(|_| {
83 pyo3::exceptions::PyAttributeError::new_err(
84 ob.error_message("<unknown>", format!("error getting {} list", context))
85 )
86 })?;
87
88 list_obj.extract().map_err(Into::into).map_err(|e: pyo3::PyErr| {
89 pyo3::exceptions::PyValueError::new_err(
90 format!("Failed to extract {} list: {}", context, e)
91 )
92 })
93}
94
95pub fn extract_optional<T>(
97 ob: &Bound<PyAny>,
98 attr_name: &str,
99) -> Option<T>
100where
101 T: for<'any, 'py> pyo3::FromPyObject<'any, 'py>,
102{
103 ob.getattr(attr_name)
104 .ok()
105 .and_then(|attr| attr.extract().ok())
106}
107
108pub fn extract_position_info(ob: &Bound<PyAny>) -> (Option<usize>, Option<usize>, Option<usize>, Option<usize>) {
110 (
111 extract_optional(&ob, "lineno"),
112 extract_optional(&ob, "col_offset"),
113 extract_optional(&ob, "end_lineno"),
114 extract_optional(&ob, "end_col_offset"),
115 )
116}
117
118pub trait ExtractFromPython<'a>: Sized {
120 fn extract_with_context(ob: &Bound<'a, PyAny>, context: &str) -> PyResult<Self>;
122}
123
124impl<'a, T> ExtractFromPython<'a> for T
125where
126 T: pyo3::conversion::FromPyObjectOwned<'a>,
127{
128 fn extract_with_context(ob: &Bound<'a, PyAny>, context: &str) -> PyResult<Self> {
129 ob.extract().map_err(Into::into).map_err(|e: pyo3::PyErr| {
130 pyo3::exceptions::PyValueError::new_err(
131 format!("Failed to extract {} from Python: {} (object: {})",
132 context, e, dump(&ob, None).unwrap_or_else(|_| "unknown".to_string()))
133 )
134 })
135 }
136}
137
138pub fn log_extraction(ob: &Bound<PyAny>, context: &str) {
140 if tracing::enabled!(tracing::Level::DEBUG) {
141 match dump(&ob, None) {
142 Ok(dump_str) => tracing::debug!("Extracting {}: {}", context, dump_str),
143 Err(_) => tracing::debug!("Extracting {} (dump failed)", context),
144 }
145 }
146}
147
148pub fn extraction_error(context: &str, details: &str) -> pyo3::PyErr {
150 pyo3::exceptions::PyValueError::new_err(
151 format!("Failed to extract {}: {}", context, details)
152 )
153}
154
155pub fn safe_ident(name: &str) -> proc_macro2::Ident {
161 use quote::format_ident;
162 match name {
163 "self" => format_ident!("{}", name),
164 "Self" | "super" | "crate" => format_ident!("{}_", name),
165 "as" | "break" | "const" | "continue" | "dyn" | "else" | "enum" | "extern" | "false"
167 | "fn" | "for" | "if" | "impl" | "in" | "let" | "loop" | "match" | "mod" | "move"
168 | "mut" | "pub" | "ref" | "return" | "static" | "struct" | "trait" | "true" | "type"
169 | "unsafe" | "use" | "where" | "while" | "async" | "await" | "abstract" | "become"
170 | "box" | "do" | "final" | "macro" | "override" | "priv" | "try" | "typeof"
171 | "unsized" | "virtual" | "yield" | "gen" => format_ident!("r#{}", name),
172 _ => format_ident!("{}", name),
173 }
174}
175
176pub fn extraction_failure(
180 kind: &str,
181 ob: &Bound<PyAny>,
182 cause: impl std::fmt::Display,
183) -> pyo3::PyErr {
184 let position = match (ob.lineno(), ob.col_offset()) {
185 (Some(line), Some(col)) => format!(" at line {}, column {}", line, col + 1),
186 (Some(line), None) => format!(" at line {}", line),
187 _ => String::new(),
188 };
189 pyo3::exceptions::PyValueError::new_err(format!(
190 "cannot convert {}{}: {}",
191 kind, position, cause
192 ))
193}
194
195#[cfg(test)]
196mod tests {
197 use super::*;
198 use pyo3::Python;
199
200 #[test]
201 fn test_extract_optional() {
202 Python::attach(|py| {
203 use std::ffi::CString;
204 let code = CString::new("42").unwrap();
206 let obj = py.eval(&code, None, None).unwrap();
207
208 let missing: Option<String> = extract_optional(&obj, "missing_attr");
210 assert_eq!(missing, None);
211
212 });
215 }
216
217 #[test]
218 fn test_log_extraction() {
219 Python::attach(|py| {
220 use std::ffi::CString;
221 let code = CString::new("42").unwrap();
222 let obj = py.eval(&code, None, None).unwrap();
223
224 log_extraction(&obj, "test object");
226 });
227 }
228
229 #[test]
230 fn test_extraction_error() {
231 let error = extraction_error("test context", "test details");
232 let error_string = format!("{}", error);
233 assert!(error_string.contains("test context"));
234 assert!(error_string.contains("test details"));
235 }
236}
237
238pub fn get_attr_with_context<'a>(
242 ob: &Bound<'a, PyAny>,
243 attr_name: &str,
244 context: &str,
245) -> PyResult<Bound<'a, PyAny>> {
246 ob.getattr(attr_name).map_err(|e| {
247 let type_name = ob.get_type().name()
248 .map(|s| s.to_string())
249 .unwrap_or_else(|_| "<unknown>".to_string());
250 let enhanced_msg = format!(
251 "Failed to get attribute '{}' from {} ({}): {}",
252 attr_name,
253 context,
254 type_name,
255 e
256 );
257 pyo3::exceptions::PyAttributeError::new_err(enhanced_msg)
258 })
259}
260
261pub fn extract_with_context<'py, T>(
263 value: &Bound<'py, PyAny>,
264 context: &str,
265 attr_name: &str,
266) -> PyResult<T>
267where
268 T: pyo3::conversion::FromPyObjectOwned<'py>,
269{
270 value.extract().map_err(Into::into).map_err(|e: pyo3::PyErr| {
271 let type_name = value.get_type().name()
272 .map(|s| s.to_string())
273 .unwrap_or_else(|_| "<unknown>".to_string());
274 let enhanced_msg = format!(
275 "Failed to extract {} for attribute '{}': {}. Expected type: {}, got: {}",
276 context,
277 attr_name,
278 e,
279 std::any::type_name::<T>(),
280 type_name
281 );
282 pyo3::exceptions::PyTypeError::new_err(enhanced_msg)
283 })
284}
285
286pub fn extract_required_attr<'py, T>(
288 ob: &Bound<'py, PyAny>,
289 attr_name: &str,
290 context: &str,
291) -> PyResult<T>
292where
293 T: pyo3::conversion::FromPyObjectOwned<'py>,
294{
295 let attr = get_attr_with_context(&ob, attr_name, context)?;
296 extract_with_context(&attr, context, attr_name)
297}