improved prime numbers implementation

Author: nikolasvargas

maths/prime_numbers.py

@@ -1,7 +1,7 @@
-from typing import List
+from typing import List, Generator
 
 
-def primes(max: int) -> List[int]:
+def primes(max: int) -> Generator[int]:
     """
     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)