Primes in Python

primes.py

import sys


def sieve(limit):
    a = [True] * limit
    a[0] = a[1] = False

    for i, isprime in enumerate(a):
        if isprime:
            yield i
            for n in xrange(i*i, limit, i):
                a[n] = False


def single_line(limit):
    return ([2] + sorted(set(xrange(3,limit,2)) -
                         {x
                          for step in xrange(3, int(limit**0.5) + 1, 2) if step % 3 or step == 3
                          for x in xrange(step * 3, limit, step * 2)}))


if __name__ == '__main__':
    if sys.argv[1] == 'sieve':
        generator = sieve
    else:
        generator = single_line

    for p in generator(int(sys.argv[2])):
        pass