From 4343ec48243dc2838d4fff5815369679533ceaed Mon Sep 17 00:00:00 2001
From: Franz Diebold <franz.diebold@tum.de>
Date: Wed, 1 Jul 2020 18:03:30 +0200
Subject: [PATCH] Improve primes helper functions.

---
 src/common/primes.py       | 20 ++++++++++++++++----
 test/common/test_primes.py | 13 +++++++++++++
 2 files changed, 29 insertions(+), 4 deletions(-)

diff --git a/src/common/primes.py b/src/common/primes.py
index 2df03a7..05257c9 100644
--- a/src/common/primes.py
+++ b/src/common/primes.py
@@ -2,11 +2,11 @@
 Prime number utility functions.
 """
 
-from typing import Iterator, Dict, List, Set
+from typing import Iterable, Dict, List, Set, Optional
 import math
 
 
-def get_prime_factors(number: int) -> Iterator[int]:
+def get_prime_factors(number: int) -> Iterable[int]:
     """Get prime factors of a given number `number` as generator."""
     if number < 0:
         raise ValueError('Negative numbers are not supported.')
@@ -50,9 +50,21 @@ def is_prime(number: int) -> bool:
         return False
 
 
+def get_primes(threshold: Optional[int] = None) -> Iterable[int]:
+    """Return all prime numbers as an increasing iterable."""
+    yield 2
+    number = 3
+    while True:
+        if is_prime(number):
+            yield number
+        number += 2
+        if threshold and number >= threshold:
+            break
+
+
 def get_sorted_primes_list(threshold: int) -> List[int]:
     """Get all prime numbers up to a threshold `threshold` (exclusive) as a sorted list."""
-    return [number for number in range(2, threshold) if is_prime(number)]
+    return list(get_primes(threshold))
 
 
 # pylint: disable=invalid-name
@@ -63,4 +75,4 @@ def get_sorted_n_digit_primes(n: int) -> List[int]:
 
 def get_primes_set(threshold: int) -> Set[int]:
     """Get all prime numbers up to a threshold `threshold` (exclusive) as set."""
-    return set(get_sorted_primes_list(threshold))
+    return set(get_primes(threshold))
diff --git a/test/common/test_primes.py b/test/common/test_primes.py
index 8597a66..c7b0744 100644
--- a/test/common/test_primes.py
+++ b/test/common/test_primes.py
@@ -54,6 +54,19 @@ def test_is_prime(test_input, expected_result):
     assert actual_result == expected_result, f'Test failed for "{test_input}"'
 
 
+def test_get_primes():
+    # arrange
+    from src.common.primes import get_primes
+
+    # act
+    actual_result_iter = get_primes()
+
+    # assert
+    expected_result = [2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47, 53, 59, 61, 67, 71, 73, 79, 83, 89, 97]
+    for expected_value in expected_result:
+        assert next(actual_result_iter) == expected_value
+
+
 def test_get_sorted_primes_list():
     # arrange
     from src.common.primes import get_sorted_primes_list