From 0abd8b1d8706e2e3a496ff9cfd0a7c1446f547a3 Mon Sep 17 00:00:00 2001 From: Jeong YunWon Date: Thu, 27 Mar 2025 20:35:08 +0900 Subject: [PATCH] Fix pystructseq --- derive-impl/src/pystructseq.rs | 30 ++++++++++++------------ vm/src/types/structseq.rs | 43 ++++++++++++++++++++-------------- 2 files changed, 40 insertions(+), 33 deletions(-) diff --git a/derive-impl/src/pystructseq.rs b/derive-impl/src/pystructseq.rs index 4a4f9276f..89dee8da1 100644 --- a/derive-impl/src/pystructseq.rs +++ b/derive-impl/src/pystructseq.rs @@ -80,16 +80,19 @@ fn field_names(input: &mut DeriveInput) -> Result<(Vec, Vec)> { } pub(crate) fn impl_pystruct_sequence(mut input: DeriveInput) -> Result { - let (not_skipped_fields, _skipped_fields) = field_names(&mut input)?; + let (not_skipped_fields, skipped_fields) = field_names(&mut input)?; let ty = &input.ident; let ret = quote! { impl ::rustpython_vm::types::PyStructSequence for #ty { - const FIELD_NAMES: &'static [&'static str] = &[#(stringify!(#not_skipped_fields)),*]; + const REQUIRED_FIELD_NAMES: &'static [&'static str] = &[#(stringify!(#not_skipped_fields),)*]; + const OPTIONAL_FIELD_NAMES: &'static [&'static str] = &[#(stringify!(#skipped_fields),)*]; fn into_tuple(self, vm: &::rustpython_vm::VirtualMachine) -> ::rustpython_vm::builtins::PyTuple { - let items = vec![#(::rustpython_vm::convert::ToPyObject::to_pyobject( - self.#not_skipped_fields, - vm, - )),*]; + let items = vec![ + #(::rustpython_vm::convert::ToPyObject::to_pyobject( + self.#not_skipped_fields, + vm, + ),)* + ]; ::rustpython_vm::builtins::PyTuple::new_unchecked(items.into_boxed_slice()) } } @@ -110,17 +113,14 @@ pub(crate) fn impl_pystruct_sequence_try_from_object( let ret = quote! { impl ::rustpython_vm::TryFromObject for #ty { fn try_from_object(vm: &::rustpython_vm::VirtualMachine, seq: ::rustpython_vm::PyObjectRef) -> ::rustpython_vm::PyResult { - const LEN: usize = #ty::FIELD_NAMES.len(); - let seq = Self::try_elements_from::(seq, vm)?; - // TODO: this is possible to be written without iterator + let seq = Self::try_elements_from(seq, vm)?; let mut iter = seq.into_iter(); Ok(Self { - #( - #not_skipped_fields: iter.next().unwrap().clone().try_into_value(vm)?, - )* - #( - #skipped_fields: vm.ctx.none(), - )* + #(#not_skipped_fields: iter.next().unwrap().clone().try_into_value(vm)?,)* + #(#skipped_fields: match iter.next() { + Some(v) => v.clone().try_into_value(vm)?, + None => vm.ctx.none(), + },)* }) } } diff --git a/vm/src/types/structseq.rs b/vm/src/types/structseq.rs index a0b445ce7..7b7f5fb9b 100644 --- a/vm/src/types/structseq.rs +++ b/vm/src/types/structseq.rs @@ -1,13 +1,14 @@ use crate::{ AsObject, Py, PyObjectRef, PyPayload, PyRef, PyResult, VirtualMachine, - builtins::{PyTuple, PyTupleRef, PyType}, + builtins::{PyBaseExceptionRef, PyTuple, PyTupleRef, PyType}, class::{PyClassImpl, StaticType}, vm::Context, }; #[pyclass] pub trait PyStructSequence: StaticType + PyClassImpl + Sized + 'static { - const FIELD_NAMES: &'static [&'static str]; + const REQUIRED_FIELD_NAMES: &'static [&'static str]; + const OPTIONAL_FIELD_NAMES: &'static [&'static str]; fn into_tuple(self, vm: &VirtualMachine) -> PyTuple; @@ -17,10 +18,16 @@ pub trait PyStructSequence: StaticType + PyClassImpl + Sized + 'static { .unwrap() } - fn try_elements_from( - obj: PyObjectRef, - vm: &VirtualMachine, - ) -> PyResult<[PyObjectRef; FIELD_LEN]> { + fn try_elements_from(obj: PyObjectRef, vm: &VirtualMachine) -> PyResult> { + #[cold] + fn sequence_length_error( + name: &str, + len: usize, + vm: &VirtualMachine, + ) -> PyBaseExceptionRef { + vm.new_type_error(format!("{name} takes a sequence of length {len}")) + } + let typ = Self::static_type(); // if !obj.fast_isinstance(typ) { // return Err(vm.new_type_error(format!( @@ -30,13 +37,13 @@ pub trait PyStructSequence: StaticType + PyClassImpl + Sized + 'static { // ))); // } let seq: Vec = obj.try_into_value(vm)?; - let seq: [PyObjectRef; FIELD_LEN] = seq.try_into().map_err(|_| { - vm.new_type_error(format!( - "{} takes a sequence of length {}", - typ.name(), - FIELD_LEN - )) - })?; + if seq.len() < Self::REQUIRED_FIELD_NAMES.len() { + return Err(sequence_length_error( + &typ.name(), + Self::REQUIRED_FIELD_NAMES.len(), + vm, + )); + } Ok(seq) } @@ -49,14 +56,14 @@ pub trait PyStructSequence: StaticType + PyClassImpl + Sized + 'static { let (body, suffix) = if let Some(_guard) = rustpython_vm::recursion::ReprGuard::enter(vm, zelf.as_object()) { - if Self::FIELD_NAMES.len() == 1 { + if Self::REQUIRED_FIELD_NAMES.len() == 1 { let value = zelf.first().unwrap(); - let formatted = format_field((value, Self::FIELD_NAMES[0]))?; + let formatted = format_field((value, Self::REQUIRED_FIELD_NAMES[0]))?; (formatted, ",") } else { let fields: PyResult> = zelf .iter() - .zip(Self::FIELD_NAMES.iter().copied()) + .zip(Self::REQUIRED_FIELD_NAMES.iter().copied()) .map(format_field) .collect(); (fields?.join(", "), "") @@ -74,7 +81,7 @@ pub trait PyStructSequence: StaticType + PyClassImpl + Sized + 'static { #[extend_class] fn extend_pyclass(ctx: &Context, class: &'static Py) { - for (i, &name) in Self::FIELD_NAMES.iter().enumerate() { + for (i, &name) in Self::REQUIRED_FIELD_NAMES.iter().enumerate() { // cast i to a u8 so there's less to store in the getter closure. // Hopefully there's not struct sequences with >=256 elements :P let i = i as u8; @@ -90,7 +97,7 @@ pub trait PyStructSequence: StaticType + PyClassImpl + Sized + 'static { class.set_attr( identifier!(ctx, __match_args__), ctx.new_tuple( - Self::FIELD_NAMES + Self::REQUIRED_FIELD_NAMES .iter() .map(|&name| ctx.new_str(name).into()) .collect::>(),