python_ast/ast/tree/
set.rs

1use proc_macro2::TokenStream;
2use pyo3::{Borrowed, FromPyObject, PyAny, PyResult};
3use quote::quote;
4use serde::{Deserialize, Serialize};
5
6use crate::{
7    CodeGen, CodeGenContext, ExprType, PythonOptions, SymbolTableScopes,
8    Node, impl_node_with_positions, extract_list
9};
10
11#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
12pub struct Set {
13    pub elts: Vec<ExprType>,
14    pub lineno: Option<usize>,
15    pub col_offset: Option<usize>,
16    pub end_lineno: Option<usize>,
17    pub end_col_offset: Option<usize>,
18}
19
20impl<'a, 'py> FromPyObject<'a, 'py> for Set {
21    type Error = pyo3::PyErr;
22    fn extract(ob: Borrowed<'a, 'py, PyAny>) -> PyResult<Self> {
23        let elts: Vec<ExprType> = extract_list(&ob, "elts", "set elements")?;
24        
25        Ok(Set {
26            elts,
27            lineno: ob.lineno(),
28            col_offset: ob.col_offset(),
29            end_lineno: ob.end_lineno(),
30            end_col_offset: ob.end_col_offset(),
31        })
32    }
33}
34
35impl_node_with_positions!(Set { lineno, col_offset, end_lineno, end_col_offset });
36
37impl CodeGen for Set {
38    type Context = CodeGenContext;
39    type Options = PythonOptions;
40    type SymbolTable = SymbolTableScopes;
41
42    fn to_rust(
43        self,
44        ctx: Self::Context,
45        options: Self::Options,
46        symbols: Self::SymbolTable,
47    ) -> Result<TokenStream, Box<dyn std::error::Error>> {
48        let elements: Result<Vec<_>, _> = self.elts
49            .into_iter()
50            .map(|elt| elt.to_rust(ctx.clone(), options.clone(), symbols.clone()))
51            .collect();
52        
53        let elements = elements?;
54        
55        Ok(quote! {
56            std::collections::HashSet::from([#(#elements),*])
57        })
58    }
59}
60
61#[cfg(test)]
62mod tests {
63    use super::*;
64    use crate::create_parse_test;
65
66    create_parse_test!(test_empty_set, "set()", "set_test.py");
67    create_parse_test!(test_simple_set, "{1, 2, 3}", "set_test.py");
68    create_parse_test!(test_set_with_variables, "{x, y, z}", "set_test.py");
69}
python_ast/ast/tree/set.rs

python_ast/ast/tree/
set.rs
