Add itertools adaptors: takewhile, dropwhile, filterfalse, starmap

crates/monty-proto/src/python/convert.rs

@@ -237,6 +237,10 @@ fn round_trip_type_table(py: Python<'_>) -> PyResult<&'static Vec<(Py<PyAny>, Mo
             MontyType::ItertoolsIslice,
             MontyType::ItertoolsChain,
             MontyType::ItertoolsCycle,
+            MontyType::ItertoolsTakeWhile,
+            MontyType::ItertoolsDropWhile,
+            MontyType::ItertoolsFilterFalse,
+            MontyType::ItertoolsStarMap,
             MontyType::Tuple,
             MontyType::Dict,
             MontyType::Set,
@@ -459,6 +463,10 @@ fn type_object_to_py(py: Python<'_>, t: MontyType) -> PyResult<Py<PyAny>> {
         MontyType::ItertoolsIslice => cached!("itertools", "islice"),
         MontyType::ItertoolsChain => cached!("itertools", "chain"),
         MontyType::ItertoolsCycle => cached!("itertools", "cycle"),
+        MontyType::ItertoolsTakeWhile => cached!("itertools", "takewhile"),
+        MontyType::ItertoolsDropWhile => cached!("itertools", "dropwhile"),
+        MontyType::ItertoolsFilterFalse => cached!("itertools", "filterfalse"),
+        MontyType::ItertoolsStarMap => cached!("itertools", "starmap"),
         // Consistent with the Path *instance* arm, which marshals as PurePosixPath
         // and is instantiable on every host OS (unlike PosixPath on Windows).
         MontyType::Path => get_pure_posix_path(py).map(|b| b.clone().unbind()),

crates/monty-types/src/object.rs

@@ -830,6 +830,14 @@ pub enum MontyType {
     ItertoolsCycle,
     #[strum(serialize = "NotImplementedType")]
     NotImplementedType,
+    #[strum(serialize = "itertools.takewhile")]
+    ItertoolsTakeWhile,
+    #[strum(serialize = "itertools.dropwhile")]
+    ItertoolsDropWhile,
+    #[strum(serialize = "itertools.filterfalse")]
+    ItertoolsFilterFalse,
+    #[strum(serialize = "itertools.starmap")]
+    ItertoolsStarMap,
 }
 
 impl fmt::Display for MontyType {

crates/monty-typeshed/custom/itertools.pyi

@@ -10,8 +10,8 @@
 # because Monty's `count()` rejects anything else, where CPython accepts any
 # number protocol.
 
-from collections.abc import Iterable
-from typing import Generic, TypeAlias, TypeVar, overload
+from collections.abc import Callable, Iterable
+from typing import Any, Generic, TypeAlias, TypeVar, overload
 
 from typing_extensions import Self
 
@@ -65,3 +65,25 @@ class cycle(Generic[_T]):
     def __new__(cls, iterable: Iterable[_T], /) -> Self: ...
     def __next__(self) -> _T: ...
     def __iter__(self) -> Self: ...
+
+class takewhile(Generic[_T]):
+    def __new__(cls, predicate: Callable[[_T], object], iterable: Iterable[_T], /) -> Self: ...
+    def __next__(self) -> _T: ...
+    def __iter__(self) -> Self: ...
+
+class dropwhile(Generic[_T]):
+    def __new__(cls, predicate: Callable[[_T], object], iterable: Iterable[_T], /) -> Self: ...
+    def __next__(self) -> _T: ...
+    def __iter__(self) -> Self: ...
+
+class filterfalse(Generic[_T]):
+    def __new__(cls, predicate: Callable[[_T], object] | None, iterable: Iterable[_T], /) -> Self: ...
+    def __next__(self) -> _T: ...
+    def __iter__(self) -> Self: ...
+
+# Upstream typeshed overloads `starmap` per arity; Monty spreads each item
+# dynamically, so one `Callable[..., _T]` covers every shape it accepts.
+class starmap(Generic[_T]):
+    def __new__(cls, function: Callable[..., _T], iterable: Iterable[Iterable[Any]], /) -> Self: ...
+    def __next__(self) -> _T: ...
+    def __iter__(self) -> Self: ...

crates/monty-typeshed/vendor/typeshed/stdlib/itertools.pyi

@@ -10,8 +10,8 @@
 # because Monty's `count()` rejects anything else, where CPython accepts any
 # number protocol.
 
-from collections.abc import Iterable
-from typing import Generic, TypeAlias, TypeVar, overload
+from collections.abc import Callable, Iterable
+from typing import Any, Generic, TypeAlias, TypeVar, overload
 
 from typing_extensions import Self
 
@@ -65,3 +65,25 @@ class cycle(Generic[_T]):
     def __new__(cls, iterable: Iterable[_T], /) -> Self: ...
     def __next__(self) -> _T: ...
     def __iter__(self) -> Self: ...
+
+class takewhile(Generic[_T]):
+    def __new__(cls, predicate: Callable[[_T], object], iterable: Iterable[_T], /) -> Self: ...
+    def __next__(self) -> _T: ...
+    def __iter__(self) -> Self: ...
+
+class dropwhile(Generic[_T]):
+    def __new__(cls, predicate: Callable[[_T], object], iterable: Iterable[_T], /) -> Self: ...
+    def __next__(self) -> _T: ...
+    def __iter__(self) -> Self: ...
+
+class filterfalse(Generic[_T]):
+    def __new__(cls, predicate: Callable[[_T], object] | None, iterable: Iterable[_T], /) -> Self: ...
+    def __next__(self) -> _T: ...
+    def __iter__(self) -> Self: ...
+
+# Upstream typeshed overloads `starmap` per arity; Monty spreads each item
+# dynamically, so one `Callable[..., _T]` covers every shape it accepts.
+class starmap(Generic[_T]):
+    def __new__(cls, function: Callable[..., _T], iterable: Iterable[Iterable[Any]], /) -> Self: ...
+    def __next__(self) -> _T: ...
+    def __iter__(self) -> Self: ...

crates/monty/src/builtins/filter.rs

@@ -5,14 +5,15 @@
 //! - `None` as predicate (filters falsy values)
 //! - Builtin functions (len, abs, etc.)
 //! - Type constructors (int, str, float, etc.)
-//! - User-defined functions (via `vm.evaluate_function`)
+//! - User-defined functions (via `call_predicate`)
 
 use crate::{
     args::ArgValues,
     bytecode::VM,
     defer_drop,
     exception_private::RunResult,
     heap::{DropGuard, HeapData},
+    predicate::call_predicate,
     types::{List, PyTrait},
     value::Value,
 };
@@ -48,12 +49,7 @@ pub fn builtin_filter(vm: &mut VM<'_>, args: ArgValues) -> RunResult<Value> {
             // No predicate - use truthiness of element
             item.py_bool(vm)?
         } else {
-            // Clone for predicate call - the clone is consumed by evaluate_function
-            let item_for_predicate = item.clone_with_heap(vm);
-            let result = vm.evaluate_function("filter()", function, ArgValues::One(item_for_predicate))?;
-            let is_truthy = result.py_bool(vm);
-            result.drop_with(vm);
-            is_truthy?
+            call_predicate(function, item, "filter()", vm)?
         };
 
         if should_include {

crates/monty/src/dump_format.rs

@@ -196,7 +196,7 @@ mod tests {
         );
         assert_eq!(
             static_strings_fingerprint(),
-            0x239c_c721_30be_eba3,
+            0xd864_f832_cad7_dea1,
             "static strings changed for dump version {DUMP_VERSION}"
         );
         assert_eq!(

crates/monty/src/intern.rs

@@ -935,6 +935,14 @@ pub enum StaticStrings {
     /// Python's `NotImplemented` singleton representation.
     #[strum(serialize = "NotImplemented")]
     NotImplementedRepr,
+    /// `itertools.takewhile()` function.
+    Takewhile,
+    /// `itertools.dropwhile()` function.
+    Dropwhile,
+    /// `itertools.filterfalse()` function.
+    Filterfalse,
+    /// `itertools.starmap()` function.
+    Starmap,
 }
 
 /// Computes an FNV-1a hash over static-string identities and serialization.

crates/monty/src/lib.rs

@@ -23,6 +23,7 @@ mod namespace;
 mod object_bridge;
 mod os_dispatch;
 mod parse;
+mod predicate;
 mod prepare;
 mod repl;
 mod resource_checks;

crates/monty/src/modules/itertools.rs

@@ -14,7 +14,9 @@ use crate::{
     modules::ModuleFunctions,
     types::{
         ItertoolsIter, Module, Type,
-        itertools::{Chain, Compress, Count, Cycle, Islice, Pairwise, Repeat},
+        itertools::{
+            Chain, Compress, Count, Cycle, DropWhile, FilterFalse, Islice, Pairwise, Repeat, StarMap, TakeWhile,
+        },
     },
     value::Value,
 };
@@ -30,6 +32,10 @@ pub(crate) enum ItertoolsFunctions {
     Islice,
     Chain,
     Cycle,
+    Takewhile,
+    Dropwhile,
+    Filterfalse,
+    Starmap,
 }
 
 /// Static mapping of attribute names to functions for module creation.
@@ -41,6 +47,10 @@ const ITERTOOLS_FUNCTIONS: &[(StaticStrings, ItertoolsFunctions)] = &[
     (StaticStrings::Islice, ItertoolsFunctions::Islice),
     (StaticStrings::Chain, ItertoolsFunctions::Chain),
     (StaticStrings::Cycle, ItertoolsFunctions::Cycle),
+    (StaticStrings::Takewhile, ItertoolsFunctions::Takewhile),
+    (StaticStrings::Dropwhile, ItertoolsFunctions::Dropwhile),
+    (StaticStrings::Filterfalse, ItertoolsFunctions::Filterfalse),
+    (StaticStrings::Starmap, ItertoolsFunctions::Starmap),
 ];
 
 /// Creates the `itertools` module on the heap.
@@ -67,6 +77,10 @@ pub(super) fn call(vm: &mut VM<'_>, function: ItertoolsFunctions, args: ArgValue
         ItertoolsFunctions::Islice => call_islice(vm, args),
         ItertoolsFunctions::Chain => call_chain(vm, args),
         ItertoolsFunctions::Cycle => call_cycle(vm, args),
+        ItertoolsFunctions::Takewhile => call_takewhile(vm, args),
+        ItertoolsFunctions::Dropwhile => call_dropwhile(vm, args),
+        ItertoolsFunctions::Filterfalse => call_filterfalse(vm, args),
+        ItertoolsFunctions::Starmap => call_starmap(vm, args),
     }
 }
 
@@ -359,3 +373,72 @@ fn call_cycle(vm: &mut VM<'_>, args: ArgValues) -> RunResult<Value> {
     let iter = ItertoolsIter::Cycle(Cycle::new(source));
     Ok(Value::Ref(vm.heap.allocate(HeapData::Itertools(iter))))
 }
+
+/// Argument shape shared by `takewhile`, `dropwhile`, `filterfalse` and
+/// `starmap`.
+///
+/// All four are `PyArg_UnpackTuple(args, name, 2, 2, ...)` in CPython, so both
+/// slots are positional-only, arity reads `takewhile expected 2 arguments, got
+/// 1`, and keywords are rejected wholesale. The macro embeds the name, hence
+/// one struct per callable rather than one shared struct.
+macro_rules! callable_and_iterable_args {
+    ($struct_name:ident, $py_name:literal) => {
+        #[derive(FromArgs)]
+        #[from_args(name = $py_name, style = unpack)]
+        struct $struct_name {
+            #[from_args(pos_only)]
+            callable: Value,
+            #[from_args(pos_only)]
+            iterable: Value,
+        }
+    };
+}
+
+callable_and_iterable_args!(TakeWhileArgs, "takewhile");
+callable_and_iterable_args!(DropWhileArgs, "dropwhile");
+callable_and_iterable_args!(FilterFalseArgs, "filterfalse");
+callable_and_iterable_args!(StarMapArgs, "starmap");
+
+/// `itertools.takewhile(predicate, iterable)` — the leading passing run.
+fn call_takewhile(vm: &mut VM<'_>, args: ArgValues) -> RunResult<Value> {
+    let TakeWhileArgs { callable, iterable } = TakeWhileArgs::from_args(args, vm)?;
+    let (predicate, source) = resolve_source(callable, iterable, vm)?;
+    let iter = ItertoolsIter::TakeWhile(TakeWhile::new(predicate, source));
+    Ok(Value::Ref(vm.heap.allocate(HeapData::Itertools(iter))))
+}
+
+/// `itertools.dropwhile(predicate, iterable)` — everything past that run.
+fn call_dropwhile(vm: &mut VM<'_>, args: ArgValues) -> RunResult<Value> {
+    let DropWhileArgs { callable, iterable } = DropWhileArgs::from_args(args, vm)?;
+    let (predicate, source) = resolve_source(callable, iterable, vm)?;
+    let iter = ItertoolsIter::DropWhile(DropWhile::new(predicate, source));
+    Ok(Value::Ref(vm.heap.allocate(HeapData::Itertools(iter))))
+}
+
+/// `itertools.filterfalse(predicate, iterable)` — the items it rejects.
+fn call_filterfalse(vm: &mut VM<'_>, args: ArgValues) -> RunResult<Value> {
+    let FilterFalseArgs { callable, iterable } = FilterFalseArgs::from_args(args, vm)?;
+    let (predicate, source) = resolve_source(callable, iterable, vm)?;
+    let iter = ItertoolsIter::FilterFalse(FilterFalse::new(predicate, source));
+    Ok(Value::Ref(vm.heap.allocate(HeapData::Itertools(iter))))
+}
+
+/// `itertools.starmap(function, iterable)` — each item spread as arguments.
+fn call_starmap(vm: &mut VM<'_>, args: ArgValues) -> RunResult<Value> {
+    let StarMapArgs { callable, iterable } = StarMapArgs::from_args(args, vm)?;
+    let (function, source) = resolve_source(callable, iterable, vm)?;
+    let iter = ItertoolsIter::StarMap(StarMap::new(function, source));
+    Ok(Value::Ref(vm.heap.allocate(HeapData::Itertools(iter))))
+}
+
+/// Resolves the iterable while keeping the callable safe from the error path.
+///
+/// CPython resolves eagerly for all four, so a non-iterable raises here rather
+/// than on the first `next()`. The callable itself is never type-checked: a
+/// non-callable is only discovered when the adaptor first applies it.
+fn resolve_source(callable: Value, iterable: Value, vm: &mut VM<'_>) -> RunResult<(Value, Value)> {
+    let mut guard = DropGuard::new(callable, vm);
+    let source = iterable.into_py_iter(guard.ctx())?;
+    let (callable, _) = guard.into_parts();
+    Ok((callable, source))
+}

crates/monty/src/object_bridge.rs

@@ -570,6 +570,10 @@ impl MontyTypeExt for MontyType {
             Self::ItertoolsIslice => Some(Type::ItertoolsIslice),
             Self::ItertoolsChain => Some(Type::ItertoolsChain),
             Self::ItertoolsCycle => Some(Type::ItertoolsCycle),
+            Self::ItertoolsTakeWhile => Some(Type::ItertoolsTakeWhile),
+            Self::ItertoolsDropWhile => Some(Type::ItertoolsDropWhile),
+            Self::ItertoolsFilterFalse => Some(Type::ItertoolsFilterFalse),
+            Self::ItertoolsStarMap => Some(Type::ItertoolsStarMap),
             Self::ItertoolsCount => Some(Type::ItertoolsCount),
             Self::ItertoolsRepeat => Some(Type::ItertoolsRepeat),
             Self::Tuple => Some(Type::Tuple),
@@ -647,6 +651,10 @@ impl MontyTypeExt for MontyType {
             Type::ItertoolsIslice => Self::ItertoolsIslice,
             Type::ItertoolsChain => Self::ItertoolsChain,
             Type::ItertoolsCycle => Self::ItertoolsCycle,
+            Type::ItertoolsTakeWhile => Self::ItertoolsTakeWhile,
+            Type::ItertoolsDropWhile => Self::ItertoolsDropWhile,
+            Type::ItertoolsFilterFalse => Self::ItertoolsFilterFalse,
+            Type::ItertoolsStarMap => Self::ItertoolsStarMap,
             Type::ItertoolsCount => Self::ItertoolsCount,
             Type::ItertoolsRepeat => Self::ItertoolsRepeat,
             Type::Tuple => Self::Tuple,

crates/monty/src/predicate.rs

@@ -0,0 +1,21 @@
+//! Calling a user-supplied predicate and taking its truthiness.
+//!
+//! Shared by `filter` and the `itertools` adaptors `takewhile`, `dropwhile`
+//! and `filterfalse`, which differ only in what they do with the answer.
+
+use crate::{args::ArgValues, bytecode::VM, defer_drop, exception_private::RunResult, types::PyTrait, value::Value};
+
+/// Applies `predicate` to `item` and returns its truthiness.
+///
+/// `item` is only borrowed — the call gets its own reference, so the caller can
+/// still yield the item afterwards. Goes through `evaluate_function`, which
+/// runs a defined function's frame to completion, so a predicate that suspends
+/// (an external function, an `os` call) is rejected rather than paused; `ctx`
+/// names the caller in that error.
+pub(crate) fn call_predicate(predicate: &Value, item: &Value, ctx: &'static str, vm: &mut VM<'_>) -> RunResult<bool> {
+    let arg = item.clone_with_heap(vm.heap);
+    let result = vm.evaluate_function(ctx, predicate, ArgValues::One(arg))?;
+    // Guarded because `py_bool` can raise through a user `__bool__`.
+    defer_drop!(result, vm);
+    result.py_bool(vm)
+}

crates/monty/src/types/iter.rs

@@ -85,7 +85,12 @@ impl<'h, I: CollectIter<'h>> Iterator for HeapedIterator<'_, 'h, I> {
             Ok(Some(value)) => {
                 self.yielded += 1;
                 let estimated = self.yielded.saturating_mul(VALUE_SIZE);
-                match check_estimated_size(estimated, &self.vm.heap.tracker) {
+                // Size alone does not bound a source whose items are cheap or
+                // interned, and the drain reaches no VM dispatch checkpoint of
+                // its own, so `max_duration` needs its own poll here.
+                let checked = check_estimated_size(estimated, &self.vm.heap.tracker)
+                    .and_then(|()| self.vm.heap.tracker.check_time_every(self.yielded));
+                match checked {
                     Ok(()) => Some(value),
                     Err(error) => {
                         *self.error = Some(error.into());

crates/monty/src/types/itertools/chain.rs

@@ -65,7 +65,14 @@ impl Chain {
 
 /// Drains the current source, then resolves the next one, until all are spent.
 pub(super) fn next<'h>(iter: &mut HeapRead<'h, ItertoolsIter>, vm: &mut VM<'h>) -> RunResult<Option<Value>> {
+    let mut steps = 0usize;
     loop {
+        // Native loop: the VM's dispatch checkpoint is per-`run()`, so a
+        // discarding pass over an infinite source reaches none. Poll the
+        // tracker so `max_duration` still bites (see `VM::run`'s
+        // `CHECK_INTERVAL`).
+        vm.heap.tracker.check_time_every(steps)?;
+        steps += 1;
         let ItertoolsIter::Chain(chain) = iter.get(vm.heap) else {
             unreachable!("dispatched on Kind::Chain")
         };

crates/monty/src/types/itertools/compress.rs

@@ -52,7 +52,14 @@ impl Compress {
 /// Pulls one item from each side per step, yielding the datum when its selector
 /// is truthy and stopping as soon as *either* side runs out.
 pub(super) fn next<'h>(iter: &mut HeapRead<'h, ItertoolsIter>, vm: &mut VM<'h>) -> RunResult<Option<Value>> {
+    let mut steps = 0usize;
     loop {
+        // Native loop: the VM's dispatch checkpoint is per-`run()`, so a
+        // discarding pass over an infinite source reaches none. Poll the
+        // tracker so `max_duration` still bites (see `VM::run`'s
+        // `CHECK_INTERVAL`).
+        vm.heap.tracker.check_time_every(steps)?;
+        steps += 1;
         let ItertoolsIter::Compress(compress) = iter.get(vm.heap) else {
             unreachable!("dispatched on Kind::Compress")
         };