python_ast/ast/tree/
unary_op.rs

1use 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        //log::info!("{}", result);
113
114        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}