Dice Rolling Utility

dice.py

import random
import math
import argparse
import re

def main():
    # Parse given arguments
    parser = argparse.ArgumentParser(description='Roll dice.')
    parser.add_argument(
        'rollArg',
        metavar='XdY',
        type=str,
        help='dice set to roll (format= XdY)'
    )
    parser.add_argument(
        '--sets',
        type=int,
        help='number of dice sets to roll'
    )
    parser.add_argument(
        '--drop',
        type=int,
        help='drops n values from given set of rolls (Defaults to dropping lowest values)'
    )
    parser.add_argument(
        '--high',
        type=bool,
        help='overrides --drop to drop the highest values given set of rolls'
    )
    args = parser.parse_args()

    # Check for proper formatting on primary argument
    if not re.match('^[0-9]+d[0-9]+$', args.rollArg):
        print 'Error reading arguments: dice set did not match "XdY" format'
        exit()

    if not args.sets or args.sets < 1:
        args.sets = 1

    # Break rollArg into its individual arguments
    quantity = int(args.rollArg.split('d')[0])
    sides = int(args.rollArg.split('d')[1])

    # Roll each set of dice
    sets = []
    for s in xrange(args.sets):
        sets.append(rollSetFn(quantity, sides, args.drop, args.high))

    # If dropping and drop val is less than total rolled
    if args.drop and args.drop < quantity:
        for s in sets:
            # Drop highest
            if args.high:
                for i in xrange(args.drop):
                    s.pop()
            # Drop lowest
            else:
                s.reverse()
                for i in xrange(args.drop):
                    s.pop()
                s.reverse()

    # Add up totals
    results = []
    for s in sets:
        total = 0
        for r in s:
            total += r
        
        results.append(total)
    
    if len(results) == 1:
        print results[0]
    else:
        print results
    
    return results

def rollDieFn(sides):
    return int(math.ceil(random.random() * sides))

def rollSetFn(quantity, sides, drop, high):
    # Roll set of dice
    rolls = []
    for i in xrange(quantity):
        rolls.append(rollDieFn(sides))

    rolls.sort()

    return rolls

main()