I have a homework assignment to program a prime number generator

Question

The rules for the assignment regarding asking for help are as follows: I am allowed to ask for help if I am confused with my code, and cannot copy and paste code from others. I must understand all code, and what each function and piece of the code is doing. We learned about the next function, and we have to use a class to make the generator. We also learned about StopIteration, but when would I use that? My code doesn't seem to work properly when I am using it. Maybe I am putting it in the wrong place.

I was having a bit of trouble with ending the generator, and it also kept returning None when the number wasn't prime, when I wanted it to keep going instead to the next prime number. How could I make it so it ends, and also keeps going until it reaches the next prime number instead of returning None? This is what I tried so far:

    class PrimeGenerator:
        def __init__(self, maximum):
            self.maximum = maximum + 1
            self.count = 0
    
        def __next__(self):
            if self.count < self.maximum:
                for num in range(self.count, self.maximum):
                    for divisor in range(2, self.count + 1):
                        if num // divisor != num / divisor:
                            print(f'number is prime, it is {num} and the count is {self.count}, and the divisor is {divisor}')
                            self.count += 1
                            return num
                        else:
                            break
    
                    self.count += 1
    
    my_gen = PrimeGenerator(7)
    print(next(my_gen))
    print(next(my_gen))
    print(next(my_gen))
    print(next(my_gen))
    print(next(my_gen))

Any help would be appreciated.

Answer 1

In the constructor of the class, define the current number and the highest number. The point of having a highest number is that when we encounter a number which is greater than the highest number, we end our prime number generation. To do that, we raise a StopIteration exception.

In the __next__ method, we check if the number is prime. If it is, we return it else we move forward to the next one and repeat the process.

def is_prime(n):
    if n == 2:
        return True
    for i in range(2, int(n ** 0.5) + 1):
        if n % i == 0:
            return False
    return True


class PrimeGenerator:
    def __init__(self, highest=100):
        self.current = 2
        self.highest_number = highest

    def __iter__(self):
        return self

    def __next__(self):
        while True:
            if self.current > self.highest_number:
                raise StopIteration
            if is_prime(self.current):
                self.current += 1
                return self.current - 1
            self.current += 1


prime_generator = PrimeGenerator(10)
print(next(prime_generator))
print(next(prime_generator))
print(next(prime_generator))
print(next(prime_generator))
print(next(prime_generator))

# For large maximum values
for prime in PrimeGenerator(100):
    print(prime)

Answer 2

Your question, and comments that follow, leave it unclear what maximum means. I.e. the number of primes the prime generator should produce or the largest number it should consider. First let's consider it limits the number of primes produced:

class PrimeGenerator:
    def __init__(self, maximum):
        self.maximum = maximum  # how many primes we'll generate
        self.count = 0  # how many primes we've generated
        self.number = 1  # prime the prime pump

    def __iter__(self):
        return self

    def __next__(self):
        if self.count < self.maximum:
            if self.count == 0:  # for speed, treat 2 as a special case
                self.count += 1
                return 2

            while True:
                self.number += 2

                for divisor in range(3, int(self.number ** 0.5) + 1, 2):
                    if self.number % divisor == 0:
                        break
                else:  # no break 'for'
                    self.count += 1
                    return self.number  # found a prime

        raise StopIteration()  # exceeded maximum

if __name__ == '__main__':

    generator = PrimeGenerator(25)

    while True:
        try:
            print(next(generator), end=" ")

        except StopIteration:
            print()
            break

Now let's modify the code so that the maximum argument represents the largest number considered:

class PrimeGenerator:
    def __init__(self, maximum):
        self.maximum = maximum  # largest number we'll consider
        self.number = 2  # prime the prime pump

    def __iter__(self):
        return self

    def __next__(self):
        if self.number <= self.maximum:
            if self.number == 2:  # for speed, treat 2 as a special case
                self.number += 1
                return 2

            while self.number < self.maximum:

                for divisor in range(3, int(self.number ** 0.5) + 1, 2):
                    if self.number % divisor == 0:
                        break
                else:  # no break 'for'
                    prime = self.number
                    self.number += 2
                    return prime  # found a prime

                self.number += 2

        raise StopIteration()  # exceeded maximum

These examples treat 2 as a special case, subsequently avoiding even numbers as candidates and/or divisors as in the other examples. (The benefit of doing this is a ~4x speed up.)

Answer 3

What you've implemented is an "iterator", but not a true "generator". Per the Python docs:

Python provides generator functions as a convenient shortcut to building iterators. Generator functions allow you to declare a function that behaves like an iterator, i.e. it can be used in a for loop.

def PrimeGenerator(maximum):
    for num in range(2, maximum + 1):
        for divisor in range(2, num):
            if num // divisor == num / divisor:
                break
        else:
            yield num


for p in PrimeGenerator(7):
    print(p)

Result:

2
3
5
7

For completeness, and to show how much of a code savings you get with true generators, here's the same logic as a classic iterator:

class PrimeGenerator:

    def __init__(self, highest):
        self.next_candidate = 2
        self.highest_number = highest

    def __iter__(self):
        return self

    def __next__(self):
        for num in range(self.next_candidate, self.highest_number + 1):
            for divisor in range(2, num // 2 + 1):
                if num // divisor == num / divisor:
                    break
            else:
                self.next_candidate = num + 1
                return num
        raise StopIteration

    def next(self):
        return self.__next__()

for p in PrimeGenerator(7):
    print(p)

to be sure it worked, I also tried using the pattern of calling next() on the iterator. This gives the same result:

p = PrimeGenerator(7)
while True:
    try:
        print(p.next())
    except StopIteration:
        break

Answer 4

In order to stop the generator, all you have to do is raise StopIteration at the end of the function.

I believe that the prime finding logic for your code is wrong, so I have written my own prime test. Each time __next__ is called, it continues to test integers from count and yields if they are prime. We only have to loop up to sqrt(num) to check if it's prime, because any factors above sqrt(num) will have a corresponding factor that is smaller. You can also use the modulo operator to easily check for factors, rather than dividing.

I also added __iter__ so that the generator can be used in a for loop.

class PrimeGenerator:
    def __init__(self, maximum):
        self.maximum = maximum
        self.count = 2

    def __iter__(self):
        return self

    def __next__(self):
        while self.count <= self.maximum:
            num = self.count
            is_prime = True
            for divisor in range(2, 1 + int(num ** 0.5)):
                if self.count % divisor == 0:
                    is_prime = False
                    break
            if is_prime:
                self.count += 1
                return num
            self.count += 1
        raise StopIteration()