itertools icon indicating copy to clipboard operation
itertools copied to clipboard

Feature request: splitting

Open feois opened this issue 5 months ago • 1 comments

Add split() that produces an iterator that takes from original iterator until a specific condition. This fixes some issues I have with take_while(). (I find the method name confusing personally when I first encountered it) The iterator returned by split() (i.e. Splitter) continues taking from the original iterator even when not fully consumed (e.g. dropped), it basically "splits/partitions" the iterator into two (but not really cuz the original iterator cannot be used until Splitter is dropped).

Runnable example:

use itertools::PeekingNext;

struct Splitter<'a, I: PeekingNext, F: FnMut(&I::Item) -> bool> {
    iter: &'a mut I,
    cond: F,
}

impl<'a, I: PeekingNext, F: FnMut(&I::Item) -> bool> Iterator for Splitter<'a, I, F> {
    type Item = I::Item;
    
    fn next(&mut self) -> Option<I::Item> {
        self.iter.peeking_next(|t| (self.cond)(t))
    }
}

impl<'a, I: PeekingNext, F: FnMut(&I::Item) -> bool> Drop for Splitter<'a, I, F> {
    fn drop(&mut self) {
        for _ in self {}
    }
}

trait Split: PeekingNext + Sized {
    fn split<'a, F: FnMut(&Self::Item) -> bool>(&'a mut self, condition: F) -> Splitter<'a, Self, F> {
        Splitter { iter: self, cond: condition }
    }
    
    fn split_n<'a>(&'a mut self, mut count: usize) -> Splitter<'a, Self, impl FnMut(&Self::Item) -> bool> {
        self.split(move |_| if count > 0 { count -= 1; true } else { false })
    }
}

impl<I: PeekingNext + Sized> Split for I {}

fn main() {
    let v = vec![0, 1, 2, 3, 4, 5];
    
    let mut i = v.iter();
    let j = i.split(|&&i| i < 3);
    
    println!("{:?}", j.collect::<Vec<_>>()); // [0, 1, 2]
    println!("{:?}", i.collect::<Vec<_>>()); // [3, 4, 5]
    
    let mut i = v.iter();
    let mut j = i.split(|&&i| i < 3);
    
    println!("{:?}", j.next().unwrap()); // 0
    std::mem::drop(j);
    // equivalent code
    // println!("{:?}", i.split(|&&i| i < 3).next().unwrap());
    println!("{:?}", i.collect::<Vec<_>>()); // [3, 4, 5]
    
    let mut i = v.iter();
    let j = i.split_n(3);
    
    println!("{:?}", j.collect::<Vec<_>>()); // [0, 1, 2]
    println!("{:?}", i.collect::<Vec<_>>()); // [3, 4, 5]
    
    let mut i = v.iter();
    
    println!("{:?}", i.split_n(3).next().unwrap()); // 0
    println!("{:?}", i.collect::<Vec<_>>()); // [3, 4, 5]
}

feois avatar Jun 22 '25 16:06 feois

Some other possible names: partition (but used by std alr 😔), take_until (might be confused with take_while)

feois avatar Jun 22 '25 16:06 feois