Skip to content

Implement missing methods for str iterator #2906

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
DimitrisJim merged 3 commits into from
Aug 18, 2021
Merged

Implement missing methods for str iterator #2906

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
93df314
builtins: impl missing methods of iterators in str
Snowapril Aug 18, 2021
04d644d
builtins: refactor reduce method body
Snowapril Aug 18, 2021
2cdb416
builtins: remove useless blank line
Snowapril Aug 18, 2021
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: 0 additions & 2 deletions Lib/test/test_iter.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -313,8 +313,6 @@ def test_iter_range(self):
self.check_for_loop(iter(range(10)), list(range(10)))

# Test a string
# TODO: RUSTPYTHON
@unittest.expectedFailure
def test_iter_string(self):
self.check_for_loop(iter("abcde"), ["a", "b", "c", "d", "e"])

Expand Down
65 changes: 59 additions & 6 deletions vm/src/builtins/pystr.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@ use std::ops::Range;
use std::string::ToString;
use std::{char, ffi, fmt};

use crossbeam_utils::atomic::AtomicCell;
use itertools::Itertools;
use num_traits::ToPrimitive;
use unic_ucd_bidi::BidiClass;
Expand All @@ -12,7 +13,11 @@ use unicode_casing::CharExt;

use super::bytes::PyBytesRef;
use super::dict::PyDict;
use super::int::{PyInt, PyIntRef};
use super::int::{try_to_primitive, PyInt, PyIntRef};
use super::iter::{
IterStatus,
IterStatus::{Active, Exhausted},
};
use super::pytype::PyTypeRef;
use crate::anystr::{self, adjust_indices, AnyStr, AnyStrContainer, AnyStrWrapper};
use crate::exceptions::IntoPyException;
Expand Down Expand Up @@ -113,6 +118,7 @@ impl TryIntoRef<PyStr> for &str {
pub struct PyStrIterator {
string: PyStrRef,
position: PyAtomic<usize>,
status: AtomicCell<IterStatus>,
}

impl PyValue for PyStrIterator {
Expand All @@ -122,20 +128,66 @@ impl PyValue for PyStrIterator {
}

#[pyimpl(with(PyIter))]
impl PyStrIterator {}
impl PyStrIterator {
#[pymethod(magic)]
fn length_hint(&self) -> usize {
match self.status.load() {
Active => {
let pos = self.position.load(atomic::Ordering::SeqCst);
self.string.len().saturating_sub(pos)
}
Exhausted => 0,
}
}

#[pymethod(magic)]
fn setstate(&self, state: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> {
// When we're exhausted, just return.
if let Exhausted = self.status.load() {
return Ok(());
}
let pos = state
.payload::<PyInt>()
.ok_or_else(|| vm.new_type_error("an integer is required.".to_owned()))?;
let pos = std::cmp::min(
try_to_primitive(pos.as_bigint(), vm).unwrap_or(0),
self.string.len(),
);
self.position.store(pos, atomic::Ordering::SeqCst);
Ok(())
}

#[pymethod(magic)]
fn reduce(&self, vm: &VirtualMachine) -> PyResult {
let iter = vm.get_attribute(vm.builtins.clone(), "iter")?;
Ok(vm.ctx.new_tuple(match self.status.load() {
Exhausted => vec![iter, vm.ctx.new_tuple(vec![vm.ctx.new_str("")])],
Active => vec![
iter,
vm.ctx.new_tuple(vec![self.string.clone().into_object()]),
vm.ctx
.new_int(self.position.load(atomic::Ordering::Relaxed)),
],
}))
}
}

impl PyIter for PyStrIterator {
fn next(zelf: &PyRef<Self>, vm: &VirtualMachine) -> PyResult {
if let Exhausted = zelf.status.load() {
return Err(vm.new_stop_iteration());
}
let value = &*zelf.string.value;
let mut start = zelf.position.load(atomic::Ordering::SeqCst);
loop {
if start == value.len() {
zelf.status.store(Exhausted);
return Err(vm.new_stop_iteration());
}
let ch = value[start..]
.chars()
.next()
.ok_or_else(|| vm.new_stop_iteration())?;
let ch = value[start..].chars().next().ok_or_else(|| {
zelf.status.store(Exhausted);
vm.new_stop_iteration()
})?;

match zelf.position.compare_exchange_weak(
start,
Expand Down Expand Up @@ -1123,6 +1175,7 @@ impl Iterable for PyStr {
Ok(PyStrIterator {
position: Radium::new(0),
string: zelf,
status: AtomicCell::new(Active),
}
.into_object(vm))
}
Expand Down