python_ast/ast/tree/
constant.rs

1use std::fmt::*;
2
3use litrs::Literal;
4use tracing::debug;
5use proc_macro2::*;
6use pyo3::{Borrowed, Bound, FromPyObject, PyAny, PyResult, prelude::PyAnyMethods};
7use quote::quote;
8
9use crate::{CodeGen, CodeGenContext, PythonOptions, SymbolTableScopes};
10
11use serde::{Deserialize, Deserializer, Serialize, Serializer};
12
13pub trait PyConstantTrait: Clone + Debug + PartialEq {
14    type RustType;
15}
16
17#[derive(Clone, Debug, PartialEq)]
18#[repr(transparent)]
19pub struct Constant(pub Option<Literal<String>>);
20
21impl Serialize for Constant {
22    fn serialize<S>(&self, serializer: S) -> std::result::Result<S::Ok, S::Error>
23    where
24        S: Serializer,
25    {
26        serializer.serialize_str(self.to_string().as_str())
27    }
28}
29
30impl<'de> Deserialize<'de> for Constant {
31    fn deserialize<D>(deserializer: D) -> std::result::Result<Self, D::Error>
32    where
33        D: Deserializer<'de>,
34    {
35        let s = String::deserialize(deserializer)?;
36        let l = Literal::parse(s).expect("[3] Parsing the literal");
37        Ok(Self(Some(l)))
38    }
39}
40
41impl std::string::ToString for Constant {
42    fn to_string(&self) -> String {
43        match self.0.clone() {
44            Some(c) => c.to_string(),
45            None => "None".to_string(),
46        }
47    }
48}
49
50pub fn try_string(value: &Bound<PyAny>) -> PyResult<Option<Literal<String>>> {
51    let v: String = value.extract()?;
52    // Debug-format the string so quotes, backslashes, and control characters
53    // come out as valid Rust escape sequences.
54    let l = Literal::parse(format!("{:?}", v)).expect("[4] Parsing the literal");
55
56    Ok(Some(l))
57}
58
59pub fn try_bytes(value: &Bound<PyAny>) -> PyResult<Option<Literal<String>>> {
60    let v: &[u8] = value.extract()?;
61    // Rust byte-string literals only allow ASCII plus escapes, so escape every
62    // byte that needs it (non-ASCII bytes become \xNN).
63    let escaped: String = v
64        .iter()
65        .flat_map(|b| std::ascii::escape_default(*b))
66        .map(char::from)
67        .collect();
68    let l = Literal::parse(format!("b\"{}\"", escaped)).expect("[4] Parsing the literal");
69
70    Ok(Some(l))
71}
72
73pub fn try_int(value: &Bound<PyAny>) -> PyResult<Option<Literal<String>>> {
74    let v: isize = value.extract()?;
75    let l = Literal::parse(format!("{}", v)).expect("[4] Parsing the literal");
76
77    Ok(Some(l))
78}
79
80pub fn try_float(value: &Bound<PyAny>) -> PyResult<Option<Literal<String>>> {
81    let v: f64 = value.extract()?;
82    // Rust's Display for integral floats drops the ".0" ("2.0" becomes
83    // "2"), which would re-parse as an INTEGER literal and silently change
84    // the generated type (and semantics — Python's 2.0 / 4 is 0.5).
85    let mut s = format!("{}", v);
86    if v.is_finite() && !s.contains('.') && !s.contains('e') && !s.contains('E') {
87        s.push_str(".0");
88    }
89    let l = Literal::parse(s).expect("[4] Parsing the literal");
90
91    Ok(Some(l))
92}
93
94pub fn try_bool(value: &Bound<PyAny>) -> PyResult<Option<Literal<String>>> {
95    let v: bool = value.extract()?;
96    let l = Literal::parse(format!("{}", v)).expect("[4] Parsing the literal");
97
98    Ok(Some(l))
99}
100
101// This will mostly be invoked when the input is None.
102pub fn try_option(value: &Bound<PyAny>) -> PyResult<Option<Literal<String>>> {
103    let v: Option<Bound<PyAny>> = value.extract()?;
104    // If we got None as a constant, return None; anything else is a constant
105    // kind we don't know how to render as a Rust literal.
106    match v {
107        None => Ok(None),
108        Some(other) => Err(crate::extraction_failure(
109            "constant",
110            &other,
111            "this constant kind cannot be represented as a Rust literal",
112        )),
113    }
114}
115
116// This is the fun bit of code that is responsible from converting from Python constants to Rust ones.
117impl<'a, 'py> FromPyObject<'a, 'py> for Constant {
118    type Error = pyo3::PyErr;
119    fn extract(ob: Borrowed<'a, 'py, PyAny>) -> PyResult<Self> {
120        // Extracts the values as a PyAny.
121        let value = ob
122            .getattr("value")
123            .map_err(|e| crate::extraction_failure("constant value", &ob, e))?;
124        debug!("[2] constant value: {}", value);
125
126        let l = if let Ok(l) = try_string(&value) {
127            l
128        } else if let Ok(l) = try_bytes(&value) {
129            l
130        // We have to evaluaet bool before int because if a bool is evaluated as it, it will be cooerced to an in.
131        } else if let Ok(l) = try_bool(&value) {
132            l
133        // Ints must be tried before floats: extracting f64 from a Python int
134        // succeeds, and would silently lose precision above 2^53.
135        } else if let Ok(l) = try_int(&value) {
136            l
137        } else if let Ok(l) = try_float(&value) {
138            l
139        } else if let Ok(l) = try_option(&value) {
140            l
141        } else {
142            return Err(crate::extraction_failure(
143                "constant",
144                &value,
145                format!("unsupported constant value `{}`", value),
146            ));
147        };
148
149        Ok(Self(l))
150    }
151}
152
153impl CodeGen for Constant {
154    type Context = CodeGenContext;
155    type Options = PythonOptions;
156    type SymbolTable = SymbolTableScopes;
157
158    fn to_rust(
159        self,
160        _ctx: Self::Context,
161        _options: Self::Options,
162        _symbols: Self::SymbolTable,
163    ) -> std::result::Result<TokenStream, Box<dyn std::error::Error>> {
164        match self.0 {
165            Some(c) => {
166                let v: TokenStream = c
167                    .to_string()
168                    .parse()
169                    .map_err(|e| format!("cannot render constant `{}` as Rust tokens: {}", c, e))?;
170                Ok(quote!(#v))
171            }
172            None => Ok(quote!(None)),
173        }
174    }
175}
176
177#[cfg(test)]
178mod tests {
179    use test_log::test;
180    //use super::*;
181    use crate::{symbols::SymbolTableScopes, CodeGen};
182    use tracing::debug;
183
184    #[test]
185    fn parse_string() {
186        let s = crate::parse("'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        debug!("ast: {}", ast.to_string());
195
196        assert_eq!("use stdpython :: * ; \"I ate a bug\"", ast.to_string());
197    }
198
199    #[test]
200    fn parse_bytes() {
201        let s = crate::parse("b'I ate a bug'", "test.py").unwrap();
202        let ast = s
203            .to_rust(
204                crate::CodeGenContext::Module("test".to_string()),
205                crate::PythonOptions::default(),
206                SymbolTableScopes::new(),
207            )
208            .unwrap();
209
210        assert_eq!("use stdpython :: * ; b\"I ate a bug\"", ast.to_string());
211    }
212
213    #[test]
214    fn parse_number_int() {
215        let s = crate::parse("871234234", "test.py").unwrap();
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
224        assert_eq!("use stdpython :: * ; 871234234", ast.to_string());
225    }
226
227    #[test]
228    fn parse_number_neg_int() {
229        let s = crate::parse("-871234234", "test.py").unwrap();
230        let ast = s
231            .to_rust(
232                crate::CodeGenContext::Module("test".to_string()),
233                crate::PythonOptions::default(),
234                SymbolTableScopes::new(),
235            )
236            .unwrap();
237
238        assert_eq!("use stdpython :: * ; - 871234234", ast.to_string());
239    }
240
241    #[test]
242    fn parse_number_float() {
243        let s = crate::parse("87123.4234", "test.py").unwrap();
244        let ast = s
245            .to_rust(
246                crate::CodeGenContext::Module("test".to_string()),
247                crate::PythonOptions::default(),
248                SymbolTableScopes::new(),
249            )
250            .unwrap();
251
252        assert_eq!("use stdpython :: * ; 87123.4234", ast.to_string());
253    }
254
255    #[test]
256    fn parse_bool() {
257        let s = crate::parse("True", "test.py").unwrap();
258        let ast = s
259            .to_rust(
260                crate::CodeGenContext::Module("test".to_string()),
261                crate::PythonOptions::default(),
262                SymbolTableScopes::new(),
263            )
264            .unwrap();
265
266        assert_eq!("use stdpython :: * ; true", ast.to_string());
267    }
268
269    #[test]
270    fn parse_none() {
271        let s = crate::parse("None", "test.py").unwrap();
272        let ast = s
273            .to_rust(
274                crate::CodeGenContext::Module("test".to_string()),
275                crate::PythonOptions::default(),
276                SymbolTableScopes::new(),
277            )
278            .unwrap();
279
280        assert_eq!("use stdpython :: * ; None", ast.to_string());
281    }
282}
python_ast/ast/tree/constant.rs

python_ast/ast/tree/
constant.rs
