python_ast/ast/tree/
unary_op.rs1use proc_macro2::TokenStream;
2use pyo3::{Bound, FromPyObject, PyAny, PyResult, prelude::PyAnyMethods, types::PyTypeMethods};
3use quote::quote;
4
5use crate::{
6 dump, CodeGen, CodeGenContext, Error, ExprType, Node, PythonOptions, SymbolTableScopes,
7};
8
9use serde::{Deserialize, Serialize};
10
11#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
12pub enum Ops {
13 Invert,
14 Not,
15 UAdd,
16 USub,
17
18 Unknown,
19}
20
21impl<'a> FromPyObject<'a> for Ops {
22 fn extract_bound(ob: &Bound<'a, PyAny>) -> PyResult<Self> {
23 let err_msg = format!("Unimplemented unary op {}", dump(ob, None)?);
24 Err(pyo3::exceptions::PyValueError::new_err(
25 ob.error_message("<unknown>", err_msg),
26 ))
27 }
28}
29
30#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
31pub struct UnaryOp {
32 op: Ops,
33 operand: Box<ExprType>,
34}
35
36impl<'a> FromPyObject<'a> for UnaryOp {
37 fn extract_bound(ob: &Bound<'a, PyAny>) -> PyResult<Self> {
38 let py = ob.py();
39
40 log::debug!("ob: {}", dump(ob, None)?);
41 let op = ob.as_unbound().getattr(py, "op").expect(
42 ob.error_message("<unknown>", "error getting unary operator")
43 .as_str(),
44 );
45
46 let bound_op = op.bind(py);
47 let op_type = bound_op.get_type().name().expect(
48 ob.error_message(
49 "<unknown>",
50 format!("extracting type name {:?} for unary operator", op),
51 )
52 .as_str(),
53 );
54
55 let operand = ob.as_unbound().getattr(py, "operand").expect(
56 ob.error_message("<unknown>", "error getting unary operand")
57 .as_str(),
58 );
59
60 let op = match op_type.extract::<String>()?.as_str() {
61 "Invert" => Ops::Invert,
62 "Not" => Ops::Not,
63 "UAdd" => Ops::UAdd,
64 "USub" => Ops::USub,
65 _ => {
66 log::debug!("{:?}", op);
67 Ops::Unknown
68 }
69 };
70
71 log::debug!("operand: {}", dump(&operand.bind(py), None)?);
72 let bound_op = operand.bind(py);
73 let operand = ExprType::extract_bound(bound_op).expect("getting unary operator operand");
74
75 return Ok(UnaryOp {
76 op: op,
77 operand: Box::new(operand),
78 });
79 }
80}
81
82impl CodeGen for UnaryOp {
83 type Context = CodeGenContext;
84 type Options = PythonOptions;
85 type SymbolTable = SymbolTableScopes;
86
87 fn to_rust(
88 self,
89 ctx: Self::Context,
90 options: Self::Options,
91 symbols: Self::SymbolTable,
92 ) -> Result<TokenStream, Box<dyn std::error::Error>> {
93 let operand = self.operand.clone().to_rust(ctx, options, symbols)?;
94 match self.op {
95 Ops::Invert | Ops::Not => Ok(quote!(!#operand)),
96 Ops::UAdd => Ok(quote!(+#operand)),
97 Ops::USub => Ok(quote!(-#operand)),
98 _ => Err(Error::UnaryOpNotYetImplemented(self).into())
99 }
100 }
101}
102
103#[cfg(test)]
104mod tests {
105 use super::*;
106
107 #[test]
108 fn test_not() {
109 let options = PythonOptions::default();
110 let result = crate::parse("not True", "test").unwrap();
111 log::info!("Python tree: {:?}", result);
112 let code = result
115 .to_rust(
116 CodeGenContext::Module("test".to_string()),
117 options,
118 SymbolTableScopes::new(),
119 )
120 .unwrap();
121 log::info!("module: {:?}", code);
122 }
123}