primes.py

View source code here on GitHub!

Includes

python.src.lib.primes.primes(stop: int | None = None) → Iterator[int]

Iterates over the prime numbers up to an (optional) limit, with caching.

This iterator leverages the modified_eratosthenes() iterator, but adds additional features. The biggest is a stop argument, which will prevent it from yielding any numbers larger than that. The next is caching of any yielded prime numbers.

python.src.lib.primes.modified_eratosthenes() → Iterator[int]

Iterates over prime numbers using the Sieve of Eratosthenes.

This function implements the Sieve of Eratosthenes. In general, it will iterate over all of the prime numbers. Below is a gif (courtesy of Wikipedia) that demonstrates the principle of the sieve.

Any animated example of the Sieve of Eratosthenes
python.src.lib.primes.prime_factors(num: int) → Iterator[int]

Iterates over the prime factors of a number.

python.src.lib.primes.is_prime(num: int, count: int = 1, distinct: bool = False) → bool

Detects if a number is prime or not.

If a count other than 1 is given, it returns True only if the number has exactly count prime factors.

python.src.lib.primes.primes_and_negatives(*args: int) → Iterator[int]

Iterate over the primes and their negatives using primes().

  1from itertools import count, filterfalse, takewhile
  2from math import ceil, sqrt
  3from pathlib import Path
  4from typing import Callable, Collection, Dict, Iterator, Optional
  5
  6from sortedcontainers import SortedSet
  7from umsgpack import load
  8
  9cache_filename = 'primes_cache.mpack'
 10
 11try:
 12    with Path(__file__).parent.joinpath(cache_filename).open('rb') as f:
 13        cache = SortedSet(load(f))
 14except Exception:  # pragma: no cover
 15    cache = SortedSet([
 16        2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47, 53, 59, 61
 17    ])
 18last_cached: int = cache[-1]
 19
 20
 21def primes(stop: Optional[int] = None) -> Iterator[int]:
 22    """Iterates over the prime numbers up to an (optional) limit, with caching.
 23
 24    This iterator leverages the :py:func:`modified_eratosthenes` iterator, but adds
 25    additional features. The biggest is a ``stop`` argument, which will prevent it
 26    from yielding any numbers larger than that. The next is caching of any yielded
 27    prime numbers."""
 28    if stop is None:
 29        yield from cache
 30    else:
 31        yield from takewhile(stop.__gt__, cache)
 32    global last_cached
 33    if stop and last_cached > stop:
 34        return
 35    if stop is None:
 36        secondary = modified_eratosthenes()
 37    else:
 38        secondary = takewhile(stop.__gt__, modified_eratosthenes())
 39    for p in filterfalse(cache.__contains__, secondary):
 40        cache.add(p)
 41        last_cached = p
 42        yield p
 43
 44
 45def modified_eratosthenes() -> Iterator[int]:
 46    """Iterates over prime numbers using the Sieve of Eratosthenes.
 47
 48    This function implements the `Sieve of Eratosthenes <https://en.wikipedia.org/wiki/Sieve_of_Eratosthenes>`_. In
 49    general, it will iterate over all of the prime numbers. Below is a gif (courtesy of Wikipedia) that demonstrates
 50    the principle of the sieve.
 51
 52    .. image:: https://upload.wikimedia.org/wikipedia/commons/9/94/Animation_Sieve_of_Eratosth.gif
 53        :alt: Any animated example of the Sieve of Eratosthenes"""
 54    # modified_eratosthenes taken, refactored from https://stackoverflow.com/a/19391111
 55    yield from (2, 3, 5, 7)
 56    sieve: Dict[int, int] = {}
 57    recurse = primes()
 58    next(recurse)
 59    prime = next(recurse)
 60    if prime != 3:
 61        raise ValueError()  # pragma: no cover
 62    prime_squared = prime * prime
 63    for candidate in count(9, 2):
 64        if candidate in sieve:  # if c is a multiple of some base prime, or
 65            step = sieve.pop(candidate)
 66        elif candidate < prime_squared:  # prime, or
 67            yield candidate
 68            continue
 69        else:  # the next base prime's square
 70            # if candidate != prime_squared:
 71            #     raise ValueError()
 72            step = prime * 2
 73            prime = next(recurse)
 74            prime_squared = prime * prime
 75        candidate += step  # the next multiple
 76        while candidate in sieve:  # make sure to not overwrite another sieve entry
 77            candidate += step
 78        sieve[candidate] = step
 79
 80
 81def prime_factors(num: int) -> Iterator[int]:
 82    """Iterates over the prime factors of a number."""
 83    if num < 0:
 84        yield -1
 85        num = -num
 86    if num == 0:
 87        yield 0
 88    else:
 89        root = ceil(sqrt(num))
 90        for factor in primes():
 91            modulo = num % factor
 92            if modulo == 0:
 93                while modulo == 0:  # double-check to call sqrt once
 94                    yield factor
 95                    num //= factor
 96                    modulo = num % factor
 97                root = ceil(sqrt(num))
 98            if num <= 1:
 99                break
100            if factor > root:
101                yield num
102                break
103
104
105def is_prime(
106    num: int,
107    count: int = 1,
108    distinct: bool = False
109) -> bool:
110    """Detects if a number is prime or not.
111
112    If a count other than 1 is given, it returns True only if the number has
113    exactly count prime factors."""
114    if num in (0, 1):
115        return False
116    factors = iter(prime_factors(num))
117    if count == 1:
118        if num in cache:  # always has 2
119            return True
120        if num % 2 == 0:
121            return False
122        return next(factors) == num
123    seen: Collection[Optional[int]]
124    seen_add: Callable[[Optional[int]], None]
125    if distinct:
126        seen = set()
127        seen_add = seen.add
128    else:
129        seen = []
130        seen_add = seen.append
131    while None not in seen and count != len(seen):
132        seen_add(next(factors, None))
133    return None not in seen and next(factors, None) is None
134
135
136def primes_and_negatives(*args: int) -> Iterator[int]:
137    """Iterate over the primes and their negatives using :py:func:`primes`."""
138    for p in primes(*args):
139        yield p
140        yield -p

Tags: python-iterator, prime-number