python_ast/ast/tree/
subscript.rs

1use proc_macro2::TokenStream;
2use pyo3::{Borrowed, FromPyObject, PyAny, PyResult, types::PyAnyMethods};
3use quote::quote;
4use serde::{Deserialize, Serialize};
5
6use crate::{extraction_failure, 
7    CodeGen, CodeGenContext, ExprType, PythonOptions, SymbolTableScopes,
8    Node, impl_node_with_positions, PyAttributeExtractor
9};
10
11/// A subscript's bracket contents: a plain index or a slice.
12#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
13pub enum SubscriptKind {
14    Index(Box<ExprType>),
15    Slice {
16        lower: Option<Box<ExprType>>,
17        upper: Option<Box<ExprType>>,
18        step: Option<Box<ExprType>>,
19    },
20}
21
22#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
23pub struct Subscript {
24    pub value: Box<ExprType>,
25    pub kind: SubscriptKind,
26    pub lineno: Option<usize>,
27    pub col_offset: Option<usize>,
28    pub end_lineno: Option<usize>,
29    pub end_col_offset: Option<usize>,
30}
31
32impl<'a, 'py> FromPyObject<'a, 'py> for Subscript {
33    type Error = pyo3::PyErr;
34    fn extract(ob: Borrowed<'a, 'py, PyAny>) -> PyResult<Self> {
35        use pyo3::types::PyTypeMethods;
36
37        let value = ob.extract_attr_with_context("value", "subscript value")?;
38        let value = value.extract().map_err(|e| extraction_failure("getting subscript value", &ob, e))?;
39
40        let slice_attr = ob.extract_attr_with_context("slice", "subscript slice")?;
41        let slice_type: String = slice_attr
42            .get_type()
43            .name()
44            .and_then(|n| n.extract())
45            .map_err(|e| extraction_failure("subscript slice type", &ob, e))?;
46
47        let kind = if slice_type == "Slice" {
48            let bound = |name: &str| -> PyResult<Option<Box<ExprType>>> {
49                match slice_attr.getattr(name) {
50                    Ok(v) if !v.is_none() => Ok(Some(Box::new(v.extract().map_err(
51                        |e| extraction_failure("slice bound", &ob, e),
52                    )?))),
53                    _ => Ok(None),
54                }
55            };
56            SubscriptKind::Slice {
57                lower: bound("lower")?,
58                upper: bound("upper")?,
59                step: bound("step")?,
60            }
61        } else {
62            let index = slice_attr
63                .extract()
64                .map_err(|e| extraction_failure("getting subscript slice", &ob, e))?;
65            SubscriptKind::Index(Box::new(index))
66        };
67
68        Ok(Subscript {
69            value: Box::new(value),
70            kind,
71            lineno: ob.lineno(),
72            col_offset: ob.col_offset(),
73            end_lineno: ob.end_lineno(),
74            end_col_offset: ob.end_col_offset(),
75        })
76    }
77}
78
79impl_node_with_positions!(Subscript { lineno, col_offset, end_lineno, end_col_offset });
80
81/// Lower the receiver of a subscript STORE as a place. Names and
82/// attributes are places already; a nested subscript (`grid[i][j] = v`)
83/// threads through py_index_mut so the store mutates the real container —
84/// the Load lowering would yield a clone and silently drop the write.
85/// Anything else (e.g. a call result) is rejected loudly.
86pub(crate) fn subscript_receiver_place(
87    expr: &ExprType,
88    ctx: CodeGenContext,
89    options: PythonOptions,
90    symbols: SymbolTableScopes,
91) -> Result<TokenStream, Box<dyn std::error::Error>> {
92    match expr {
93        ExprType::Subscript(sub) => {
94            let inner = subscript_receiver_place(
95                &sub.value,
96                ctx.clone(),
97                options.clone(),
98                symbols.clone(),
99            )?;
100            match &sub.kind {
101                SubscriptKind::Index(index) => {
102                    let index = index
103                        .clone()
104                        .to_rust(ctx, options, symbols)?;
105                    Ok(quote!((#inner).py_index_mut(#index)?))
106                }
107                SubscriptKind::Slice { .. } => Err(
108                    "cannot assign through a slice (`x[a:b][...] = ...`)"
109                        .to_string()
110                        .into(),
111                ),
112            }
113        }
114        ExprType::Name(_) | ExprType::Attribute(_) => {
115            expr.clone().to_rust(ctx, options, symbols)
116        }
117        other => Err(format!(
118            "cannot assign into a subscript of this expression: {:?} (only \
119             variables, attributes, and nested subscripts can be stored \
120             through)",
121            other
122        )
123        .into()),
124    }
125}
126
127impl CodeGen for Subscript {
128    type Context = CodeGenContext;
129    type Options = PythonOptions;
130    type SymbolTable = SymbolTableScopes;
131
132    fn to_rust(
133        self,
134        ctx: Self::Context,
135        options: Self::Options,
136        symbols: Self::SymbolTable,
137    ) -> Result<TokenStream, Box<dyn std::error::Error>> {
138        let value = self.value.to_rust(ctx.clone(), options.clone(), symbols.clone())?;
139        match self.kind {
140            // Python index rules via PyIndex: negatives from the end, a
141            // catchable IndexError/KeyError instead of a Rust panic.
142            SubscriptKind::Index(index) => {
143                let index = index.to_rust(ctx, options, symbols)?;
144                Ok(quote! { (#value).py_index(#index)? })
145            }
146            // Slices clamp and never raise.
147            SubscriptKind::Slice { lower, upper, step } => {
148                let bound = |b: Option<Box<ExprType>>| -> Result<TokenStream, Box<dyn std::error::Error>> {
149                    Ok(match b {
150                        Some(e) => {
151                            let t = e.to_rust(ctx.clone(), options.clone(), symbols.clone())?;
152                            quote!(Some(#t))
153                        }
154                        None => quote!(None),
155                    })
156                };
157                let lower = bound(lower)?;
158                let upper = bound(upper)?;
159                let step = bound(step)?;
160                Ok(quote! { (#value).py_slice(#lower, #upper, #step) })
161            }
162        }
163    }
164}
165
166#[cfg(test)]
167mod tests {
168    use super::*;
169    use crate::create_parse_test;
170
171    create_parse_test!(test_list_subscript, "a[0]", "subscript_test.py");
172    create_parse_test!(test_dict_subscript, "d['key']", "subscript_test.py");
173    create_parse_test!(test_nested_subscript, "matrix[i][j]", "subscript_test.py");
174}
python_ast/ast/tree/subscript.rs

python_ast/ast/tree/
subscript.rs
