Skip to content

Fix itertools.count to take PyNumber instead of PyInt #3822

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
merged 8 commits into from
Jun 27, 2022
Merged

Fix itertools.count to take PyNumber instead of PyInt #3822

Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
aad60e8
Fix to get number of itertools count
yonmilk Jun 26, 2022
79cb376
Fix itertools count
yonmilk Jun 27, 2022
d3485f4
Fix itertools count
yonmilk Jun 27, 2022
e7c1eb2
Fix itertools count
yonmilk Jun 27, 2022
89ae703
Fix itertools count step
yonmilk Jun 27, 2022
4422a54
Fix itertools count PyResult
yonmilk Jun 27, 2022
e1212b7
Fix itertools count next
yonmilk Jun 27, 2022
fd85e71
Fix clippy error
yonmilk Jun 27, 2022
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
29 changes: 13 additions & 16 deletions vm/src/stdlib/itertools.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,14 +11,13 @@ mod decl {
convert::ToPyObject,
function::{ArgCallable, FuncArgs, OptionalArg, OptionalOption, PosArgs},
identifier,
protocol::{PyIter, PyIterReturn},
protocol::{PyIter, PyIterReturn, PyNumber},
stdlib::sys,
types::{Constructor, IterNext, IterNextIterable},
AsObject, Py, PyObjectRef, PyPayload, PyRef, PyResult, PyWeakRef, VirtualMachine,
};
use crossbeam_utils::atomic::AtomicCell;
use num_bigint::BigInt;
use num_traits::{One, Signed, ToPrimitive, Zero};
use num_traits::{Signed, ToPrimitive};
use std::fmt;

#[pyattr]
Expand Down Expand Up @@ -174,14 +173,14 @@ mod decl {
#[pyclass(name = "count")]
#[derive(Debug, PyPayload)]
struct PyItertoolsCount {
cur: PyRwLock<BigInt>,
step: BigInt,
cur: PyRwLock<PyObjectRef>,
step: PyIntRef,
}

#[derive(FromArgs)]
struct CountNewArgs {
#[pyarg(positional, optional)]
start: OptionalArg<PyIntRef>,
start: OptionalArg<PyObjectRef>,

#[pyarg(positional, optional)]
step: OptionalArg<PyIntRef>,
Expand All @@ -195,14 +194,11 @@ mod decl {
Self::Args { start, step }: Self::Args,
vm: &VirtualMachine,
) -> PyResult {
let start = match start.into_option() {
Some(int) => int.as_bigint().clone(),
None => BigInt::zero(),
};
let step = match step.into_option() {
Some(int) => int.as_bigint().clone(),
None => BigInt::one(),
};
let start: PyObjectRef = start.into_option().unwrap_or_else(|| vm.new_pyobj(0));
let step: PyIntRef = step.into_option().unwrap_or_else(|| vm.new_pyref(1));
if !PyNumber::check(&start, vm) {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
if !PyNumber::check(&start, vm) {
if !PyNumber::check(&start, vm) || !PyNumber::check(&step, vm) {

It seems step also need to be checked via PyNumber::check.
https://github.com/python/cpython/blob/3.10/Modules/itertoolsmodule.c#L4178-L4182

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I didn't because it's a PyIntRef. I'll fix it.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

PyIntRef is always a PyNumber.

return Err(vm.new_value_error("a number is require".to_owned()));
}

PyItertoolsCount {
cur: PyRwLock::new(start),
Expand All @@ -219,7 +215,7 @@ mod decl {
// if (lz->cnt == PY_SSIZE_T_MAX)
// return Py_BuildValue("0(00)", Py_TYPE(lz), lz->long_cnt, lz->long_step);
#[pymethod(magic)]
fn reduce(zelf: PyRef<Self>) -> (PyTypeRef, (BigInt,)) {
fn reduce(zelf: PyRef<Self>) -> (PyTypeRef, (PyObjectRef,)) {
(zelf.class().clone(), (zelf.cur.read().clone(),))
}

Expand All @@ -234,8 +230,9 @@ mod decl {
impl IterNext for PyItertoolsCount {
fn next(zelf: &Py<Self>, vm: &VirtualMachine) -> PyResult<PyIterReturn> {
let mut cur = zelf.cur.write();
let step = zelf.step.clone();
let result = cur.clone();
*cur += &zelf.step;
*cur = vm._iadd(&*cur, step.as_object())?;
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

is iadd intended to be used with return value?

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

oh, it does. I didn't know that.

Ok(PyIterReturn::Return(result.to_pyobject(vm)))
}
}
Expand Down