From 5a8ecf402bc79976f74c39b7ff4b558b65fe099d Mon Sep 17 00:00:00 2001
From: Ari Koivula <ari@koivu.la>
Date: Fri, 9 Aug 2013 00:35:22 +0300
Subject: [PATCH] [dice] Add back arithmetic expression functionality.

-Also allows for using many dice expressions in one command, just like before.
---
 willie/modules/dice.py | 103 +++++++++++++++++++++++++++--------------
 1 file changed, 68 insertions(+), 35 deletions(-)

diff --git a/willie/modules/dice.py b/willie/modules/dice.py
index cfd428d099..650721008d 100644
--- a/willie/modules/dice.py
+++ b/willie/modules/dice.py
@@ -6,11 +6,12 @@
 
 http://willie.dftba.net/
 """
-
 import random
-import willie.module
 import re
 
+import willie.module
+from willie.tools import eval_equation
+
 
 class DicePouch:
     def __init__(self, num_of_die, type_of_die, addition):
@@ -108,60 +109,92 @@ def get_number_of_faces(self):
         return len(self.dice) + len(self.dropped)
 
 
-@willie.module.commands("roll")
-@willie.module.commands("dice")
-@willie.module.commands("d")
-@willie.module.priority("medium")
-@willie.module.example(".roll 3d1+1", 'You roll 3d1+1: (1+1+1)+1 = 4')
-@willie.module.example(".roll 3d1v2+1", 'You roll 3d1v2+1: (1[+1+1])+1 = 2')
-@willie.module.example(".roll 2d4", re='You roll 2d4: \(\d\+\d\) = \d')
-@willie.module.example(".roll 100d1", re='[^:]*: \(100x1\) = 100')
-def roll(bot, trigger):
-    """.dice XdY[vZ][+N], rolls dice and reports the result.
-
-    X is the number of dice. Y is the number of faces in the dice. Z is the
-    number of lowest dice to be dropped from the result. N is the constant to
-    be applied to the end result.
-    """
+def _roll_dice(dice_expression):
     result = re.search(r"""
-            (?P<dice_num>\d+)
+            (?P<dice_num>\d*)
             d
             (?P<dice_type>\d+)
             (v(?P<drop_lowest>\d+))?
-            (?P<plus>(-|\+)\d+)?
             $""",
-            trigger.group(2),
+            dice_expression,
             re.IGNORECASE | re.VERBOSE)
-    if not result:
-        bot.reply("Syntax for rolling dice is XdY[vZ][+N].")
-        return
 
-    dice_num = int(result.group('dice_num'))
+    dice_num = int(result.group('dice_num') or 1)
     dice_type = int(result.group('dice_type'))
-    addition = int(result.group('plus') or 0)
 
     # Upper limit for dice should be at most a million. Creating a dict with
     # more than a million elements already takes a noticeable amount of time
     # on a fast computer and ~55kB of memory.
     if dice_num > 1000:
-        bot.reply("I only have 1000 dice. =(")
-        return
+        return None
 
-    dice = DicePouch(dice_num, dice_type, addition)
+    dice = DicePouch(dice_num, dice_type, 0)
 
     if result.group('drop_lowest'):
         drop = int(result.group('drop_lowest'))
         dice.drop_lowest(drop)
 
-    if dice_num <= 10:
-        dice_str = dice.get_simple_string()
-    elif dice.get_number_of_faces() <= 10:
-        dice_str = dice.get_compressed_string()
-    else:
-        dice_str = "(...)"
+    return dice
+
+
+@willie.module.commands("roll")
+@willie.module.commands("dice")
+@willie.module.commands("d")
+@willie.module.priority("medium")
+@willie.module.example(".roll 3d1+1", 'You roll 3d1+1: (1+1+1)+1 = 4')
+@willie.module.example(".roll 3d1v2+1", 'You roll 3d1v2+1: (1[+1+1])+1 = 2')
+@willie.module.example(".roll 2d4", re='You roll 2d4: \(\d\+\d\) = \d')
+@willie.module.example(".roll 100d1", re='[^:]*: \(100x1\) = 100')
+@willie.module.example(".roll 1001d1", 'I only have 1000 dice. =(')
+@willie.module.example(".roll 1d1 + 1d1", 'You roll 1d1 + 1d1: (1) + (1) = 2')
+@willie.module.example(".roll 1d1+1d1", 'You roll 1d1+1d1: (1)+(1) = 2')
+def roll(bot, trigger):
+    """.dice XdY[vZ][+N], rolls dice and reports the result.
+
+    X is the number of dice. Y is the number of faces in the dice. Z is the
+    number of lowest dice to be dropped from the result. N is the constant to
+    be applied to the end result.
+    """
+    # This regexp is only allowed to have one captured group, because having
+    # more would alter the output of re.findall.
+    dice_regexp = r"\d*d\d+(?:v\d+)?"
+
+    # Get a list of all dice expressions, evaluate them and then replace the
+    # expressions in the original string with the results. Replacing is done
+    # using string formatting, so %-characters must be escaped.
+    arg_str = trigger.group(2)
+    dice_expressions = re.findall(dice_regexp, arg_str)
+    arg_str = arg_str.replace("%", "%%")
+    arg_str = re.sub(dice_regexp, "%s", arg_str)
+    dice = map(_roll_dice, dice_expressions)
+    if None in dice:
+        bot.reply("I only have 1000 dice. =(")
+        return
+
+    def _get_eval_str(dice):
+        return "(%d)" % (dice.get_sum(),)
+
+    def _get_pretty_str(dice):
+        if dice.num <= 10:
+            return dice.get_simple_string()
+        elif dice.get_number_of_faces() <= 10:
+            return dice.get_compressed_string()
+        else:
+            return "(...)"
+
+    eval_str = arg_str % (tuple(map(_get_eval_str, dice)))
+    pretty_str = arg_str % (tuple(map(_get_pretty_str, dice)))
+
+    # Showing the actual error will hopefully give a better hint of what is
+    # wrong with the syntax than a generic error message.
+    try:
+        result = eval_equation(eval_str)
+    except Exception as e:
+        bot.reply("SyntaxError, eval(%s), %s" % (eval_str, e))
+        return
 
     bot.reply("You roll %s: %s = %d" % (
-            trigger.group(2), dice_str, dice.get_sum()))
+            trigger.group(2), pretty_str, result))
 
 
 @willie.module.commands("choice")