Is_prime check in Python: can this be more pythonic?

Question

I have the following function that checks if a number is prime:

def is_prime(n):
    """ Is n prime? """

    if n < 2:
        return False
    if n in (2, 3):
        return True
    if n % 2 == 0 or n % 3 == 0:
        return False

    limit   = int(math.sqrt(n))
    divisor = 5

    while divisor <= limit:
        if n % divisor == 0 or n % (divisor + 2) == 0:
            return False
        divisor += 6

    return True

Most of it are checks and calculations to save up time, to make it faster than a simple loop through a range.

Can it be written in a more pythonic way? Can it be improved in some way?

Any other feedback is welcome too.

Answer 1

Your code looks nice, seems to follow PEP 8 and has docstrings. I don't have much to say about it but here are a few changes to make it look more pythonic.

The while loop can be rewritten using a for loop with range (or xrange).

for divisor in xrange(5, 1+int(math.sqrt(n)), 6):
    if n % divisor == 0 or n % (divisor + 2) == 0:
        return False
return True

Then, this looks like a scenario where we could use all/any. The obvious way to convert the code above is :

return not any(n % divisor == 0 or n % (divisor + 2) == 0 for divisor in xrange(5, int(math.sqrt(n))+1, 6))

but it looks better (and makes more sense) to use all with reversed conditions :

return all(n % divisor != 0 and n % (divisor + 2) != 0 for divisor in xrange(5, int(math.sqrt(n))+1, 6))

then, you can use the fact that non-zero numbers evaluates to True in a boolean context (and zero evaluates to False) :

return all(n % divisor and n % (divisor + 2) for divisor in xrange(5, int(math.sqrt(n))+1, 6))

Then, you could use the same kind of idea in the previous if then merge duplicated code :

return n % 2 and n % 3 and all(n % divisor and n % (divisor + 2) for divisor in xrange(5, int(math.sqrt(n))+1, 6))

This is pretty much all I have to say about the code. Now, I think it might be worth commenting the algorithm a bit more.

def is_prime(n):
    """ Check if n (assumed to be an int) is a prime number. """
    # Handle values up to 3
    if n < 2:
        return False
    if n in (2, 3):
        return True
    # Handles bigger values : divisors for n have to be
    # * multiples of 2
    # * multiples of 3
    # * of the form 5 + 6*i or 5 + 6*i + 2 (other cases already handled)  
    return n % 2 and n % 3 and all(n % divisor and n % (divisor + 2) for divisor in xrange(5, int(math.sqrt(n))+1, 6))

Answer 2

if n < 2:
    return False
if n in (2, 3):
    return True

These can be combined:

if n < 4:
    return n in (2, 3)

This:

divisor = 5

while divisor <= limit:
    if n % divisor == 0 or n % (divisor + 2) == 0:
        return False
    divisor += 6
return True

Really wants to be either a for loop:

for divisor in xrange(5, limit+1, 6):
    if n % divisor == 0 or n % (divisor + 2) == 0:
        return False
return True

or a generator:

return all(n % divisor and n % (divisor + 2) for divisor in xrange(5, limit+1, 6))

Answer 3

This can definitely be improved by using any one of of a multitude of prime number algorithms, e.g the Sieve of Atkin or the Sieve of Eratosthenes.

However, without having you implement one of those algorithms, your code looks to be pretty well optimized.

I have a couple suggestions that will help your code look and feel more Pythonic:

1. Take advantage of the fact that 0 == False. So instead of n % 2 == 0 we can do not n % 2. We can then logically factor out the not to make your if statements like this:

   if not(n % 2 and n % 3):
   ...
   if not(n % divisor and n % (divisor + 2)):
   

2. Don't tab align equal signs. This is actually a Pythonic convention pet peeve:

   # No
   limit   = int(math.sqrt(n))
   divisor = 5
   
   # Yes
   limit = int(math.sqrt(n))
   divisor = 5
   

3. Its good you use docstrings. They can be really helpful. However, I would suggest making yours more descriptive instead of asking a question.

4. I would probably use num instead of n as its slightly less ambiguous.

Answer 4

I would change this:

while divisor <= limit:
    if n % divisor == 0 or n % (divisor+2) == 0:
        return False
    divisor += 6

To this:

for pivot in range(6,limit+1,6):
    if n % (pivot-1) == 0 or n % (pivot+1) == 0:
        return False