# Is_prime check in Python: can this be more pythonic?

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.

-
If you want to save time you may want to use a faster, probabilistic test, like Baillie-PSW or Miller-Rabin for larger inputs (usually you quickly check for divisibility by some small primes, then move on to them, but it depends how your function will be used). –  Thomas Jul 1 '14 at 3:18
@Thomas I have no idea of how to implement something like that. I know probabilistic tests exist, but I don't know enough to make use of them, I don't even know where to start. –  JCPedroza Jul 1 '14 at 11:13

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))

-
Super redundant comments. is_prime """Check if it's a prime""" What else would a method called "is_prime" do? Also # Handle values up to 3 followed by if n < 2 etc, not necessary when the code is that straight forward. Comments are good up to a point, but should be describing concepts or operations that aren't apparent. You have just documented self-documenting code. –  user11177 Jun 30 '14 at 17:20
Whoah, "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))" I always aim for readable code and nothing like this :X –  RvdK Jul 1 '14 at 10:24
This I think is a prime example of how you can obfuscate great, easy to read code in an attempt to write a one-liner (which also violates PEP8 by being 115 chars long) Not pythonic in my book but more.. perlithic. –  Voo Jul 1 '14 at 11:13
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))

-

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.

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.

-
Your first point is great for Golfing, but seems to hurt readability. Maybe it's just a Python thing that I'm not used to seeing. –  RubberDuck Jul 1 '14 at 1:08
If you do use not like that, you should leave a space before (, because it isn't a function call. –  otus Jul 2 '14 at 9:28

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

-
I see the appeal of initializing to 6 if you're going to be incrementing by 6. However, now divisor never acts as a divisor, so it probably should be renamed. –  200_success Jun 30 '14 at 15:23
@200_success: I certainly agree, thanks :) –  barak manos Jun 30 '14 at 15:28
I reckon limit+2 is what you should be using :-) (Also, I suggest you'd have a look at the other answers and they somehow apply to your code - this doesn't make your answer wrong in any way but it is just for your personal information). –  Josay Jun 30 '14 at 16:59