prime_sieve.py

import unittest


def is_prime(n):
    if n in (2, 3):
        return True
    if n < 2 or n % 2 == 0:
        return False
    if n < 9:
        return True
    if n % 3 == 0:
        return False
    sqr = int(n**0.5)
    f = 5
    while f <= sqr:
        if n % f == 0:
            return False
        if n % (f+2) == 0:
            return False
        f += 6
    return True


class SieveTestCase(unittest.TestCase):
    def test_primes(self):
        get_primes = sieve_primes(121)
        self.assertTrue(all([is_prime(n) for n in get_primes]))


def sieve_primes(target):
    primes = {2}
    nums = set(range(3, target+1, 2))
    size = len(nums)

    while size:
        smallest_num = min(nums)
        for i in primes:
            if not smallest_num % i:
                multiples = set(range(smallest_num, target+1, smallest_num))
                nums = nums - multiples
            elif smallest_num < i**2:
                primes.add(smallest_num)
                multiples = set(range(smallest_num, target+1, smallest_num))
                nums = nums - multiples
                break
        size = len(nums)
    return primes

if __name__ == '__main__':
    unittest.main()