sadrand-py/sadrand.py

#!/usr/bin/python3

import random
import argparse
import math
import sys
import warnings

tqdm = lambda it, **kwargs: iter(it)

def std_rand(args):
    result = [0 for _ in range(args.resolution)];
    for _ in tqdm(range(args.iter), unit_scale=True, total=args.iter):
        result[random.randint(0, args.resolution-1)] += 1
    return result

def std_rand_fast(args):
    result = [0 for _ in range(args.resolution)];
    nbits = math.ceil(math.log2(args.resolution))
    for _ in tqdm(range(args.iter), unit_scale=True, total=args.iter):
        while True:
            rnd = random.getrandbits(nbits)
            if rnd < args.resolution:
                break
        result[rnd] += 1
    return result

def rnd_bytes(count, chunk_size=4096):
    to_read = count
    with open('/dev/urandom', 'br') as rfd:
        for _ in range((count // chunk_size) + 1):
            read_sz = chunk_size if to_read > chunk_size else to_read
            chunk = rfd.read(read_sz)
            to_read -= len(chunk)
            for res_byte in chunk:
                yield res_byte

def generic_rand(args, rnd_fun):
    result = [0 for _ in range(args.resolution)];
    for rbyte in tqdm(rnd_bytes(args.iter), total=args.iter, unit_scale=True):
        result[rnd_fun(rbyte)] += 1
    return result

mod_rand = lambda args: generic_rand(args,
                                 lambda rbyte: rbyte % args.resolution)
prop_rand = lambda args: generic_rand(args,
                                  lambda rbyte:(rbyte*(args.resolution)) // 256)

if __name__ == "__main__":
    parser = argparse.ArgumentParser();
    parser.add_argument('-i', '--iter', type=lambda v: int(v,0),
                        default=0x100000);
    parser.add_argument('-r', '--resolution', type=lambda v: int(v,0),
                        default=100);
    parser.add_argument('-c', '--chunk-size', type=int, default=4096);
    parser.add_argument('-R', '--relative', action="store_true",
                        default=False);
    parser.add_argument('-p', '--progress', action="store_true",
                        default=False);
    parser.add_argument('--std', action="store_true", default=False,
                        help="Use random.randint instead of random.getrandbits")

    args = parser.parse_args();

    if args.progress:
        try:
            from tqdm import tqdm
        except (NameError, ImportError):
            warnings.warn("tqdm module not found, no progress bar")

    methods = (std_rand if args.std else std_rand_fast, mod_rand, prop_rand)
    results = [method(args) for method in methods]

    if args.relative:
        for i, result in enumerate(results):
            results[i] = [(r / args.iter) * args.resolution for r in result]

    for i, res in enumerate(zip(*results)):
        if args.relative:
            i = i / args.resolution
        sys.stdout.write(str(i) + ' ' + (' '.join([str(r) for r in res]))+'\n')