roll.py

#!/usr/bin/python

import random
import re
import sys

options = {}

class RollSpec:
    def __init__(self, spec):
        """Initializes a RollSpec given the number & type of dice, and an offset.

           `spec` may be either a string (e.g. "2d8+5") or an iterable
           (e.g. (2, 8, 5))."""
        if hasattr(spec, "isupper"):
            self.dicenum, self.dicetype, self.diceoffset = self._parse_string(spec)
        else:
            self.dicenum, self.dicetype, self.diceoffset = spec

    def _parse_string(self, spec_string):
        """Parses a string and returns the corresponding numbers.

           For example, `_parse_string("2d8+5")` would return `(2, 8, 5)`."""
        m = re.match("^([0-9]+)d([0-9]+)(\\+[0-9]+)?$", spec_string)
        if m is None:
            raise ValueError("{0} is not a valid dice spec".format(spec_string))
        return (int(m.group(1)), int(m.group(2)), (int(m.group(3) or 0)))

    def to_string(self):
        return "{0}d{1}+{2}".format(self.dicenum, self.dicetype, self.diceoffset)


def print_usage():
    print("""USAGE: roll.py [--verbose] [TRANSFORMS] [SPEC]

TRANSFORMS:
    Transforms apply different effects to the roll. They are processed in the
    order they're specified on the command line

    --advantage: Gives advantage on the roll.
    --disadvantage: Gives disadvantage on the roll.
    --drop-lowest: Drops the lowest value from a multi-dice roll.
    --reroll-ones: Reroll any ones that come up. Only do this once; if another 1
                   comes up, we keep it.
          
SPEC:
    A string in the format '<DICE NUMBER>d<DICE TYPE>+<OFFSET>', e.g. '2d8+5'""")

def debug(msg):
    global options
    if options["verbose"]:
        print("DEBUG: {0}".format(msg))

def roll(spec):
    """Rolls dice according to the given spec and returns the resulting list of ints."""
    rslt = []
    for i in range(spec.dicenum):
        indrslt = random.randint(1, spec.dicetype)
        debug("Rolled a {0}".format(indrslt))
        rslt.append(indrslt)
    debug("Rolled {0}; got {1}".format(spec.to_string(), rslt))
    return rslt

def evaluate(dice, spec):
    """Evaluates the given array of die values according to the given RollSpec."""
    return sum(dice) + spec.diceoffset

def t_reroll_ones(spec, scores):
    debug("Rerolling ones")
    rslt = []
    for sc in scores:
        if sc == 1:
            rslt.append(roll(RollSpec((1, spec.dicetype, 0)))[0])
            debug("Rerolled a one and got {0}".format(rslt[-1]))
        else:
            rslt.append(sc)
    debug("Die values after rerolling ones: {0}".format(rslt))
    return rslt

def t_drop_lowest(spec, scores):
    debug("Dropping the lowest")
    scores.sort()
    rslt = scores[1:]
    debug("Got {0} after dropping the lowest".format(rslt))
    return rslt

if __name__ == "__main__":
    options["verbose"] = False
    options["advantage"] = False
    options["disadvantage"] = False

    args = sys.argv[1:]
    transforms = []
    spec = RollSpec("1d20")
    for arg in args:
        if arg == "--reroll-ones":
            transforms.append(t_reroll_ones)
        elif arg == "--drop-lowest":
            transforms.append(t_drop_lowest)
        elif arg == "--advantage":
            options["advantage"] = True
        elif arg == "--disadvantage":
            options["disadvantage"] = True
        elif arg in ("-v", "--verbose"):
            options["verbose"] = True
        elif arg in ("-h", "--help"):
            print_usage()
            sys.exit(0)
        else:
            spec = RollSpec(arg)

    rslt = []
    if options["advantage"] or options["disadvantage"]:
        rslt_a = roll(spec)
        rslt_b = roll(spec)
        total_a = 0
        total_b = 0
        for t in transforms:
            rslt_a = t(rslt_a)
            rslt_b = t(rslt_b)
        total_a = evaluate(rslt_a, spec)
        total_b = evaluate(rslt_b, spec)
        debug("Rolled {0}=>{1} and {2}=>{3}".format(rslt_a, total_a, rslt_b, total_b))

        if options["advantage"]:
            if total_a < total_b:
                rslt = rslt_b
            else:
                rslt = rslt_a
            debug("Advantage: picked {0}".format(rslt))
        else:
            if total_a > total_b:
                rslt = rslt_b
            else:
                rslt = rslt_a
            debug("Disadvantage: picked {0}".format(rslt))
    else:
        rslt = roll(spec)
        for t in transforms:
            rslt = t(spec, rslt)
    print(evaluate(rslt, spec))