support +/- modifiers on dice !roll

1 files changed, 27 insertions, 14 deletions

diff --git a/modules/dice.py b/modules/dice.py
index 6fb967a7..c808490e 100644
--- a/modules/dice.py
+++ b/modules/dice.py
@@ -1,9 +1,11 @@
 #--depends-on commands
 
-import random
+import random, re
 from src import ModuleManager, utils
 
 ERROR_FORMAT = "Incorrect format! Format must be [number]d[number], e.g. 1d20"
+RE_DICE = re.compile("([1-9]\d*)d([1-9]\d*)((?:[-+]\d+)*)", re.I)
+RE_MODIFIERS = re.compile("([-+]\d+)")
 
 class Module(ModuleManager.BaseModule):
     @utils.hook("received.command.roll", min_args=1)
@@ -12,20 +14,31 @@ class Module(ModuleManager.BaseModule):
         :help: Roll some dice, DND style
         :usage: [1-5]d[1-20]
         """
-        roll = event["args_split"][0].lower()
-        count, _, sides = roll.partition("d")
-        if not count.isdigit() or not sides.isdigit():
-            raise utils.EventError(ERROR_FORMAT)
+        match = RE_DICE.match(event["args_split"][0])
+        if match:
+            roll = match.group(0)
+            dice_count = int(match.group(1))
+            side_count = int(match.group(2))
+            modifiers = RE_MODIFIERS.findall(match.group(3))
 
-        count_n = min(5, int(count))
-        sides_n = min(20, int(sides))
+            if dice_count > 6:
+                raise utils.EventError("Max number of dice is 6")
+            if side_count > 20:
+                raise utils.EventError("Max number of sides is 20")
 
-        results = random.choices(range(1, sides_n+1), k=count_n)
+            results = random.choices(range(1, side_count+1), k=dice_count)
 
-        total_n = sum(results)
-        total = ""
-        if len(results) > 1:
-            total = " (total: %d)" % total_n
+            total_n = sum(results)
+            for modifier in modifiers:
+                if modifier[0] == "+":
+                    total_n += int(modifier[1:])
+                else:
+                    total_n -= int(modifier[1:])
 
-        event["stdout"].write("Rolled %s and got %s%s" % (
-            roll, ", ".join(str(r) for r in results), total))
+            total = ""
+            if len(results) > 1:
+                total = " (total: %d)" % total_n
+
+            results_str = ", ".join(str(r) for r in results)
+            event["stdout"].write("Rolled %s and got %s%s" % (
+                roll, results_str, total))