python_ast/ast/tree/
set.rs

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

python_ast/ast/tree/
set.rs
