judge n is prime or not

prime.py

#!/usr/bin/env python
# coding: utf-8
import math

def prime(n):
    if n <= 2:
        return True
    if n % 2 == 0:
        return False
    for i in range(3, int(math.sqrt(n)) + 1, 2):
        if n % i == 0:
            return False
    return True

if __name__ == '__main__':
    while True:
        n = int(raw_input('enter n:\n'))
        if n == 0:
            break
        if prime(n):
            print 'prime\n'
        else:
            print 'not prime\n'