python_ast/ast/tree/
unary_op.rs

1use proc_macro2::TokenStream;
2use pyo3::{FromPyObject, PyAny, PyResult};
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(ob: &'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(ob: &'a PyAny) -> PyResult<Self> {
38        log::debug!("ob: {}", dump(ob, None)?);
39        let op = ob.getattr("op").expect(
40            ob.error_message("<unknown>", "error getting unary operator")
41                .as_str(),
42        );
43
44        let op_type = op.get_type().name().expect(
45            ob.error_message(
46                "<unknown>",
47                format!("extracting type name {:?} for unary operator", op),
48            )
49            .as_str(),
50        );
51
52        let operand = ob.getattr("operand").expect(
53            ob.error_message("<unknown>", "error getting unary operand")
54                .as_str(),
55        );
56
57        let op = match op_type.as_ref() {
58            "Invert" => Ops::Invert,
59            "Not" => Ops::Not,
60            "UAdd" => Ops::UAdd,
61            "USub" => Ops::USub,
62            _ => {
63                log::debug!("{:?}", op);
64                Ops::Unknown
65            }
66        };
67
68        log::debug!("operand: {}", dump(operand, None)?);
69        let operand = ExprType::extract(operand).expect("getting unary operator operand");
70
71        return Ok(UnaryOp {
72            op: op,
73            operand: Box::new(operand),
74        });
75    }
76}
77
78impl CodeGen for UnaryOp {
79    type Context = CodeGenContext;
80    type Options = PythonOptions;
81    type SymbolTable = SymbolTableScopes;
82
83    fn to_rust(
84        self,
85        ctx: Self::Context,
86        options: Self::Options,
87        symbols: Self::SymbolTable,
88    ) -> Result<TokenStream, Box<dyn std::error::Error>> {
89        let operand = self.operand.clone().to_rust(ctx, options, symbols)?;
90        match self.op {
91            Ops::Invert | Ops::Not => Ok(quote!(!#operand)),
92            Ops::UAdd => Ok(quote!(+#operand)),
93            Ops::USub => Ok(quote!(-#operand)),
94            _ => Err(Error::UnaryOpNotYetImplemented(self).into())
95        }
96    }
97}
98
99#[cfg(test)]
100mod tests {
101    use super::*;
102
103    #[test]
104    fn test_not() {
105        let options = PythonOptions::default();
106        let result = crate::parse("not True", "test").unwrap();
107        log::info!("Python tree: {:?}", result);
108        //log::info!("{}", result);
109
110        let code = result
111            .to_rust(
112                CodeGenContext::Module("test".to_string()),
113                options,
114                SymbolTableScopes::new(),
115            )
116            .unwrap();
117        log::info!("module: {:?}", code);
118    }
119}