4etools/dice.py

#!/usr/bin/python
#
# Command-line dice roller
#
# It assumes that the Mersenne Twister is acceptable.
# Support for real entropy is an RFE

import random
import sys
import re
import argparse


# Implements a "dice" - that is, a series of dice plus modifiers
class Dice():
    # This builds dice out of a simple input format
    # All of the dark regex magic is contained to this function,
    # lest it corrupt innocent code
    @classmethod
    def from_desc(cls, desc):
        if not re.match(r'^\d+d\d+', desc):
            raise Exception('Dice format invalid. See --help')

        (n, s) = re.match(r'^(\d+)d(\d+)', desc).groups()
        num = int(n)
        sides = int(s)
        mod = 0
        drop_low = 0
        drop_high = 0
        reroll = 0
        reroll_times = 0

        for m in re.findall(r'[dD+-]\d+', desc):
            if m[0] == '-':
                mod += int(m)
            elif m[0] == '+':
                mod += int(m[1:])
            elif m[0] == 'D':
                drop_high = int(m[1:])
            elif m[0] == 'd':
                drop_low = int(m[1:])

        m = re.search(r'r(\d+)(x\d+)?', desc)
        if m:
            (r, rt) = m.groups()
            reroll = int(r)
            if rt:
                reroll_times = int(rt[1:])

        return cls(num, sides, mod, drop_low, drop_high, reroll, reroll_times)


    def __init__(self, num, sides, mod, drop_low=0, drop_high=0, reroll=0, reroll_times=0):
        self.num = num
        self.sides = sides
        self.mod = mod
        self.drop_low = drop_low
        self.drop_high = drop_high
        self.reroll = reroll
        self.reroll_times = reroll_times
        self.times_rolled = 0


    def num_rolls(self):
        return self.times_rolled

        
    # Rolls the dice, returns a structure
    # containing the result plus metadata
    def roll(self, verbose=False):
        results = []
        rerolled = []

        # Roll the actual dice, handling rerolls
        for i in range(0, self.num):
            count = 0
            result = 0
            while result <= self.reroll or result == 0:
                if self.reroll_times and count >= self.reroll_times:
                    break

                result = random.randint(1, self.sides)
                count += 1
                if result <= self.reroll:
                    rerolled.append(result)

            results.append(result)
            

        # Drop any high or low, if specified
        if self.drop_low or self.drop_high:
            results.sort()

            dropped = []
            if self.drop_low:
                dropped.extend(results[:self.drop_low])
                results = results[self.drop_low:]

            if self.drop_high:
                dropped.extend(results[-1 * self.drop_high:])
                results = results[:-1 * self.drop_high]
            
            random.shuffle(results)
            
        # Get the results
        total = sum(results) + self.mod

        if verbose:
            drop_info = ''
            reroll_info = ''

            if dropped:
                drop_info = ' [dropped {}]'.format(','.join(['{}'.format(x) for x in dropped]))
            if rerolled:
                reroll_info = ' [rerolled {}]'.format(','.join(['{}'.format(x) for x in rerolled]))
                
            print('{dice}: {results}{drop}{reroll} {total}'.format(dice=self, results=results, total=total, drop=drop_info, reroll=reroll_info))

        else:
            print('{dice}: {total}'.format(dice=self, total=total))

        self.times_rolled += 1


    def __str__(self):
        reroll_info = ''
        drop_info = ''
        mod_info = ''

        if self.reroll:
            reroll_info = ' [reroll {}'.format(self.reroll)
            if self.reroll_times:
                reroll_info = reroll_info + 'x{}'.format(self.reroll_times)
            reroll_info = reroll_info + ']'

        if self.drop_low:
            drop_info = drop_info + ' [drop low {}]'.format(self.drop_low)
        if self.drop_high:
            drop_info = drop_info + ' [drop high {}]'.format(self.drop_high)

        if self.mod > 0:
            mod_info = '+{}'.format(self.mod)
        elif self.mod < 0:
            mod_info = '-{}'.format(self.mod)

        return "{num}d{sides}{mod}{drop}{reroll}".format(num=self.num, sides=self.sides, mod=mod_info, drop=drop_info, reroll=reroll_info)
        
        
# This takes command-line input as dice description strings
# and builds a list of Dice objects
def parse_input(args):
    dice_list = []
    for arg in args:
        dice_list.append(Dice.from_desc(arg))

    return dice_list


def parse_args():
    parser = argparse.ArgumentParser(description='Roll dice based on descriptions passed in on the command line', formatter_class=argparse.RawTextHelpFormatter)
    parser.add_argument('--repeat', '-r', metavar='N', type=int, default=1, help='repeat the given rolls N times')
    parser.add_argument('--verbose', '-v', action='store_true', help='print detailed information about each roll')
    parser.add_argument('dice', nargs=argparse.REMAINDER, help="""
Dice are input in the following form:

XdY[modifiers]

This will roll X Y-sided dice and apply the specified modifiers.

Modifiers can be any of the following (where N and M are integers):

(+|-)N   Add or subtract N from the total
dN       Drop the lowest-rolling N dice from the total
DN       Drop the highest-rolling N dice from the total
rN[xM]   Any dice that roll <= N will be rerolled.
         If the optional 'xM' option is specified, dice will be rerolled a maximum of M times.
         Otherwise each die will be rerolled until the result is > N

Examples:

1d20+5      roll 1 twenty-sided die, and add 5 to the result
6d6l1h1     roll 6 six-sided dice, and drop both the highest and the lowest roll
4d6l1r2x1   roll 4 six-sided dice. Any dice rolling a 1 or 2 will be rerolled once.
            If the result is still 1 or 2 it is kept. The lowest die is dropped from the result
""")
    return parser.parse_args()
    


def main():
    settings = parse_args()
    dice_list = parse_input(settings.dice)

    for i in range(settings.repeat):
        for dice in dice_list:
            dice.roll(verbose=settings.verbose)


if __name__ == '__main__':
    main()