Python while loop for finding prime numbers

Question

As a first exercise with Python, I'm trying to write a program using loops to find primes. Everything works with a for loop so I am trying to use a while loop. This works but the program returns a few incorrect numbers.

import math
# looking for all primes below this number
max_num = int(input("max number?: "))

primes = [2]  # start with 2
test_num = 3  # which means testing starts with 3

while test_num < max_num:
    i = 0
    # It's only necessary to check with the primes smaller than the square
    # root of the test_num
    while primes[i] < math.sqrt(test_num):
        # using modulo to figure out if test_num is prime or not
        if (test_num % primes[i]) == 0:
            test_num += 1
            break
        else:
            i += 1
    else:
        primes.append(test_num)
        test_num += 1

print(primes)

So the weird thing is that for max_num=100 it returns:

[2, 3, 5, 7, 9, 11, 13, 17, 19, 23, 25, 29, 31, 37, 41, 43, 47, 49, 53, 59, 61, 67, 71, 73, 79, 83, 89, 97]

which is correct except for 9, 25 and 49 and I can't figure out why.

Solution 1

You need to go up to and including the square root. Otherwise your algorithm will miss the family of prime squares (9, 25, and 49 are prime squares).

The quick fix is to replace < with <= as your stopping condition.

But consider changing the stopping condition to

primes[i] * primes[i] <= test_num

With this test, you don't dip in and out of floating point.

Solution 2

If you would like to find the next primes for each iteration, probably this is a better function since it bypasses the procedure of giving input.

import math
def primes():
    (h,n) = (2,1)
    k = 2
    while True:
        if any(k % i == 0 for i in range(2,k))==False: # h is composite number, so continue iteration
            (h,n) = (k,n+1) # After the for loop we can find out the next prime number
            k += 1
        else:
            k += 1 # Remember to continue the loop, if not then next function may return the same number
            continue        
        yield h
x = prime()

Then you may use the followings to iterate:

next(x)

or

[next(x) for _ in range(20)]

which gives the output

[2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47, 53, 59, 61, 67, 71]

Hope this is a graceful function for writing while loop for prime numbers aiming at beginners.