Take the first N items of an iterable (or while a predicate holds). drop is the complement. Lazy via islice + itertools.dropwhile — works on infinite generators.
Python
from itertools import islice, dropwhile, takewhile, count
from typing import Iterable, Iterator, TypeVar, Callable
T = TypeVar("T")
def take[T](items: Iterable[T], n: int) -> list[T]:
return list(islice(items, n))
def drop[T](items: Iterable[T], n: int) -> Iterator[T]:
return islice(items, n, None)
# Works on infinite sources
take(count(), 5) # [0, 1, 2, 3, 4]
list(take(takewhile(lambda n: n < 10, count()), 4)) # [0, 1, 2, 3]
list(dropwhile(lambda n: n < 3, [1, 2, 3, 4, 1])) # [3, 4, 1]Save your own code snippets
Create a free account and build your private vault. Share publicly whenever you want.