python_ast/ast/tree/
constant.rs

1use std::fmt::*;
2
3use encoding::{all::ISO_8859_6, DecoderTrap, Encoding};
4use litrs::Literal;
5use log::debug;
6use proc_macro2::*;
7use pyo3::{FromPyObject, PyAny, PyResult};
8use quote::quote;
9
10use crate::{CodeGen, CodeGenContext, Node, PythonOptions, SymbolTableScopes};
11
12use serde::{Deserialize, Deserializer, Serialize, Serializer};
13
14#[derive(Clone, Debug, PartialEq)]
15#[repr(transparent)]
16pub struct Constant(pub Option<Literal<String>>);
17
18impl Serialize for Constant {
19    fn serialize<S>(&self, serializer: S) -> std::result::Result<S::Ok, S::Error>
20    where
21        S: Serializer,
22    {
23        serializer.serialize_str(self.to_string().as_str())
24    }
25}
26
27impl<'de> Deserialize<'de> for Constant {
28    fn deserialize<D>(deserializer: D) -> std::result::Result<Self, D::Error>
29    where
30        D: Deserializer<'de>,
31    {
32        let s = String::deserialize(deserializer)?;
33        let l = Literal::parse(s).expect("[3] Parsing the literal");
34        Ok(Self(Some(l)))
35    }
36}
37
38impl std::string::ToString for Constant {
39    fn to_string(&self) -> String {
40        match self.0.clone() {
41            Some(c) => c.to_string(),
42            None => "None".to_string(),
43        }
44    }
45}
46
47pub fn try_string(value: &PyAny) -> PyResult<Option<Literal<String>>> {
48    let v: String = value.extract()?;
49    let l = Literal::parse(format!("\"{}\"", v)).expect("[4] Parsing the literal");
50
51    Ok(Some(l))
52}
53
54pub fn try_bytes(value: &PyAny) -> PyResult<Option<Literal<String>>> {
55    let v: &[u8] = value.extract()?;
56    let l = Literal::parse(format!(
57        "b\"{}\"",
58        ISO_8859_6
59            .decode(v, DecoderTrap::Replace)
60            .expect("decoding byte string")
61    ))
62    .expect("[4] Parsing the literal");
63
64    Ok(Some(l))
65}
66
67pub fn try_int(value: &PyAny) -> PyResult<Option<Literal<String>>> {
68    let v: isize = value.extract()?;
69    let l = Literal::parse(format!("{}", v)).expect("[4] Parsing the literal");
70
71    Ok(Some(l))
72}
73
74pub fn try_float(value: &PyAny) -> PyResult<Option<Literal<String>>> {
75    let v: f64 = value.extract()?;
76    let l = Literal::parse(format!("{}", v)).expect("[4] Parsing the literal");
77
78    Ok(Some(l))
79}
80
81pub fn try_bool(value: &PyAny) -> PyResult<Option<Literal<String>>> {
82    let v: bool = value.extract()?;
83    let l = Literal::parse(format!("{}", v)).expect("[4] Parsing the literal");
84
85    Ok(Some(l))
86}
87
88// This will mostly be invoked when the input is None.
89pub fn try_option(value: &PyAny) -> PyResult<Option<Literal<String>>> {
90    let v: Option<&PyAny> = value.extract()?;
91    debug!("extracted value {:?}", v);
92    // If we got None as a constant, return None
93    match v {
94        None => Ok(None),
95        // See if we can parse whatever we got that wasn't None.
96        Some(_c) => {
97            let l = Literal::parse(format!("{:?}", v)).expect("[5] Parsing the literal");
98            Ok(Some(l))
99        }
100    }
101}
102
103// This is the fun bit of code that is responsible from converting from Python constants to Rust ones.
104impl<'a> FromPyObject<'a> for Constant {
105    fn extract(ob: &'a PyAny) -> PyResult<Self> {
106        // Extracts the values as a PyAny.
107        let value = ob.getattr("value").expect(
108            ob.error_message("<unknown>", "error getting constant value")
109                .as_str(),
110        );
111        debug!("[2] constant value: {}", value);
112
113        let l = if let Ok(l) = try_string(value) {
114            l
115        } else if let Ok(l) = try_bytes(value) {
116            l
117        // We have to evaluaet bool before int because if a bool is evaluated as it, it will be cooerced to an in.
118        } else if let Ok(l) = try_bool(value) {
119            l
120        } else if let Ok(l) = try_float(value) {
121            l
122        } else if let Ok(l) = try_int(value) {
123            l
124        } else if let Ok(l) = try_option(value) {
125            l
126        } else {
127            panic!("Failed to parse literal values {}", value);
128        };
129
130        Ok(Self(l))
131    }
132}
133
134impl CodeGen for Constant {
135    type Context = CodeGenContext;
136    type Options = PythonOptions;
137    type SymbolTable = SymbolTableScopes;
138
139    fn to_rust(
140        self,
141        _ctx: Self::Context,
142        _options: Self::Options,
143        _symbols: Self::SymbolTable,
144    ) -> std::result::Result<TokenStream, Box<dyn std::error::Error>> {
145        match self.0 {
146            Some(c) => {
147                let v: TokenStream = c
148                    .to_string()
149                    .parse()
150                    .expect(format!("parsing Constant {}", c).as_str());
151                Ok(quote!(#v))
152            }
153            None => Ok(quote!(None)),
154        }
155    }
156}
157
158#[cfg(test)]
159mod tests {
160    use test_log::test;
161    //use super::*;
162    use crate::{symbols::SymbolTableScopes, CodeGen};
163    use log::debug;
164
165    #[test]
166    fn parse_string() {
167        let s = crate::parse("'I ate a bug'", "test.py").unwrap();
168        let ast = s
169            .to_rust(
170                crate::CodeGenContext::Module("test".to_string()),
171                crate::PythonOptions::default(),
172                SymbolTableScopes::new(),
173            )
174            .unwrap();
175        debug!("ast: {}", ast.to_string());
176
177        assert_eq!("use stdpython :: * ; \"I ate a bug\"", ast.to_string());
178    }
179
180    #[test]
181    fn parse_bytes() {
182        let s = crate::parse("b'I ate a bug'", "test.py").unwrap();
183        println!("parsed value: {:?}", s);
184        let ast = s
185            .to_rust(
186                crate::CodeGenContext::Module("test".to_string()),
187                crate::PythonOptions::default(),
188                SymbolTableScopes::new(),
189            )
190            .unwrap();
191        println!("ast: {:?}", ast);
192
193        assert_eq!("use stdpython :: * ; b\"I ate a bug\"", ast.to_string());
194    }
195
196    #[test]
197    fn parse_number_int() {
198        let s = crate::parse("871234234", "test.py").unwrap();
199        println!("parsed value: {:?}", s);
200        let ast = s
201            .to_rust(
202                crate::CodeGenContext::Module("test".to_string()),
203                crate::PythonOptions::default(),
204                SymbolTableScopes::new(),
205            )
206            .unwrap();
207        println!("ast: {:?}", ast);
208
209        assert_eq!("use stdpython :: * ; 871234234", ast.to_string());
210    }
211
212    #[test]
213    fn parse_number_neg_int() {
214        let s = crate::parse("-871234234", "test.py").unwrap();
215        println!("parsed value: {:?}", s);
216        let ast = s
217            .to_rust(
218                crate::CodeGenContext::Module("test".to_string()),
219                crate::PythonOptions::default(),
220                SymbolTableScopes::new(),
221            )
222            .unwrap();
223        println!("ast: {:?}", ast);
224
225        assert_eq!("use stdpython :: * ; - 871234234", ast.to_string());
226    }
227
228    #[test]
229    fn parse_number_float() {
230        let s = crate::parse("87123.4234", "test.py").unwrap();
231        println!("parsed value: {:?}", s);
232        let ast = s
233            .to_rust(
234                crate::CodeGenContext::Module("test".to_string()),
235                crate::PythonOptions::default(),
236                SymbolTableScopes::new(),
237            )
238            .unwrap();
239        println!("ast: {:?}", ast);
240
241        assert_eq!("use stdpython :: * ; 87123.4234", ast.to_string());
242    }
243
244    #[test]
245    fn parse_bool() {
246        let s = crate::parse("True", "test.py").unwrap();
247        println!("parsed value: {:?}", s);
248        let ast = s
249            .to_rust(
250                crate::CodeGenContext::Module("test".to_string()),
251                crate::PythonOptions::default(),
252                SymbolTableScopes::new(),
253            )
254            .unwrap();
255        println!("ast: {:?}", ast);
256
257        assert_eq!("use stdpython :: * ; true", ast.to_string());
258    }
259
260    #[test]
261    fn parse_none() {
262        let s = crate::parse("None", "test.py").unwrap();
263        println!("parsed value: {:?}", s);
264        let ast = s
265            .to_rust(
266                crate::CodeGenContext::Module("test".to_string()),
267                crate::PythonOptions::default(),
268                SymbolTableScopes::new(),
269            )
270            .unwrap();
271        println!("ast: {:?}", ast);
272
273        assert_eq!("use stdpython :: * ; None", ast.to_string());
274    }
275}