Dealing with primes

primes.py

from bisect import bisect_right, bisect_left
from itertools import count, islice
from math import isqrt
from typing import Iterator, Union, Generator


class Primes:
    known = [2, 3]

    def __iter__(self) -> Generator[int, None, None]:
        yield from self.known
        # start generating primes (saving while we do it)
        number = self.known[-1] + 2
        cutoff = bisect_right(self.known, isqrt(number))
        # we only want to look until there's a square root
        for number in count(number, 2):
            # has the square root grown enough to cover a new prime?
            if self.known[cutoff] ** 2 <= number:
                cutoff += 1
            # it's a prime if we have remainders for all relevant primes
            if all(number % prime for prime in islice(self.known, 1, cutoff)):
                self.known.append(number)
                yield number

    def factorize(self, value: int) -> Generator[int, None, None]:
        number = int(value)
        if value != number or number <= 0:
            raise ValueError("Expected positive integer.")
        for prime in self:
            while number % prime == 0:
                yield prime
                number //= prime
            if number == 1:
                return

    def __contains__(self, value: int) -> bool:
        number = int(value)
        if value != number or number <= 0:
            raise ValueError("Expected positive integer.")
        if value == 1:
            return True  # special case :)
        # is it small enough for us to know about it directly? O(log n)
        index = bisect_left(self.known, number)
        if index < len(self.known):
            return self.known[index] == number
        # uh-oh, new number! O(n**2) worst case, O(n) usually
        square_root = isqrt(number)
        for prime in self:
            if prime > square_root:
                return True  # relevant primes exhausted
            if number % prime == 0:
                return False  # found a factor

    def __getitem__(self, item: Union[int, slice]) -> Union[int, Iterator[int]]:
        if isinstance(item, int):
            if item < 0:
                raise ValueError("Can't walk down from infinity! :)")
            return next(islice(self, item, item + 1))
        elif isinstance(item, slice):
            return islice(self, item.start, item.stop, item.step)
        else:
            raise TypeError(f"Expected int or slice, got {type(item)}")


primes = Primes()