Provide basic number theory functions isprime, isprimepower, nextprime, primes, prevprime

isprime.py

""" isprime.py
    May / June 2022 by M. F. Hasler

SYNOPSIS: very simple implementation of basic Number Theory functions:
    isprime(n): return True or False according to whether n is prime or not
    nextprime(n, i=1): return the i-th prime > n (i.e., >= n + 1).
    prevprime(n, i=1): return the i-th prime < n or 0 if n < prime(n).
    primes(n=math.inf, start=2, end=math.inf): return a generator 
              producing at most n primes in [start, end). 
    isprimepower(n): return True if n is prime, or
                     an integer k > 1 if n is a higher power of a prime,
                     or False if n is not a power of a prime.
REMARK: one might implement a function prime(n) to return the n-th prime,
  in which case one might store the primes as they are computed on demand
  and use the list of precomputed primes wherever useful.
"""
import math # for isqrt

def isprime(n: int):
    "Return True if n is a prime number, else False."
    if not n & 1:       # even: prime iff n = 2.
        return n == 2
    if n < 9:           # up to here, primes = odd numbers > 1.
        return n > 1
    return all(n%p for p in range(3, math.isqrt(n)+1, 2))

def isprimepower(n: int):
    "Return True or k > 1 if n is a prime or k-th power of a prime, else False."
    if isprime(n):
        return True
    for k in range(2,n):
        if not n >> k:      # n < 2^k
            return False
        # Only need to check prime exponents k: If  n = p^k  with  k = a*b,
        # a > b > 1, then we found  n = (p^a)^b  already for b < k.
        # Naive check of whether  n = r ^ k  with  r = round( n ^ 1/k ),
        # and r must also be a power of a prime.
        if isprime(k) and n == (r := round(n**(1/k)))**k and (p := isprimepower(r)):
            return k*p

def nextprime(n: int, i=1):
    "Return the i-th prime > n."
    # TO DO (?): use prevprime if i < 0.
    if n <= 2:
        return 2 if n < 2 else 3
    n += 2 if n & 1 else 1
    while not isprime(n) or (i := i - 1):
        n += 2
    return n    

def prevprime(n: int, i=1):
    "Return the i-th prime < n, or 0 if n <= prime(i)."
    if n & 1 == 0 and n > 2:
        n += 1
    while n > 3:
        if isprime(n := n - 2) and not(i := i - 1):
            return n
    return 0 if i > 1 or n < 3 else 2

def primes(n=math.inf, start=2, end=math.inf):
    "Return generator of at most n primes >= start and < end."
    start = nextprime(start-1) if start > 2 else 2
    while n and start < end:
        yield start
        start = nextprime(start)
        n -= 1

if __name__ == '__main__':
    print("The list of prime powers up to 30 should be:\n",
          L := list(filter(isprimepower,range(30))),
          "\nThe corresponding k-values are: (True <=> prime; 16 = (2²)²)\n",
          [isprimepower(n) for n in L])

in the while loop of nextprime(), one might use n+=4 if n%6==1 else 2.