Simplify/optimize partial_product

[pochmann]

I improved it further, it's now super simple, just dumping the initial and subsequent input items into a list and yielding tuple versions of it. PR in a moment.

The simplicity is my main objective, but it also turned out to be faster. Here are some test results with different numbers of iterables and iterable lengths. For example the first test uses three iterables of length 10:

lengths: [10] * 3
    4.3 ± 0.0 μs  improved
    6.0 ± 0.0 μs  current
    6.8 ± 0.0 μs  old_buggy

lengths: [5] * 5
    4.0 ± 0.0 μs  improved
    5.8 ± 0.0 μs  current
    6.9 ± 0.0 μs  old_buggy

lengths: [1000] * 2
  199.8 ± 0.8 μs  improved
  264.8 ± 1.3 μs  current
  267.3 ± 0.4 μs  old_buggy

lengths: [2] * 500
  546.1 ± 0.8 μs  improved
 1208.2 ± 5.7 μs  current
 1567.1 ± 7.9 μs  old_buggy

Benchmark code

Attempt This Online!

def partial_product_old_buggy(*args):

    iterables = [iter(it) for it in args]

    if len(iterables) == 0:
        return

    if len(iterables) == 1:
        yield from iterables[0]
        return

    previous, current, future = [], [], []

    try:
        current, *future = [next(i) for i in iterables]
    except StopIteration:
        raise ValueError("iterable must have at least one value.")

    yield (*previous, current, *future)

    for i in iterables:
        for current in i:
            yield (*previous, current, *future)
        previous.append(current)
        current = future[0] if len(future) > 0 else None
        future = future[1:]


def partial_product_current(*iterables):

    iterators = list(map(iter, iterables))

    try:
        future = [next(it) for it in iterators]
    except StopIteration:
        return
    yield (*future,)

    previous = []
    for it in iterators:
        current = future.pop(0)
        for current in it:
            yield (*previous, current, *future)
        previous.append(current)


def partial_product_improved(*iterables):

    iterators = list(map(iter, iterables))

    try:
        prod = [next(it) for it in iterators]
    except StopIteration:
        return
    yield tuple(prod)

    for i, it in enumerate(iterators):
        for prod[i] in it:
            yield tuple(prod)


funcs = partial_product_old_buggy, partial_product_current, partial_product_improved


from timeit import timeit
from statistics import mean, stdev
from itertools import product
from collections import deque
from time import time

def test(lengths, timeit_number):
    print('lengths:', lengths)
    lengths= eval(lengths)

    timeit_number //= 1
    times = {f: [] for f in funcs}
    def stats(f):
        ts = [t * 1e6 for t in sorted(times[f])[:10]]
        return f'{mean(ts):7.1f} ± {stdev(ts):3.1f} μs '

    iterables = [
        [object() for _ in range(m)]
        for m in lengths
    ]

    consume = deque(maxlen=0).extend
    for _ in range(100):
        for f in funcs:
        
            t = timeit(
                lambda: consume(f(*iterables)),
                'gc.enable()',
                number=timeit_number
            ) / timeit_number
            times[f].append(t)

    for f in sorted(funcs, key=stats):
        print(stats(f), f.__name__.removeprefix('partial_product_'))

    print()

test('[10] * 3', 250)
test('[5] * 5', 250)
test('[1000] * 2', 1)
test('[2] * 500', 1)