prime_bench.py

prime_bench.py

from __future__ import annotations
import time
import math
from contextlib import contextmanager


@contextmanager
def timer():
    st = time.perf_counter()
    try:
        yield
    finally:
        t = time.perf_counter() - st
        print(f'Total time: {t:.3f}s')


def get_primes7(n: int) -> list[int]:
    if n < 2:
        # return [0 for _ in range(0)]  # for numba
        return []
    if n == 2:
        return [2]

    s = list(range(3, n + 1, 2))

    mroot = math.sqrt(n)
    half = len(s)
    i = 0
    m = 3
    while m <= mroot:
        if s[i]:
            j = (m**2 - 3) // 2
            s[j] = 0
            while j < half:
                s[j] = 0
                j += m
        i += 1
        m = 2 * i + 3

    # res = [x for x in s if x]
    res = list(filter(lambda x: x or False, s))
    return [2] + res


def main():
    lim = int(10e6)

    # warmup
    for i in range(4):
        get_primes7(lim)
        print('.', end='', flush=True)
    print()

    with timer():
        for i in range(10):
            res = get_primes7(lim)
            print(f'[{i}] Found {len(res)} prime numbers')


if __name__ == '__main__':
    main()