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

python_ast/ast/tree/
constant.rs
