improved prime numbers implementation

data_structures/data_structures/heap/heap_generic.py

@@ -30,7 +30,8 @@ def _swap(self, i, j):
         """Performs changes required for swapping two elements in the heap"""
         # First update the indexes of the items in index map.
         self.pos_map[self.arr[i][0]], self.pos_map[self.arr[j][0]] = (
-            self.pos_map[self.arr[j][0]], self.pos_map[self.arr[i][0]]
+            self.pos_map[self.arr[j][0]],
+            self.pos_map[self.arr[i][0]],
         )
         # Then swap the items in the list.
         self.arr[i], self.arr[j] = self.arr[j], self.arr[i]

maths/prime_numbers.py

@@ -1,7 +1,7 @@
-from typing import List
+from typing import Generator, List
 
 
-def primes(max: int) -> List[int]:
+def primes(max: int) -> Generator[int, None, None]:
     """
     Return a list of all primes numbers up to max.
     >>> primes(10)
@@ -13,16 +13,16 @@ def primes(max: int) -> List[int]:
     >>> primes(1_000_000)[-1]
     999983
     """
-    max += 1
-    numbers = [False] * max
-    ret = []
-    for i in range(2, max):
-        if not numbers[i]:
-            for j in range(i, max, i):
-                numbers[j] = True
-            ret.append(i)
-    return ret
+    numbers: Generator = (i for i in range(1, (max + 1)))
+    for i in filter(lambda x: x > 1, numbers):
+        for j in range(2, i):
+            if (i % j) == 0:
+                break
+        else:
+            yield i
 
 
 if __name__ == "__main__":
-    print(primes(int(input("Calculate primes up to:\n>> "))))
+    number = int(input("Calculate primes up to:\n>> "))
+    for ret in primes(number):
+        print(ret)