From 66aa6ffe07d234a1b8a1b5d2fcf39a29aaea89d1 Mon Sep 17 00:00:00 2001 From: ed cuss Date: Sun, 29 Mar 2026 10:49:12 +0100 Subject: [PATCH] perf: run fn once in loop --- src/danom/_stream.py | 17 +++++++++++------ 1 file changed, 11 insertions(+), 6 deletions(-) diff --git a/src/danom/_stream.py b/src/danom/_stream.py index e85bb21..f301dc4 100644 --- a/src/danom/_stream.py +++ b/src/danom/_stream.py @@ -7,7 +7,7 @@ from collections.abc import Awaitable, Callable, Iterable, Sequence from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor from copy import deepcopy -from enum import Enum, auto +from enum import Enum from functools import reduce from itertools import batched from typing import ParamSpec, TypeVar, cast @@ -312,10 +312,15 @@ def partition( seq_tuple = self.par_collect(workers=workers, use_threads=use_threads) else: seq_tuple = self.collect() - return ( - Stream(seq=tuple(x for x in seq_tuple if fn(x))), - Stream(seq=tuple(x for x in seq_tuple if not fn(x))), - ) + + pos, neg = [], [] + + for x in seq_tuple: + if fn(x): + pos.append(x) + else: + neg.append(x) + return (Stream.from_iterable(pos), Stream.from_iterable(neg)) def fold( self, initial: T, fn: Callable[[T, U], T], *, workers: int = 1, use_threads: bool = False @@ -442,7 +447,7 @@ async def async_collect(self) -> Awaitable[tuple[U, ...]]: class _Nothing(Enum): - NOTHING = auto() + NOTHING = 0 PlannedOps = tuple[str, StreamFn]