Generate and test prime numbers

prime.py

#! /usr/bin/env python
# -*- coding: utf-8 -*-
# Copyright 2014 Choongmin Lee
# Licensed under the MIT License
import argparse
from random import SystemRandom
import sys


parser = argparse.ArgumentParser(description='Generate or test prime numbers.')
subparsers = parser.add_subparsers()

gen_parser = subparsers.add_parser('gen', help='generate random primes')
gen_parser.set_defaults(func='generate_primes')
gen_parser.add_argument('bits', type=int,
                        help='specify the maximum number of bits')
gen_parser.add_argument('--hex', action='store_true',
                        help='print numbers in hexadecimal form')
gen_parser.add_argument('-n', '--number', type=int, default=1,
                        help='specify the number of primes to generate')
gen_parser.add_argument('-c', '--certainty', metavar='NUMBER', type=int,
                        default=20,
                        help='specify the certainty such that the probability '
                             'of each number is generated correctly would '
                             'exceed (1 - 1/2^certainty)')
gen_parser.add_argument('-s', '--safe', action='store_true',
                        help='generate safe primes')

test_parser = subparsers.add_parser('test', help='test primality of numbers')
test_parser.set_defaults(func='test_primes')
test_parser.add_argument('numbers', type=str, nargs='+',
                         help='numbers whose primality is to be tested; '
                              'read from stdin if -')
test_parser.add_argument('--hex', action='store_true',
                         help='read numbers in hexadecimal form')
test_parser.add_argument('-c', '--certainty', metavar='NUMBER', type=int,
                         default=20,
                         help='specify the certainty such that the '
                              'probability of each result being correct would '
                              'exceed (1 - 1/2^certainty)')

doctest_parser = subparsers.add_parser('doctest', help='run doctests')
doctest_parser.set_defaults(func='run_doctests')
doctest_parser.add_argument('-v', '--verbose', action='store_true',
                            help='print lots of stuff while testing')


def main(args=None):
    opts = parser.parse_args(args)
    sys.exit(globals()[opts.func](opts))


def generate_primes(opts):
    rng = SystemRandom()
    generator = random_prime if not opts.safe else random_safe_prime
    for _ in range(opts.number):
        p = generator(opts.bits, opts.certainty, rng)
        print ('{}' if not opts.hex else '{:x}').format(p)
        try:
            sys.stdout.flush()
        except IOError:
            pass
    return 0


def test_primes(opts):
    rng = SystemRandom()
    numbers = line_reader(sys.stdin) if opts.numbers == ['-'] else opts.numbers
    radix = 10 if not opts.hex else 16
    fmt = '{} is {}' if not opts.hex else '{} (hex) is {}'
    for n in numbers:
        n_is_prime = is_probable_prime(int(n, radix), opts.certainty, rng)
        print fmt.format(n, 'prime' if n_is_prime else 'composite')
        try:
            sys.stdout.flush()
        except IOError:
            pass
    return 0


def line_reader(f):
    while True:
        x = f.readline()
        if x == '':
            raise StopIteration()
        yield x.rstrip()


def run_doctests(opts):
    import doctest
    results = doctest.testmod(verbose=opts.verbose)
    print results
    return 0 if results.failed == 0 else 1


def random_prime(bits, certainty=20, rng=None):
    """Returns a random prime less than 2\ :sup:`bits + 1`. The probability
    that the returned number is a prime will exceed
    (1 - 1/2\ :sup:`certainty`).

    >>> p = random_prime(100, certainty=100)
    >>> p < 2 ** 100
    True
    >>> is_probable_prime(p, certainty=100)
    True

    """
    if rng is None:
        rng = SystemRandom()
    p = rng.randint(1, 1 << bits)
    while not is_probable_prime(p, certainty, rng):
        p = rng.randint(1, 1 << bits)
    return p


def random_safe_prime(bits, certainty=20, rng=None):
    """Returns a random prime number p of the form of 2q + 1 where q is also a
    prime. The probability of p being a safe prime will exceed
    (1 - 1/2\ :sup:`certainty`).

    >>> p = random_safe_prime(100, certainty=100)
    >>> p < 2 ** 100
    True
    >>> is_probable_prime(p, certainty=100)
    True
    >>> is_probable_prime(p >> 1, certainty=100)
    True

    """
    if rng is None:
        rng = SystemRandom()
    bits = bits - 1
    certainty = certainty + 1
    while True:
        q = random_prime(bits, certainty, rng)
        p = (q << 1) + 1
        if is_probable_prime(p, certainty, rng):
            return p


def is_probable_prime(n, certainty=20, rng=None):
    """Returns ``True`` if *n* is probably a prime, ``False`` if it's
    definitely composite. If certainty is ``<= 0``, ``True`` is returned.

    If the call returns ``True``, the probability that *n* is prime will exceed
    (1 - 1/2\ :sup:`certainty`). The execution time of this function is
    proportional to the value of certainty.

    *Implementation detail*: This function runs ⌊(*certainty* + 1)/2⌋ rounds of
    Miller-Rabin primality test.

    >>> is_probable_prime(11)
    True
    >>> is_probable_prime(2147483647)
    True
    >>> is_probable_prime(221, certainty=100)
    False

    """
    if n == 2 or n == 3:
        return True
    if n < 2 or not n & 1:
        return False
    m = n - 1
    s = 1
    d = m >> 1
    while not d & 1:
        s = s + 1
        d = d >> 1
    # assert n - 1 == (2 ** s) * d
    if rng is None:
        rng = SystemRandom()
    k = (certainty + 1) >> 1
    for i in range(k):
        a = rng.randint(2, n - 2)
        x = expmod(a, d, n)
        if x == 1 or x == m:
            continue
        for _ in range(s - 1):
            x = x * x % n
            if x == 1:
                return False
            if x == m:
                break
        else:
            return False
    return True


def expmod(b, e, m):
    """Computes b\ :sup:`e` mod m where *b*, *e*, and *m* are all positive
    integers.

    >>> expmod(2, 5, 10)
    2
    >>> expmod(4776, 1923, 201)
    81

    """
    r = 1
    while e:
        if e & 1:
            r = r * b % m
        e = e >> 1
        b = b * b % m
    return r


if __name__ == '__main__':
    main()