Skip to main content

python_ast/
parser_utils.rs

1/// Generic utilities for parsing Python AST objects with consistent error handling.
2
3use pyo3::{Bound, PyAny, PyResult, prelude::PyAnyMethods, types::PyTypeMethods};
4use crate::{Node, dump};
5
6/// Generic function for extracting Python operator types with consistent error handling.
7pub 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
33/// Generic function for extracting operands from binary operations.
34pub 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
73/// Generic function for extracting lists of items with error handling.
74pub 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
95/// Generic function to safely extract optional attributes.
96pub 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
108/// Generic function to extract position information from AST nodes.
109pub 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
118/// Trait for types that can be extracted from Python with improved error messages.
119pub trait ExtractFromPython<'a>: Sized {
120    /// Extract from Python object with context for better error messages.
121    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
138/// Utility function for consistent logging during Python object extraction.
139pub 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
148/// Helper function to create standardized error messages for failed extractions.
149pub 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
155/// Build an identifier for generated code from a Python name, escaping Rust
156/// keywords so identifiers like `type`, `match`, or `move` don't produce
157/// invalid Rust. Keywords that cannot be raw identifiers (`self`, `Self`,
158/// `super`, `crate`) get a trailing underscore instead — except `self`,
159/// which is left as-is so method receivers keep working.
160pub 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        // Strict and reserved keywords, all legal as raw identifiers.
166        "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
176/// Build a PyErr for a failed AST-node conversion, carrying the node's source
177/// position so the error can be reported against the user's Python code
178/// instead of panicking with an internal dump.
179pub 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            // Create a simple Python integer which has some attributes
205            let code = CString::new("42").unwrap();
206            let obj = py.eval(&code, None, None).unwrap();
207            
208            // Try to extract something that probably won't exist 
209            let missing: Option<String> = extract_optional(&obj, "missing_attr");
210            assert_eq!(missing, None);
211            
212            // This test mainly checks that extract_optional doesn't panic
213            // and properly returns None for missing attributes
214        });
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            // Should not panic
225            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
238/// Enhanced error handling utilities for parsing Python AST objects
239
240/// Get an attribute from a Python object with better error messaging
241pub 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
261/// Extract a value from PyAny with better error messaging
262pub 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
286/// Extract a required attribute with enhanced error messaging  
287pub 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}