diff --git a/crates/vm/src/stdlib/itertools.rs b/crates/vm/src/stdlib/itertools.rs index 4433cf470..e72b4f08d 100644 --- a/crates/vm/src/stdlib/itertools.rs +++ b/crates/vm/src/stdlib/itertools.rs @@ -224,7 +224,7 @@ mod decl { } } - #[pyclass(with(IterNext, Iterable, Constructor, Representable))] + #[pyclass(with(IterNext, Iterable, Constructor, Representable), flags(BASETYPE))] impl PyItertoolsCount {} impl SelfIter for PyItertoolsCount {} @@ -638,7 +638,7 @@ mod decl { } } - #[pyclass(with(IterNext, Iterable, Constructor))] + #[pyclass(with(IterNext, Iterable, Constructor), flags(BASETYPE))] impl PyItertoolsGroupBy { pub(super) fn advance( &self, @@ -989,7 +989,7 @@ mod decl { } } - #[pyclass(with(IterNext, Iterable, Constructor))] + #[pyclass(with(IterNext, Iterable, Constructor), flags(BASETYPE))] impl PyItertoolsAccumulate {} impl SelfIter for PyItertoolsAccumulate {} @@ -1251,7 +1251,7 @@ mod decl { } } - #[pyclass(with(IterNext, Iterable, Constructor))] + #[pyclass(with(IterNext, Iterable, Constructor), flags(BASETYPE))] impl PyItertoolsProduct { fn update_idxs(&self, mut idxs: PyRwLockWriteGuard<'_, Vec>) { if idxs.is_empty() { @@ -1367,7 +1367,7 @@ mod decl { } } - #[pyclass(with(IterNext, Iterable, Constructor))] + #[pyclass(with(IterNext, Iterable, Constructor), flags(BASETYPE))] impl PyItertoolsCombinations {} impl SelfIter for PyItertoolsCombinations {} @@ -1479,7 +1479,7 @@ mod decl { } } - #[pyclass(with(IterNext, Iterable, Constructor))] + #[pyclass(with(IterNext, Iterable, Constructor), flags(BASETYPE))] impl PyItertoolsCombinationsWithReplacement {} impl SelfIter for PyItertoolsCombinationsWithReplacement {} @@ -1595,7 +1595,7 @@ mod decl { } } - #[pyclass(with(IterNext, Iterable, Constructor))] + #[pyclass(with(IterNext, Iterable, Constructor), flags(BASETYPE))] impl PyItertoolsPermutations {} impl SelfIter for PyItertoolsPermutations {} @@ -1702,7 +1702,7 @@ mod decl { fillvalue: PyRwLock, } - #[pyclass(with(IterNext, Iterable, Constructor))] + #[pyclass(with(IterNext, Iterable, Constructor), flags(BASETYPE))] impl PyItertoolsZipLongest {} impl SelfIter for PyItertoolsZipLongest {} @@ -1743,6 +1743,8 @@ mod decl { impl Constructor for PyItertoolsPairwise { type Args = IterablePosArg; + const DROP_KWARGS_WHEN_INIT_OVERRIDDEN: bool = true; + fn py_new(_cls: &Py, args: Self::Args, _vm: &VirtualMachine) -> PyResult { let iterator = args.iterable; Ok(Self { @@ -1752,7 +1754,7 @@ mod decl { } } - #[pyclass(with(IterNext, Iterable, Constructor))] + #[pyclass(with(IterNext, Iterable, Constructor), flags(BASETYPE))] impl PyItertoolsPairwise {} impl SelfIter for PyItertoolsPairwise {} diff --git a/extra_tests/snippets/stdlib_itertools.py b/extra_tests/snippets/stdlib_itertools.py index 04917fe99..38e213ea9 100644 --- a/extra_tests/snippets/stdlib_itertools.py +++ b/extra_tests/snippets/stdlib_itertools.py @@ -648,3 +648,34 @@ def check_compress_consumption(): check_compress_consumption() + + +# These public iterator types also support Python subclasses. +for base, args in ( + (itertools.count, ()), + (itertools.groupby, ([1],)), + (itertools.accumulate, ([1],)), + (itertools.product, ([1],)), + (itertools.combinations, ([1], 1)), + (itertools.combinations_with_replacement, ([1], 1)), + (itertools.permutations, ([1],)), + (itertools.zip_longest, ([1],)), + (itertools.pairwise, ([1, 2],)), +): + + class Subclass(base): + pass + + instance = Subclass(*args) + assert type(instance) is Subclass + next(instance) + + +class PairwiseWithLabel(itertools.pairwise): + def __init__(self, iterable, *, label): + self.label = label + + +pairs = PairwiseWithLabel([1, 2], label="pairs") +assert pairs.label == "pairs" +assert next(pairs) == (1, 2)