diff --git a/vm/src/protocol/sequence.rs b/vm/src/protocol/sequence.rs index 9951fa7e8..90c115d29 100644 --- a/vm/src/protocol/sequence.rs +++ b/vm/src/protocol/sequence.rs @@ -366,6 +366,48 @@ impl PySequence<'_> { Err(vm.new_value_error("sequence.index(x): x not in sequence".to_string())) } + + pub fn extract(&self, mut f: F, vm: &VirtualMachine) -> PyResult> + where + F: FnMut(&PyObject) -> PyResult, + { + if let Some(tuple) = self.obj.payload_if_exact::(vm) { + tuple.as_slice().iter().map(|x| f(x.as_ref())).collect() + } else if let Some(list) = self.obj.payload_if_exact::(vm) { + list.borrow_vec().iter().map(|x| f(x.as_ref())).collect() + } else { + let iter = self.obj.to_owned().get_iter(vm)?; + let iter = iter.iter::(vm)?; + let len = self.length(vm).unwrap_or(0); + let mut v = Vec::with_capacity(len); + for x in iter { + v.push(f(x?.as_ref())?); + } + v.shrink_to_fit(); + Ok(v) + } + } + + pub fn extract_cloned(&self, mut f: F, vm: &VirtualMachine) -> PyResult> + where + F: FnMut(PyObjectRef) -> PyResult, + { + if let Some(tuple) = self.obj.payload_if_exact::(vm) { + tuple.as_slice().iter().map(|x| f(x.clone())).collect() + } else if let Some(list) = self.obj.payload_if_exact::(vm) { + list.borrow_vec().iter().map(|x| f(x.clone())).collect() + } else { + let iter = self.obj.to_owned().get_iter(vm)?; + let iter = iter.iter::(vm)?; + let len = self.length(vm).unwrap_or(0); + let mut v = Vec::with_capacity(len); + for x in iter { + v.push(f(x?)?); + } + v.shrink_to_fit(); + Ok(v) + } + } } const NOT_IMPLEMENTED: PySequenceMethods = PySequenceMethods {