diceroller: first stab at a robust command-line dice rolling utility

diceroller.py

@@ -0,0 +1,156 @@
+#!/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
+
+# Implements a "dice" - that is, a series of dice plus modifiers
+# fixme: implement rerolls (currently they do nothing!)
+class Dice():
+    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
+        
+    # Rolls the dice, prints the result
+    def roll(self, verbose=True): # fixme: default verbose to False and add -v flag
+        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 = None
+            reroll_info = None
+
+            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 + ' up to {} times'.format(self.reroll_times)
+
+        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 builds dice out of a simple input format
+# All of the dark regex magic is contained to this function,
+# lest it corrupt innocent code
+def parse_input(args):
+    dice_list = []
+    for arg in args:
+        if not re.match(r'^\d+d\d+', arg):
+            continue
+
+        (n, s) = re.match(r'^(\d+)d(\d+)', arg).groups()
+        num = int(n)
+        sides = int(s)
+        mod = 0
+        dlow = 0
+        dhigh = 0
+        reroll = 0
+        reroll_times = 0
+
+        for m in re.findall(r'[hl+-]\d+', arg):
+            if m[0] == '-':
+                mod += int(m)
+            elif m[0] == '+':
+                mod += int(m[1:])
+            elif m[0] == 'h':
+                dhigh = int(m[1:])
+            elif m[0] == 'l':
+                dlow = int(m[1:])
+
+        m = re.search(r'r(\d+)(x\d+)?', arg)
+        if m:
+            (r, rt) = m.groups()
+            reroll = int(r)
+            if rt:
+                reroll_times = int(rt[1:])
+
+        dice_list.append(Dice(num, sides, mod, drop_low=dlow, drop_high=dhigh, reroll=reroll, reroll_times=reroll_times))
+
+    return dice_list
+
+
+def main():
+    # fixme: command line arguments! get some argparse over here
+
+    dice_list = parse_input(sys.argv)
+
+    for dice in dice_list:
+        dice.roll()
+
+
+if __name__ == '__main__':
+    main()