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        let ast = s
188            .to_rust(
189                crate::CodeGenContext::Module("test".to_string()),
190                crate::PythonOptions::default(),
191                SymbolTableScopes::new(),
192            )
193            .unwrap();
194
195        assert_eq!("use stdpython :: * ; b\"I ate a bug\"", ast.to_string());
196    }
197
198    #[test]
199    fn parse_number_int() {
200        let s = crate::parse("871234234", "test.py").unwrap();
201        let ast = s
202            .to_rust(
203                crate::CodeGenContext::Module("test".to_string()),
204                crate::PythonOptions::default(),
205                SymbolTableScopes::new(),
206            )
207            .unwrap();
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        let ast = s
216            .to_rust(
217                crate::CodeGenContext::Module("test".to_string()),
218                crate::PythonOptions::default(),
219                SymbolTableScopes::new(),
220            )
221            .unwrap();
222
223        assert_eq!("use stdpython :: * ; - 871234234", ast.to_string());
224    }
225
226    #[test]
227    fn parse_number_float() {
228        let s = crate::parse("87123.4234", "test.py").unwrap();
229        let ast = s
230            .to_rust(
231                crate::CodeGenContext::Module("test".to_string()),
232                crate::PythonOptions::default(),
233                SymbolTableScopes::new(),
234            )
235            .unwrap();
236
237        assert_eq!("use stdpython :: * ; 87123.4234", ast.to_string());
238    }
239
240    #[test]
241    fn parse_bool() {
242        let s = crate::parse("True", "test.py").unwrap();
243        let ast = s
244            .to_rust(
245                crate::CodeGenContext::Module("test".to_string()),
246                crate::PythonOptions::default(),
247                SymbolTableScopes::new(),
248            )
249            .unwrap();
250
251        assert_eq!("use stdpython :: * ; true", ast.to_string());
252    }
253
254    #[test]
255    fn parse_none() {
256        let s = crate::parse("None", "test.py").unwrap();
257        let ast = s
258            .to_rust(
259                crate::CodeGenContext::Module("test".to_string()),
260                crate::PythonOptions::default(),
261                SymbolTableScopes::new(),
262            )
263            .unwrap();
264
265        assert_eq!("use stdpython :: * ; None", ast.to_string());
266    }
267}
python_ast/ast/tree/constant.rs

python_ast/ast/tree/
constant.rs
