Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion crates/vm/src/builtins/bytearray.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1504,7 +1504,7 @@ impl AsSequence for PyByteArray {
PyByteArray::sequence_downcast(seq)
.borrow_buf()
.getitem_by_index(vm, i)
.map(|x| vm.ctx.new_bytes(vec![x]).into())
.map(|x| vm.ctx.new_int(x).into())
}),
ass_item: atomic_func!(|seq, i, value, vm| {
let zelf = PyByteArray::sequence_downcast(seq);
Expand Down
2 changes: 1 addition & 1 deletion crates/vm/src/builtins/bytes.rs
Original file line number Diff line number Diff line change
Expand Up @@ -698,7 +698,7 @@ impl AsSequence for PyBytes {
PyBytes::sequence_downcast(seq)
.as_bytes()
.getitem_by_index(vm, i)
.map(|x| vm.ctx.new_bytes(vec![x]).into())
.map(|x| vm.ctx.new_int(x).into())
}),
contains: atomic_func!(|seq, other, vm| {
PyBytes::sequence_downcast(seq).__contains__(other.to_owned(), vm)
Expand Down
127 changes: 99 additions & 28 deletions crates/vm/src/builtins/enumerate.rs
Original file line number Diff line number Diff line change
@@ -1,15 +1,15 @@
use super::{
IterStatus, PositionIterInternal, PyGenericAlias, PyTupleRef, PyType, PyTypeRef,
iter::builtins_reversed, locked_rev_next,
PyGenericAlias, PyTupleRef, PyType, PyTypeRef,
iter::{SequenceIter, SequenceIterDirection},
};
use crate::common::lock::{PyMutex, PyRwLock};
use crate::common::lock::PyRwLock;
use crate::{
AsObject, Context, Py, PyObjectRef, PyPayload, PyResult, VirtualMachine,
class::PyClassImpl,
function::ArgIndex,
function::{ArgIndex, FuncArgs},
protocol::{PyIter, PyIterReturn},
raise_if_stop,
types::{Constructor, IterNext, Iterable, SelfIter},
types::{Constructor, IterNext, Iterable, SelfIter, drop_kwargs_if_init_overridden},
};
use malachite_bigint::BigInt;
use num_traits::ToPrimitive;
Expand Down Expand Up @@ -139,7 +139,7 @@ impl IterNext for PyEnumerate {
#[pyclass(module = false, name = "reversed", traverse)]
#[derive(Debug)]
pub(crate) struct PyReverseSequenceIterator {
internal: PyMutex<PositionIterInternal<PyObjectRef>>,
internal: SequenceIter,
}

impl PyPayload for PyReverseSequenceIterator {
Expand All @@ -150,50 +150,121 @@ impl PyPayload for PyReverseSequenceIterator {
}

impl PyReverseSequenceIterator {
pub(crate) fn from_object(
cls: PyTypeRef,
sequence: &PyObjectRef,
vm: &VirtualMachine,
) -> PyResult {
// Keep the original sequence alive while binding, calling, and releasing
// the special method, including when it does not retain its receiver.
if let Some(method) = vm.get_method(sequence.clone(), identifier!(vm, __reversed__)) {
let method = method?;
if vm.is_none(&method) {
return Err(Self::not_reversible(sequence, vm));
}
let result = method.call((), vm);
drop(method);
return result;
}

Self::py_new(
&cls,
ReversedArgs {
object: sequence.clone(),
},
vm,
)?
.into_ref_with_type(vm, cls)
.map(Into::into)
}

fn not_reversible(
sequence: &PyObjectRef,
vm: &VirtualMachine,
) -> crate::builtins::PyBaseExceptionRef {
vm.new_type_error(format!(
"'{}' object is not reversible",
sequence.class().slot_name()
))
}

pub(crate) const fn new(obj: PyObjectRef, len: usize) -> Self {
let position = len.saturating_sub(1);
let position = len as isize - 1;
Self {
internal: PyMutex::new(PositionIterInternal::new(obj, position)),
internal: SequenceIter::new(obj, position, SequenceIterDirection::Reverse),
}
}
}

#[pyclass(with(IterNext, Iterable))]
#[derive(FromArgs)]
pub(crate) struct ReversedArgs {
#[pyarg(positional)]
object: PyObjectRef,
}

impl Constructor for PyReverseSequenceIterator {
type Args = ReversedArgs;

fn slot_new(cls: PyTypeRef, args: FuncArgs, vm: &VirtualMachine) -> PyResult {
// A subtype's custom initializer may consume keywords, but the exact
// type and subtypes inheriting its initializer reject them first.
// Retain the original arguments through callbacks even if keywords
// are removed from the constructor's view of the call.
let constructor_args =
drop_kwargs_if_init_overridden(&cls, vm.ctx.types.reverse_iter_type, args.clone());
if !constructor_args.kwargs.is_empty() {
return Err(vm.new_type_error("reversed() takes no keyword arguments"));
}
let ReversedArgs { object } = constructor_args.bind_for(vm, "reversed")?;
let result = Self::from_object(cls, &object, vm);
drop(object);
drop(args);
result
}

fn py_new(
_cls: &Py<PyType>,
Self::Args { object }: Self::Args,
vm: &VirtualMachine,
) -> PyResult<Self> {
// The slot handles arbitrary __reversed__ results before constructing
// a payload. Checking sequence support must not bind __getitem__.
let protocol = object.sequence_unchecked();
if !protocol.check() {
return Err(Self::not_reversible(&object, vm));
}
let len = protocol.length(vm)?;
Ok(Self::new(object, len))
}
}

#[pyclass(with(IterNext, Iterable, Constructor), flags(BASETYPE))]
impl Py<PyReverseSequenceIterator> {
#[pymethod]
fn __length_hint__(&self, vm: &VirtualMachine) -> PyResult<usize> {
let internal = self.internal.lock();
if let IterStatus::Active(obj) = &internal.status
&& internal.position <= obj.length(vm)?
{
return Ok(internal.position + 1);
}
Ok(0)
vm.with_recursion(" in __length_hint__", || {
self.internal
.length_hint(vm)
.map(|length| length.unwrap_or(0))
})
}

#[pymethod]
fn __setstate__(&self, state: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> {
self.internal.lock().set_state(&state, |_, pos| pos, vm)
self.internal.set_state(&state, vm)
}

#[pymethod]
fn __reduce__(&self, vm: &VirtualMachine) -> PyResult<PyTupleRef> {
let func = builtins_reversed(vm)?;
Ok(self.internal.lock().reduce(
func,
|x| x.clone(),
|vm| vm.ctx.empty_tuple.clone().into(),
vm,
))
fn __reduce__(&self, vm: &VirtualMachine) -> PyTupleRef {
let func = self.class().to_owned().into();
self.internal.reduce(func, vm)
}
}

impl SelfIter for PyReverseSequenceIterator {}
impl IterNext for PyReverseSequenceIterator {
fn next(zelf: &Py<Self>, vm: &VirtualMachine) -> PyResult<PyIterReturn> {
locked_rev_next(&zelf.internal, |obj, pos| {
PyIterReturn::from_getitem_result(obj.get_item(&pos, vm), vm)
})
zelf.internal.next(vm)
}
}

Expand Down
35 changes: 18 additions & 17 deletions crates/vm/src/builtins/filter.rs
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
use super::{PyType, PyTypeRef};
use crate::{
Context, Py, PyObjectRef, PyPayload, PyResult, VirtualMachine,
AsObject, Context, Py, PyObjectRef, PyPayload, PyResult, VirtualMachine,
class::PyClassImpl,
protocol::{PyIter, PyIterReturn},
raise_if_stop,
Expand All @@ -26,7 +26,7 @@ pub struct FilterArgs {
#[pyarg(positional)]
function: PyObjectRef,
#[pyarg(positional)]
iterable: PyIter,
iterable: PyObjectRef,
}

impl Constructor for PyFilter {
Expand All @@ -35,25 +35,25 @@ impl Constructor for PyFilter {

fn py_new(
_cls: &Py<PyType>,
Self::Args {
function,
iterable: iterator,
}: Self::Args,
_vm: &VirtualMachine,
Self::Args { function, iterable }: Self::Args,
vm: &VirtualMachine,
) -> PyResult<Self> {
Ok(Self {
predicate: function,
let result = iterable.get_iter(vm).map(|iterator| Self {
predicate: function.clone(),
iterator,
})
});
drop(iterable);
drop(function);
result
}
}

#[pyclass(with(IterNext, Iterable, Constructor), flags(BASETYPE))]
impl Py<PyFilter> {
#[pymethod]
fn __reduce__(&self, vm: &VirtualMachine) -> (PyTypeRef, (PyObjectRef, PyIter)) {
fn __reduce__(&self) -> (PyTypeRef, (PyObjectRef, PyIter)) {
(
vm.ctx.types.filter_type.to_owned(),
self.class().to_owned(),
(self.predicate.clone(), self.iterator.clone()),
)
}
Expand All @@ -66,17 +66,18 @@ impl IterNext for PyFilter {
let predicate = &zelf.predicate;
loop {
let next_obj = raise_if_stop!(zelf.iterator.next(vm)?);
let predicate_value = if vm.is_none(predicate) {
next_obj.clone()
let is_true = if vm.is_none(predicate) || predicate.is(vm.ctx.types.bool_type) {
next_obj.try_to_bool(vm)?
} else {
// the predicate itself can raise StopIteration which does stop the filter iteration
raise_if_stop!(PyIterReturn::from_pyresult(
let predicate_value = raise_if_stop!(PyIterReturn::from_pyresult(
predicate.call((next_obj.clone(),), vm),
vm
)?)
)?);
predicate_value.try_to_bool(vm)?
};

if predicate_value.try_to_bool(vm)? {
if is_true {
return Ok(PyIterReturn::Return(next_obj));
}
}
Expand Down
Loading
Loading