python_ast/
traits.rs

1use proc_macro2::TokenStream;
2use quote::quote;
3use pyo3::types::{PyAnyMethods, PyTypeMethods};
4use crate::{CodeGen, CodeGenContext, PythonOptions, SymbolTableScopes, ExprType};
5
6/// Common trait for Python operators that can be converted to Rust tokens.
7pub trait PythonOperator: Clone + std::fmt::Debug {
8    /// Convert the operator to its Rust equivalent TokenStream.
9    fn to_rust_op(&self) -> Result<TokenStream, Box<dyn std::error::Error>>;
10    
11    /// Get the operator precedence for proper parenthesization.
12    fn precedence(&self) -> u8 {
13        0 // Default precedence
14    }
15    
16    /// Check if this operator is unknown/unimplemented.
17    fn is_unknown(&self) -> bool;
18}
19
20/// Common trait for binary operations (binary ops, bool ops, comparisons).
21pub trait BinaryOperation: Clone + std::fmt::Debug {
22    type OperatorType: PythonOperator;
23    
24    /// Get the operator type.
25    fn operator(&self) -> &Self::OperatorType;
26    
27    /// Get the left operand.
28    fn left(&self) -> &ExprType;
29    
30    /// Get the right operand.
31    fn right(&self) -> &ExprType;
32    
33    /// Generate Rust code for this binary operation.
34    fn generate_rust_code(
35        &self,
36        ctx: CodeGenContext,
37        options: PythonOptions,
38        symbols: SymbolTableScopes,
39    ) -> Result<TokenStream, Box<dyn std::error::Error>> {
40        let left = self.left()
41            .clone()
42            .to_rust(ctx.clone(), options.clone(), symbols.clone())?;
43        let right = self.right()
44            .clone()
45            .to_rust(ctx, options, symbols)?;
46        let op = self.operator().to_rust_op()?;
47        
48        Ok(quote!((#left) #op (#right)))
49    }
50}
51
52/// Trait for extracting Python attributes with consistent error handling.
53pub trait PyAttributeExtractor {
54    /// Extract an attribute with context-aware error messages.
55    fn extract_attr_with_context(&self, attr: &str, context: &str) -> pyo3::PyResult<pyo3::Bound<'_, pyo3::PyAny>>;
56    
57    /// Extract a type name with error handling.
58    fn extract_type_name(&self, context: &str) -> pyo3::PyResult<String>;
59}
60
61impl<'py> PyAttributeExtractor for pyo3::Bound<'py, pyo3::PyAny> {
62    fn extract_attr_with_context(&self, attr: &str, context: &str) -> pyo3::PyResult<pyo3::Bound<'_, pyo3::PyAny>> {
63        use crate::Node;
64        self.getattr(attr).map_err(|_| {
65            pyo3::exceptions::PyAttributeError::new_err(
66                self.error_message("<unknown>", format!("error getting {}", context))
67            )
68        })
69    }
70    
71    fn extract_type_name(&self, context: &str) -> pyo3::PyResult<String> {
72        use crate::Node;
73        
74        let type_name = self.get_type().name().map_err(|_| {
75            pyo3::exceptions::PyTypeError::new_err(
76                self.error_message("<unknown>", format!("extracting type name for {}", context))
77            )
78        })?;
79        
80        type_name.extract()
81    }
82}
83
84/// Trait for consistent handling of position information in AST nodes.
85pub trait PositionInfo {
86    /// Get all position fields as a tuple (lineno, col_offset, end_lineno, end_col_offset).
87    fn position_info(&self) -> (Option<usize>, Option<usize>, Option<usize>, Option<usize>);
88    
89    /// Check if this node has position information.
90    fn has_position(&self) -> bool {
91        let (lineno, col_offset, _, _) = self.position_info();
92        lineno.is_some() && col_offset.is_some()
93    }
94}
95
96/// Trait for AST nodes that can provide debugging information.
97pub trait DebugInfo {
98    /// Get a human-readable description of this node.
99    fn debug_description(&self) -> String;
100    
101    /// Get the node type name.
102    fn node_type(&self) -> &'static str;
103}
104
105/// Trait for operations that can be chained (like comparison operations).
106pub trait ChainableOperation {
107    type Operand;
108    type Operator;
109    
110    /// Get all operands in the chain.
111    fn operands(&self) -> Vec<&Self::Operand>;
112    
113    /// Get all operators in the chain.
114    fn operators(&self) -> Vec<&Self::Operator>;
115    
116    /// Generate chained Rust code.
117    fn generate_chained_rust(
118        &self,
119        ctx: CodeGenContext,
120        options: PythonOptions,
121        symbols: SymbolTableScopes,
122    ) -> Result<TokenStream, Box<dyn std::error::Error>>;
123}
124
125/// Helper trait for converting Python string literals to enum variants.
126pub trait FromPythonString: Sized {
127    /// Convert a Python operator string to the enum variant.
128    fn from_python_string(s: &str) -> Option<Self>;
129    
130    /// Get the unknown/default variant.
131    fn unknown() -> Self;
132    
133    /// Parse from string with fallback to unknown.
134    fn parse_or_unknown(s: &str) -> Self {
135        Self::from_python_string(s).unwrap_or_else(Self::unknown)
136    }
137}
138
139/// Trait for generating consistent error messages across the codebase.
140pub trait ErrorContext {
141    /// Generate a standardized error message with context.
142    fn with_context(&self, operation: &str) -> String;
143}
144
145impl<T: std::fmt::Debug> ErrorContext for T {
146    fn with_context(&self, operation: &str) -> String {
147        format!("Error during {}: {:?}", operation, self)
148    }
149}