Skip to content

Instantly share code, notes, and snippets.

@realeroberto
Last active June 6, 2021 18:29
Show Gist options
  • Save realeroberto/41d787b2a80068d29cff to your computer and use it in GitHub Desktop.
Save realeroberto/41d787b2a80068d29cff to your computer and use it in GitHub Desktop.
A simple Python iterator that returns the sequence of prime numbers on successive calls to its next() method...
class PrimeGenerator:
def __init__(self):
self.primes = []
self.current_prime = 1
def __iter__(self):
return self
def __next__(self) -> int:
candidate = self.current_prime + 1
while True:
is_prime = True
for prime in self.primes:
if candidate % prime == 0:
is_prime = False
break
if is_prime:
self.primes.append(candidate)
self.current_prime = candidate
break
else:
candidate += 1
return self.current_prime
@avi1mizrahi
Copy link

avi1mizrahi commented Jun 6, 2021

My suggestion:

class PrimeGenerator:
    def __init__(self):
        self.primes = []
        self.current = 1

    def __iter__(self):
        return self

    def __next__(self) -> int:
        candidate = self.current + 1
        while True:
            for prime in itertools.takewhile(lambda p: candidate >= p**2, self.primes):
                if candidate % prime == 0:
                    break
            else:
                self.primes.append(candidate)
                self.current = candidate
                break

            candidate += 1

        return self.current

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment