python_ast/ast/tree/
unary_op.rs1use 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 Ops::Invert => Ok(quote!(!#operand)),
88 Ops::Not => Ok(quote!(!(#operand).is_truthy())),
89 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 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}