primes.py

#!/usr/bin/env python
import time
from math import sqrt

def sieve(stop):
    '''
    That sieve of eratosthenes thing. Implemented directly from the
    pseudocode found here:
    http://en.wikipedia.org/wiki/Sieve_of_Eratosthenes#Implementation
    '''
    stop += 1
    counter = 0
    num_list = range(2, stop)
    for i in num_list:
        counter += 1
        if i:
            args = (i * 2, stop, i)
            for j in range(*args):
                counter += 1
                idx = j - 2
                num_list[idx] = False
    primes = [n for n in num_list if n]
    return primes, counter

def trial_div(stop, lim_sqrt=False, counter=0):
    '''
    trial divide each candidate by each known prime
    OR if lim_sqrt=True is passed as an argument, 
    trial divide each candidate by  each known prime < sqrt(candidate)
    '''
    counter = 0
    start, step = 3, 2
    args = start, stop, step
    primes = [2]
    for n in range(*args):
        counter += 1
        prime = True
        for p in primes:
            # assume prime is True
            counter += 1
            # limit trial division to known primes < sqrt(n)
            if lim_sqrt:
                if not p <= sqrt(n):
                    break
                if n % p == 0:
                    prime = False
            # trial divide each candidate by all known primes
            else:
                if n % p == 0:
                    prime = False
        # if the prime flag wasn't changed
        if prime:
            primes.append(n)
    return primes, counter


def format_output(name, stop, primes, steps, elapsed):
    '''
    Return a formatted string with result data filled in
    '''
    return '\n'.join(
        ('------------------------------',
        "%s" % name,
        "",
        "primes found: %s" % len(primes),
        "steps taken: %s" % steps,
        "running time: %ss" % elapsed,
        "",
        )
    )

def prime_wrapper(stop, f, arg=None):
    '''
    A wrapper function for calling each prime function and
    timing how long it takes. It calls trial_div(max_prime, lim_sqrt=True)
    if arg is True
    '''
    then = time.time()
    if arg:
        primes, steps = f(stop, arg)
    else:
        primes, steps = f(stop)
    now = time.time()
    elapsed = str(now - then)[:8]
    return primes, steps, elapsed

if __name__ == '__main__':
    import sys
    # CSV = True for csv-ish output in the following format:
    # 'name','max prime candidate','time'
    CSV = False
    # ugly ass separator
    sep = '##################################################'

    # results are printed out in the order of these tuples
    # each element is (func, arg, name)
    prime_funcs = [(sieve, None, 'sieve')]
                   (trial_div, True, 'trial division - sqrt'),
                   (trial_div, None, 'trial division - all')]

    # find 2 to n primes for each max candidate in this list
    max_prime_list = [1000, 5000, 10000, 25000, 50000]

    # or pass an argument in to check all 3 methods vs one range of numbers
    if len(sys.argv) > 1 and sys.argv[1].isdigit():
        max_prime_list = [int(sys.argv[1])]

    if CSV: 
        stat_data = []

    for max_prime in max_prime_list:
        # output header
        print ''.join([sep, '\n', 'Primes from 2 to %s' % max_prime, '\n'])
        # iterate over the prime_funcs list from above and call
        # the wrapper function with once with each prime function
        # then output the results.
        for prime_func, arg, name in prime_funcs:
            stats = prime_wrapper(max_prime, prime_func, arg)
            if CSV:
                stat_data.append((name, str(max_prime), str(stats[2])))
            print format_output(name, max_prime, *stats)

    if CSV:
        for data in stat_data:
            print ','.join(data)