python-ast 1.1.0

/// Generic utilities for parsing Python AST objects with consistent error handling.

use pyo3::{Bound, PyAny, PyResult, prelude::PyAnyMethods, types::PyTypeMethods};
use crate::{Node, dump};

/// Generic function for extracting Python operator types with consistent error handling.
pub fn extract_operator_type<T>(
    ob: &Bound<PyAny>,
    attr_name: &str,
    context: &str,
) -> PyResult<String>
where
    T: std::fmt::Debug,
{
    let op = ob.getattr(attr_name).map_err(|_| {
        pyo3::exceptions::PyAttributeError::new_err(
            ob.error_message("<unknown>", format!("error getting {}", context))
        )
    })?;

    let op_type = op.get_type().name().map_err(|_| {
        pyo3::exceptions::PyTypeError::new_err(
            ob.error_message(
                "<unknown>",
                format!("extracting type name for {}", context),
            )
        )
    })?;

    op_type.extract()
}

/// Generic function for extracting operands from binary operations.
pub fn extract_binary_operands<L, R>(
    ob: &Bound<PyAny>,
    left_attr: &str,
    right_attr: &str,
    context: &str,
) -> PyResult<(L, R)>
where
    L: for<'any, 'py> pyo3::FromPyObject<'any, 'py>,
    R: for<'any, 'py> pyo3::FromPyObject<'any, 'py>,
{
    let left = ob.getattr(left_attr).map_err(|_| {
        pyo3::exceptions::PyAttributeError::new_err(
            ob.error_message("<unknown>", format!("error getting {} left operand", context))
        )
    })?;

    let right = ob.getattr(right_attr).map_err(|_| {
        pyo3::exceptions::PyAttributeError::new_err(
            ob.error_message("<unknown>", format!("error getting {} right operand", context))
        )
    })?;

    let left = left.extract().map_err(Into::into).map_err(|e: pyo3::PyErr| {
        let e: pyo3::PyErr = e.into();
        pyo3::exceptions::PyValueError::new_err(
            format!("Failed to extract {} left operand: {}", context, e)
        )
    })?;

    let right = right.extract().map_err(Into::into).map_err(|e: pyo3::PyErr| {
        let e: pyo3::PyErr = e.into();
        pyo3::exceptions::PyValueError::new_err(
            format!("Failed to extract {} right operand: {}", context, e)
        )
    })?;

    Ok((left, right))
}

/// Generic function for extracting lists of items with error handling.
pub fn extract_list<T>(
    ob: &Bound<PyAny>,
    attr_name: &str,
    context: &str,
) -> PyResult<Vec<T>>
where
    T: for<'any, 'py> pyo3::FromPyObject<'any, 'py>,
{
    let list_obj = ob.getattr(attr_name).map_err(|_| {
        pyo3::exceptions::PyAttributeError::new_err(
            ob.error_message("<unknown>", format!("error getting {} list", context))
        )
    })?;

    list_obj.extract().map_err(Into::into).map_err(|e: pyo3::PyErr| {
        pyo3::exceptions::PyValueError::new_err(
            format!("Failed to extract {} list: {}", context, e)
        )
    })
}

/// Generic function to safely extract optional attributes.
pub fn extract_optional<T>(
    ob: &Bound<PyAny>,
    attr_name: &str,
) -> Option<T>
where
    T: for<'any, 'py> pyo3::FromPyObject<'any, 'py>,
{
    ob.getattr(attr_name)
        .ok()
        .and_then(|attr| attr.extract().ok())
}

/// Generic function to extract position information from AST nodes.
pub fn extract_position_info(ob: &Bound<PyAny>) -> (Option<usize>, Option<usize>, Option<usize>, Option<usize>) {
    (
        extract_optional(&ob, "lineno"),
        extract_optional(&ob, "col_offset"),
        extract_optional(&ob, "end_lineno"),
        extract_optional(&ob, "end_col_offset"),
    )
}

/// Trait for types that can be extracted from Python with improved error messages.
pub trait ExtractFromPython<'a>: Sized {
    /// Extract from Python object with context for better error messages.
    fn extract_with_context(ob: &Bound<'a, PyAny>, context: &str) -> PyResult<Self>;
}

impl<'a, T> ExtractFromPython<'a> for T
where
    T: pyo3::conversion::FromPyObjectOwned<'a>,
{
    fn extract_with_context(ob: &Bound<'a, PyAny>, context: &str) -> PyResult<Self> {
        ob.extract().map_err(Into::into).map_err(|e: pyo3::PyErr| {
            pyo3::exceptions::PyValueError::new_err(
                format!("Failed to extract {} from Python: {} (object: {})", 
                    context, e, dump(&ob, None).unwrap_or_else(|_| "unknown".to_string()))
            )
        })
    }
}

/// Utility function for consistent logging during Python object extraction.
pub fn log_extraction(ob: &Bound<PyAny>, context: &str) {
    if tracing::enabled!(tracing::Level::DEBUG) {
        match dump(&ob, None) {
            Ok(dump_str) => tracing::debug!("Extracting {}: {}", context, dump_str),
            Err(_) => tracing::debug!("Extracting {} (dump failed)", context),
        }
    }
}

/// Helper function to create standardized error messages for failed extractions.
pub fn extraction_error(context: &str, details: &str) -> pyo3::PyErr {
    pyo3::exceptions::PyValueError::new_err(
        format!("Failed to extract {}: {}", context, details)
    )
}

/// Build an identifier for generated code from a Python name, escaping Rust
/// keywords so identifiers like `type`, `match`, or `move` don't produce
/// invalid Rust. Keywords that cannot be raw identifiers (`self`, `Self`,
/// `super`, `crate`) get a trailing underscore instead — except `self`,
/// which is left as-is so method receivers keep working.
pub fn safe_ident(name: &str) -> proc_macro2::Ident {
    use quote::format_ident;
    match name {
        "self" => format_ident!("{}", name),
        "Self" | "super" | "crate" => format_ident!("{}_", name),
        // Strict and reserved keywords, all legal as raw identifiers.
        "as" | "break" | "const" | "continue" | "dyn" | "else" | "enum" | "extern" | "false"
        | "fn" | "for" | "if" | "impl" | "in" | "let" | "loop" | "match" | "mod" | "move"
        | "mut" | "pub" | "ref" | "return" | "static" | "struct" | "trait" | "true" | "type"
        | "unsafe" | "use" | "where" | "while" | "async" | "await" | "abstract" | "become"
        | "box" | "do" | "final" | "macro" | "override" | "priv" | "try" | "typeof"
        | "unsized" | "virtual" | "yield" | "gen" => format_ident!("r#{}", name),
        _ => format_ident!("{}", name),
    }
}

/// Build a PyErr for a failed AST-node conversion, carrying the node's source
/// position so the error can be reported against the user's Python code
/// instead of panicking with an internal dump.
pub fn extraction_failure(
    kind: &str,
    ob: &Bound<PyAny>,
    cause: impl std::fmt::Display,
) -> pyo3::PyErr {
    let position = match (ob.lineno(), ob.col_offset()) {
        (Some(line), Some(col)) => format!(" at line {}, column {}", line, col + 1),
        (Some(line), None) => format!(" at line {}", line),
        _ => String::new(),
    };
    pyo3::exceptions::PyValueError::new_err(format!(
        "cannot convert {}{}: {}",
        kind, position, cause
    ))
}

#[cfg(test)]
mod tests {
    use super::*;
    use pyo3::Python;

    #[test]
    fn test_extract_optional() {
        Python::attach(|py| {
            use std::ffi::CString;
            // Create a simple Python integer which has some attributes
            let code = CString::new("42").unwrap();
            let obj = py.eval(&code, None, None).unwrap();
            
            // Try to extract something that probably won't exist 
            let missing: Option<String> = extract_optional(&obj, "missing_attr");
            assert_eq!(missing, None);
            
            // This test mainly checks that extract_optional doesn't panic
            // and properly returns None for missing attributes
        });
    }

    #[test]
    fn test_log_extraction() {
        Python::attach(|py| {
            use std::ffi::CString;
            let code = CString::new("42").unwrap();
            let obj = py.eval(&code, None, None).unwrap();
            
            // Should not panic
            log_extraction(&obj, "test object");
        });
    }

    #[test]
    fn test_extraction_error() {
        let error = extraction_error("test context", "test details");
        let error_string = format!("{}", error);
        assert!(error_string.contains("test context"));
        assert!(error_string.contains("test details"));
    }
}

/// Enhanced error handling utilities for parsing Python AST objects

/// Get an attribute from a Python object with better error messaging
pub fn get_attr_with_context<'a>(
    ob: &Bound<'a, PyAny>,
    attr_name: &str,
    context: &str,
) -> PyResult<Bound<'a, PyAny>> {
    ob.getattr(attr_name).map_err(|e| {
        let type_name = ob.get_type().name()
            .map(|s| s.to_string())
            .unwrap_or_else(|_| "<unknown>".to_string());
        let enhanced_msg = format!(
            "Failed to get attribute '{}' from {} ({}): {}",
            attr_name,
            context,
            type_name,
            e
        );
        pyo3::exceptions::PyAttributeError::new_err(enhanced_msg)
    })
}

/// Extract a value from PyAny with better error messaging
pub fn extract_with_context<'py, T>(
    value: &Bound<'py, PyAny>,
    context: &str,
    attr_name: &str,
) -> PyResult<T>
where
    T: pyo3::conversion::FromPyObjectOwned<'py>,
{
    value.extract().map_err(Into::into).map_err(|e: pyo3::PyErr| {
        let type_name = value.get_type().name()
            .map(|s| s.to_string())
            .unwrap_or_else(|_| "<unknown>".to_string());
        let enhanced_msg = format!(
            "Failed to extract {} for attribute '{}': {}. Expected type: {}, got: {}",
            context,
            attr_name,
            e,
            std::any::type_name::<T>(),
            type_name
        );
        pyo3::exceptions::PyTypeError::new_err(enhanced_msg)
    })
}

/// Extract a required attribute with enhanced error messaging  
pub fn extract_required_attr<'py, T>(
    ob: &Bound<'py, PyAny>,
    attr_name: &str,
    context: &str,
) -> PyResult<T>
where
    T: pyo3::conversion::FromPyObjectOwned<'py>,
{
    let attr = get_attr_with_context(&ob, attr_name, context)?;
    extract_with_context(&attr, context, attr_name)
}