import itertools


# The original from the video
def prime_numbers1():
    # Handle the first prime
    yield 2
    prime_cache = [2]  # Primes we've already found

    # Loop over positive, odd integers
    for n in itertools.count(3, 2):
        # Assume prime until proven otherwise
        is_prime = True

        # Check to see if any prime number divides n
        for p in prime_cache:
            if n % p == 0:  # p divides n evenly
                is_prime = False
                break

        if is_prime:
            prime_cache.append(n)
            yield n


# My modified, much more efficient version
def prime_numbers2():
    # Handle the first prime
    yield 2
    prime_cache = [2]  # Primes we've already found

    # Loop over positive, odd integers
    for n in itertools.count(3, 2):
        # Assume prime until proven otherwise
        is_prime = True

        # Check to see if any prime number divides n
        for p in prime_cache:
            quotient, remainder = divmod(n, p)
            if remainder == 0:  # p divides n evenly
                is_prime = False
                break
            if quotient <= p:  # p**2 >= n
                is_prime = True
                break

        if is_prime:
            prime_cache.append(n)
            yield n


def primes_up_to(n, primegen):
    lst = []
    for p in primegen:
        if p > n:
            return lst
        lst.append(p)


# In [1]: %timeit primegen.primes_up_to(200000, primegen.prime_numbers1())
# 7.03 s ± 15.5 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

# In [2]: %timeit primegen.primes_up_to(200000, primegen.prime_numbers2())
# 160 ms ± 1.07 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)