rustpython_vm/types/
structseq.rs

1use crate::{
2    builtins::{PyTuple, PyTupleRef, PyType},
3    class::{PyClassImpl, StaticType},
4    vm::Context,
5    AsObject, Py, PyObjectRef, PyPayload, PyRef, PyResult, VirtualMachine,
6};
7
8#[pyclass]
9pub trait PyStructSequence: StaticType + PyClassImpl + Sized + 'static {
10    const FIELD_NAMES: &'static [&'static str];
11
12    fn into_tuple(self, vm: &VirtualMachine) -> PyTuple;
13
14    fn into_struct_sequence(self, vm: &VirtualMachine) -> PyTupleRef {
15        self.into_tuple(vm)
16            .into_ref_with_type(vm, Self::static_type().to_owned())
17            .unwrap()
18    }
19
20    fn try_elements_from<const FIELD_LEN: usize>(
21        obj: PyObjectRef,
22        vm: &VirtualMachine,
23    ) -> PyResult<[PyObjectRef; FIELD_LEN]> {
24        let typ = Self::static_type();
25        // if !obj.fast_isinstance(typ) {
26        //     return Err(vm.new_type_error(format!(
27        //         "{} is not a subclass of {}",
28        //         obj.class().name(),
29        //         typ.name(),
30        //     )));
31        // }
32        let seq: Vec<PyObjectRef> = obj.try_into_value(vm)?;
33        let seq: [PyObjectRef; FIELD_LEN] = seq.try_into().map_err(|_| {
34            vm.new_type_error(format!(
35                "{} takes a sequence of length {}",
36                typ.name(),
37                FIELD_LEN
38            ))
39        })?;
40        Ok(seq)
41    }
42
43    #[pymethod(magic)]
44    fn repr(zelf: PyRef<PyTuple>, vm: &VirtualMachine) -> PyResult<String> {
45        let format_field = |(value, name): (&PyObjectRef, _)| {
46            let s = value.repr(vm)?;
47            Ok(format!("{name}={s}"))
48        };
49        let (body, suffix) = if let Some(_guard) =
50            rustpython_vm::recursion::ReprGuard::enter(vm, zelf.as_object())
51        {
52            if Self::FIELD_NAMES.len() == 1 {
53                let value = zelf.first().unwrap();
54                let formatted = format_field((value, Self::FIELD_NAMES[0]))?;
55                (formatted, ",")
56            } else {
57                let fields: PyResult<Vec<_>> = zelf
58                    .iter()
59                    .zip(Self::FIELD_NAMES.iter().copied())
60                    .map(format_field)
61                    .collect();
62                (fields?.join(", "), "")
63            }
64        } else {
65            (String::new(), "...")
66        };
67        Ok(format!("{}({}{})", Self::TP_NAME, body, suffix))
68    }
69
70    #[pymethod(magic)]
71    fn reduce(zelf: PyRef<PyTuple>, vm: &VirtualMachine) -> PyTupleRef {
72        vm.new_tuple((zelf.class().to_owned(), (vm.ctx.new_tuple(zelf.to_vec()),)))
73    }
74
75    #[extend_class]
76    fn extend_pyclass(ctx: &Context, class: &'static Py<PyType>) {
77        for (i, &name) in Self::FIELD_NAMES.iter().enumerate() {
78            // cast i to a u8 so there's less to store in the getter closure.
79            // Hopefully there's not struct sequences with >=256 elements :P
80            let i = i as u8;
81            class.set_attr(
82                ctx.intern_str(name),
83                ctx.new_readonly_getset(name, class, move |zelf: &PyTuple| {
84                    zelf.fast_getitem(i.into())
85                })
86                .into(),
87            );
88        }
89
90        class.set_attr(
91            identifier!(ctx, __match_args__),
92            ctx.new_tuple(
93                Self::FIELD_NAMES
94                    .iter()
95                    .map(|&name| ctx.new_str(name).into())
96                    .collect::<Vec<_>>(),
97            )
98            .into(),
99        );
100    }
101}