Improve return type hints for prime number generators

maths/prime_numbers.py

@@ -1,109 +1,73 @@
-import math
-from collections.abc import Generator
+"""
+Generate prime numbers up to a given limit using different approaches.
+"""
 
+from collections.abc import Iterator
+from math import isqrt
 
-def slow_primes(max_n: int) -> Generator[int]:
+
+def slow_primes(max_n: int) -> Iterator[int]:
     """
-    Return a list of all primes numbers up to max.
-    >>> list(slow_primes(0))
-    []
-    >>> list(slow_primes(-1))
+    Generate prime numbers up to max_n using a slow approach.
+
+    >>> list(slow_primes(10))
+    [2, 3, 5, 7]
+    >>> list(slow_primes(1))
     []
-    >>> list(slow_primes(-10))
+    >>> list(slow_primes(-5))
     []
-    >>> list(slow_primes(25))
-    [2, 3, 5, 7, 11, 13, 17, 19, 23]
-    >>> list(slow_primes(11))
-    [2, 3, 5, 7, 11]
-    >>> list(slow_primes(33))
-    [2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31]
-    >>> list(slow_primes(1000))[-1]
-    997
     """
-    numbers: Generator = (i for i in range(1, (max_n + 1)))
-    for i in (n for n in numbers if n > 1):
-        for j in range(2, i):
-            if (i % j) == 0:
+    for num in range(2, max_n + 1):
+        for i in range(2, num):
+            if num % i == 0:
                 break
         else:
-            yield i
+            yield num
 
 
-def primes(max_n: int) -> Generator[int]:
+def primes(max_n: int) -> Iterator[int]:
     """
-    Return a list of all primes numbers up to max.
-    >>> list(primes(0))
-    []
-    >>> list(primes(-1))
+    Generate prime numbers up to max_n using an optimized approach.
+
+    >>> list(primes(10))
+    [2, 3, 5, 7]
+    >>> list(primes(1))
     []
-    >>> list(primes(-10))
+    >>> list(primes(0))
     []
-    >>> list(primes(25))
-    [2, 3, 5, 7, 11, 13, 17, 19, 23]
-    >>> list(primes(11))
-    [2, 3, 5, 7, 11]
-    >>> list(primes(33))
-    [2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31]
-    >>> list(primes(1000))[-1]
-    997
     """
-    numbers: Generator = (i for i in range(1, (max_n + 1)))
-    for i in (n for n in numbers if n > 1):
-        # only need to check for factors up to sqrt(i)
-        bound = int(math.sqrt(i)) + 1
-        for j in range(2, bound):
-            if (i % j) == 0:
+    for num in range(2, max_n + 1):
+        for i in range(2, isqrt(num) + 1):
+            if num % i == 0:
                 break
         else:
-            yield i
+            yield num
 
 
-def fast_primes(max_n: int) -> Generator[int]:
+def fast_primes(max_n: int) -> Iterator[int]:
     """
-    Return a list of all primes numbers up to max.
-    >>> list(fast_primes(0))
-    []
-    >>> list(fast_primes(-1))
-    []
-    >>> list(fast_primes(-10))
+    Generate prime numbers up to max_n using the Sieve of Eratosthenes.
+
+    >>> list(fast_primes(10))
+    [2, 3, 5, 7]
+    >>> list(fast_primes(1))
     []
-    >>> list(fast_primes(25))
-    [2, 3, 5, 7, 11, 13, 17, 19, 23]
-    >>> list(fast_primes(11))
-    [2, 3, 5, 7, 11]
-    >>> list(fast_primes(33))
-    [2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31]
-    >>> list(fast_primes(1000))[-1]
-    997
     """
-    numbers: Generator = (i for i in range(1, (max_n + 1), 2))
-    # It's useless to test even numbers as they will not be prime
-    if max_n > 2:
-        yield 2  # Because 2 will not be tested, it's necessary to yield it now
-    for i in (n for n in numbers if n > 1):
-        bound = int(math.sqrt(i)) + 1
-        for j in range(3, bound, 2):
-            # As we removed the even numbers, we don't need them now
-            if (i % j) == 0:
-                break
-        else:
-            yield i
+    if max_n < 2:
+        return iter(())
 
+    sieve = [True] * (max_n + 1)
+    sieve[0] = sieve[1] = False
 
-def benchmark():
-    """
-    Let's benchmark our functions side-by-side...
-    """
-    from timeit import timeit
+    for num in range(2, isqrt(max_n) + 1):
+        if sieve[num]:
+            for multiple in range(num * num, max_n + 1, num):
+                sieve[multiple] = False
 
-    setup = "from __main__ import slow_primes, primes, fast_primes"
-    print(timeit("slow_primes(1_000_000_000_000)", setup=setup, number=1_000_000))
-    print(timeit("primes(1_000_000_000_000)", setup=setup, number=1_000_000))
-    print(timeit("fast_primes(1_000_000_000_000)", setup=setup, number=1_000_000))
+    return (num for num in range(2, max_n + 1) if sieve[num])
 
 
 if __name__ == "__main__":
-    number = int(input("Calculate primes up to:\n>> ").strip())
-    for ret in primes(number):
-        print(ret)
-    benchmark()
+    import doctest
+
+    doctest.testmod()