Skip to main content

python_ast/ast/tree/
unary_op.rs

1use proc_macro2::TokenStream;
2use pyo3::{Borrowed, FromPyObject, PyAny, PyResult, prelude::PyAnyMethods, types::PyTypeMethods};
3use quote::quote;
4
5use crate::{extraction_failure,     dump, err_from, CodeGen, CodeGenContext, ExprType, PythonOptions, SymbolTableScopes,
6    UnaryOpNotYetImplemented,
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
21#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
22pub struct UnaryOp {
23    pub op: Ops,
24    pub operand: Box<ExprType>,
25}
26
27impl<'a, 'py> FromPyObject<'a, 'py> for UnaryOp {
28    type Error = pyo3::PyErr;
29    fn extract(ob: Borrowed<'a, 'py, PyAny>) -> PyResult<Self> {
30        let py = ob.py();
31
32        tracing::debug!("ob: {}", dump(&ob, None)?);
33        let op = ob
34            .as_unbound()
35            .getattr(py, "op")
36            .map_err(|e| extraction_failure("unary operator", &ob, e))?;
37
38        let bound_op = op.bind(py);
39        let op_type = bound_op
40            .get_type()
41            .name()
42            .map_err(|e| extraction_failure("unary operator type", &ob, e))?;
43
44        let operand = ob
45            .as_unbound()
46            .getattr(py, "operand")
47            .map_err(|e| extraction_failure("unary operand", &ob, e))?;
48
49        let op = match op_type.extract::<String>()?.as_str() {
50            "Invert" => Ops::Invert,
51            "Not" => Ops::Not,
52            "UAdd" => Ops::UAdd,
53            "USub" => Ops::USub,
54            _ => {
55                tracing::debug!("{:?}", op);
56                Ops::Unknown
57            }
58        };
59
60        tracing::debug!("operand: {}", dump(&operand.bind(py), None)?);
61        let bound_op = operand.bind(py);
62        let operand = ExprType::extract(bound_op.as_borrowed())
63            .map_err(|e| extraction_failure("unary operator operand", &ob, e))?;
64
65        return Ok(UnaryOp {
66            op: op,
67            operand: Box::new(operand),
68        });
69    }
70}
71
72impl CodeGen for UnaryOp {
73    type Context = CodeGenContext;
74    type Options = PythonOptions;
75    type SymbolTable = SymbolTableScopes;
76
77    fn to_rust(
78        self,
79        ctx: Self::Context,
80        options: Self::Options,
81        symbols: Self::SymbolTable,
82    ) -> Result<TokenStream, Box<dyn std::error::Error>> {
83        let operand = self.operand.clone().to_rust(ctx, options, symbols)?;
84        match self.op {
85            // `~x` is Rust's bitwise complement, but `not x` is a
86            // TRUTHINESS test: `not 5` is False, where `!5i64` is -6.
87            Ops::Invert => Ok(quote!(!#operand)),
88            Ops::Not => Ok(quote!(!(#operand).is_truthy())),
89            // Rust has no unary plus; Python's `+x` is the identity for
90            // numbers, so emit the operand alone (parenthesized).
91            Ops::UAdd => Ok(quote!((#operand))),
92            Ops::USub => Ok(quote!(-#operand)),
93            _ => Err(err_from(UnaryOpNotYetImplemented(self)).into())
94        }
95    }
96}
97
98#[cfg(test)]
99mod tests {
100    use super::*;
101
102    #[test]
103    fn test_not() {
104        let options = PythonOptions::default();
105        let result = crate::parse("not True", "test").unwrap();
106        tracing::info!("Python tree: {:?}", result);
107        //tracing::info!("{}", result);
108
109        let code = result
110            .to_rust(
111                CodeGenContext::Module("test".to_string()),
112                options,
113                SymbolTableScopes::new(),
114            )
115            .unwrap();
116        tracing::info!("module: {:?}", code);
117    }
118}