python_ast/ast/tree/
unary_op.rs1use 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 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}