From 362ae1444853371e4c4e71fa2842a89a8b64936a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Nikola=20Forr=C3=B3?= <nikola@mraveniste.cc>
Date: Wed, 13 Mar 2019 18:30:40 +0100
Subject: [PATCH] Implement !roll command

---
 clients/discord.py |  10 +
 commands.py        |   9 +
 requirements.txt   |   1 +
 services/roll20.py | 516 +++++++++++++++++++++++++++++++++++++++++++++
 4 files changed, 536 insertions(+)
 create mode 100644 services/roll20.py

diff --git a/clients/discord.py b/clients/discord.py
index e8ba840..21d6d9e 100644
--- a/clients/discord.py
+++ b/clients/discord.py
@@ -63,6 +63,7 @@ class DiscordClient(discord.Client):
             (re.compile(r'^!yt\s+(?P<q>")?(?P<query>.+)(?(q)")$'), self._do_yt),
             (re.compile(r'^!(find)?clip\s+(?P<q>")?(?P<filter>.+)(?(q)")$'), self._do_clip),
             (re.compile(r'^!cheese\s+(?P<q>")?(?P<query>.+)(?(q)")$'), self._do_cheese),
+            (re.compile(r'^!roll\s+(?P<q>")?(?P<formula>.+)(?(q)")$'), self._do_roll),
             (re.compile(r'^!nextstream$'), self._do_nextstream),
         ]
         super(DiscordClient, self).__init__()
@@ -247,6 +248,15 @@ class DiscordClient(discord.Client):
                 embed.add_field(name=name.strip(), value=value.strip(), inline=True)
             await self.send_message(message.channel, embed=embed)
 
+    async def _do_roll(self, server, user, message, formula, **kwargs):
+        try:
+            result = self.commands.roll(formula)
+        except CommandError as e:
+            await self.send_message(message.channel, 'Sorry {0}, {1}'.format(message.author.mention, e))
+        else:
+            embed = discord.Embed(description=result, color=0xaaaccc)
+            await self.send_message(message.channel, embed=embed)
+
     async def _do_nextstream(self, server, user, message, **kwargs):
         def format_delta(d):
             if d.total_seconds() <= 0:
diff --git a/commands.py b/commands.py
index 3658464..94e8151 100644
--- a/commands.py
+++ b/commands.py
@@ -5,6 +5,7 @@ import requests
 
 from services.youtube import Youtube, YoutubeError
 from services.cheesecom import CheeseCom, CheeseComError
+from services.roll20 import Roll20, Roll20Error
 
 
 INSTAGRAM_BASE_URL = 'https://www.instagram.com'
@@ -190,6 +191,14 @@ class Commands(object):
         else:
             return event
 
+    def roll(self, formula):
+        try:
+            result = Roll20.execute(formula)
+        except Roll20Error:
+            raise CommandError('failed to interpret the formula')
+        else:
+            return result
+
     def _get_instagram_media(self, params):
         api_url = self.config['Instagram'].get('api_url')
         r = requests.get('{0}/media'.format(api_url), params=params)
diff --git a/requirements.txt b/requirements.txt
index debd2c8..2c2bd37 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -2,6 +2,7 @@ discord.py
 fuzzywuzzy[speedup]
 google-api-python-client
 irc
+parsy
 pyquery
 python-dateutil
 python-twitter
diff --git a/services/roll20.py b/services/roll20.py
new file mode 100644
index 0000000..5abd8bc
--- /dev/null
+++ b/services/roll20.py
@@ -0,0 +1,516 @@
+import random
+
+import parsy
+
+
+# TODO: functions
+# TODO: dice matching
+# FIXME: handle infinite loops
+
+
+def num(x):
+    if int(x) == x:
+        return int(x)
+    return float(x)
+
+
+class Group(object):
+    def __init__(self, items):
+        self.items = items
+        self.keep = None
+        self.drop = None
+        self.succ = None
+        self.fail = None
+
+    def __repr__(self):
+        return '<group {0} keep:{1} drop:{2} succ:{3} fail:{4}>'.format(
+            ' '.join([repr(x) for x in self.items]), self.keep, self.drop, self.succ, self.fail)
+
+    def __str__(self):
+        kept = self._kept()
+        return '**{{** {0} **}}**'.format(' + '.join(
+            [str(x) if i in kept else '~~*{0}*~~'.format(x) for i, x in enumerate(self.items)]))
+
+    # TODO: handle single-sub-rolls - len(self.items) == 1
+
+    def _calculated(self):
+        result = []
+        for item in self.items:
+            try:
+                result.append(item.calc())
+            except AttributeError:
+                result.append(num(item))
+        return result
+
+    def _kept(self):
+        calculated = self._calculated()
+        result = sorted(enumerate(calculated), key=lambda x: (x[1], len(calculated) - x[0]))
+        if self.keep:
+            if self.keep[1]:
+                result = result[:round(self.keep[0])]
+            else:
+                result = result[-round(self.keep[0]):]
+        if self.drop:
+            if self.drop[1]:
+                result = result[:-round(self.drop[0])]
+            else:
+                result = result[round(self.drop[0]):]
+        return list(list(zip(*sorted(result)))[0])
+
+    def filtered(self):
+        calculated = self._calculated()
+        kept = self._kept()
+        return [x for i, x in enumerate(calculated) if i in kept]
+
+    def calc(self):
+        filtered = self.filtered()
+        if self.succ:
+            result = len([x for x in filtered if self.succ(x)])
+            if self.fail:
+                result -= len([x for x in filtered if self.fail(x)])
+            return result
+        return sum(filtered)
+
+
+class Operation(object):
+    def __init__(self, op, func, left, right):
+        self.op = op
+        self.func = func
+        self.left = left
+        self.right = right
+
+    def __repr__(self):
+        return '<{0} {1} {2}>'.format(self.op, repr(self.left), repr(self.right))
+
+    def __str__(self):
+        # FIXME: drop unneeded parentheses
+        return '( {0} {1} {2} )'.format(str(self.left), self.op.replace('*', '\\*'), str(self.right))
+
+    def calc(self):
+        try:
+            left = self.left.calc()
+        except AttributeError:
+            left = num(self.left)
+        try:
+            right = self.right.calc()
+        except AttributeError:
+            right = num(self.right)
+        return self.func(left, right)
+
+
+class Roll(object):
+    def __init__(self, result):
+        self.result = result
+        self.label = None
+        self.keep = None
+        self.drop = None
+        self.succ = None
+        self.fail = None
+
+    def __repr__(self):
+        return '<roll {0} label:{1} keep:{2} drop:{3} succ:{4} fail:{5}>'.format(
+            self.result, self.label, self.keep, self.drop, self.succ, self.fail)
+
+    def __str__(self):
+        kept = self._kept()
+        return '**(** {0} **)**'.format(' + '.join(
+            [str(x) if i in kept else '~~*{0}*~~'.format(x) for i, x in enumerate(self.result)]))
+
+    def _kept(self):
+        result = sorted(enumerate(self.result), key=lambda x: (x[1], len(self.result) - x[0]))
+        if self.keep:
+            if self.keep[1]:
+                result = result[:round(self.keep[0])]
+            else:
+                result = result[-round(self.keep[0]):]
+        if self.drop:
+            if self.drop[1]:
+                result = result[:-round(self.drop[0])]
+            else:
+                result = result[round(self.drop[0]):]
+        return list(list(zip(*sorted(result)))[0])
+
+    def filtered(self):
+        kept = self._kept()
+        return [x for i, x in enumerate(self.result) if i in kept]
+
+    def calc(self):
+        filtered = self.filtered()
+        if self.succ:
+            result = len([x for x in filtered if self.succ(x)])
+            if self.fail:
+                result -= len([x for x in filtered if self.fail(x)])
+            return result
+        else:
+            return sum(filtered)
+
+
+class Parser(object):
+    @classmethod
+    def tokenize(cls, formula):
+        whitespace = parsy.regex(r'\s*')
+        number = parsy.regex(r'(0|[1-9][0-9]*)([.][0-9]+)?([eE][+-]?[0-9]+)?').map(float)
+        expression = (whitespace >> (
+            (parsy.regex(r'\*{2}|[()*/%+-]') | number | parsy.regex(r'\[.*?\]') |
+             parsy.regex(r'\!{2}|\!p|mt|ro|k[hl]|d[hl]|s[ad]|[{}dFfmkdrs!,<>=]')
+            ) << whitespace)).many()
+        return expression.parse(formula)
+
+    @classmethod
+    def parse(cls, tokens):
+        @parsy.generate
+        def group_failures():
+            result = yield group_successes
+            yield parsy.match_item('f')
+            comparison.operator_required = False
+            condition = yield comparison
+            result.fail = condition
+            return result
+
+        @parsy.generate
+        def group_successes():
+            result = yield group_drop
+            comparison.operator_required = True
+            condition = yield comparison | parsy.success('')
+            if not condition:
+                return result
+            result.succ = condition
+            return result
+
+        @parsy.generate
+        def group_drop():
+            result = yield group_keep
+            modifier = yield parsy.test_item(lambda x: x in ['d', 'dh', 'dl'], 'd[h|l]') | parsy.success('')
+            if not modifier:
+                return result
+            count = yield number
+            result.drop = (count, modifier == 'dh')
+            return result
+
+        @parsy.generate
+        def group_keep():
+            result = yield group
+            modifier = yield parsy.test_item(lambda x: x in ['k', 'kh', 'kl'], 'k[h|l]') | parsy.success('')
+            if not modifier:
+                return result
+            count = yield number
+            result.keep = (count, modifier == 'kl')
+            return result
+
+        @parsy.generate
+        def group():
+            yield parsy.match_item('{')
+            result = yield group_simple
+            result = Group([result])
+            while True:
+                end = yield parsy.match_item('}') | parsy.success('')
+                if end:
+                    break
+                yield parsy.match_item(',')
+                other = yield group_simple
+                result.items.append(other)
+            return result
+
+        @parsy.generate
+        def expression_additive():
+            result = yield expression_multiplicative
+            sign = parsy.match_item('+') | parsy.match_item('-')
+            while True:
+                operation = yield sign | parsy.success('')
+                if not operation:
+                    break
+                operand = yield expression_multiplicative
+                if operation == '+':
+                    result = Operation(operation, lambda x, y: x + y, result, operand)
+                elif operation == '-':
+                    result = Operation(operation, lambda x, y: x - y, result, operand)
+            return result
+
+        @parsy.generate
+        def expression_multiplicative():
+            result = yield expression_exponential
+            operator = parsy.match_item('*') | parsy.match_item('/') | parsy.match_item('%')
+            while True:
+                operation = yield operator | parsy.success('')
+                if not operation:
+                    break
+                operand = yield expression_exponential
+                if operation == '*':
+                    result = Operation(operation, lambda x, y: x * y, result, operand)
+                elif operation == '/':
+                    result = Operation(operation, lambda x, y: x / y, result, operand)
+                elif operation == '%':
+                    result = Operation(operation, lambda x, y: x % y, result, operand)
+            return result
+
+        @parsy.generate
+        def expression_exponential():
+            result = yield expression_simple
+            operator = parsy.match_item('**')
+            while True:
+                operation = yield operator | parsy.success('')
+                if not operation:
+                    break
+                operand = yield expression_simple
+                if operation == '**':
+                    result = Operation(operation, lambda x, y: x ** y, result, operand)
+            return result
+
+        @parsy.generate
+        def roll():
+            result = yield failures | successes
+            label = parsy.test_item(lambda x: x.startswith('[') and x.endswith(']'), '[LABEL]')
+            label = yield label | parsy.success(None)
+            if label:
+                label = label[1:-1]
+            result.label = label
+            return result
+
+        @parsy.generate
+        def failures():
+            result = yield successes
+            yield parsy.match_item('f')
+            comparison.operator_required = False
+            condition = yield comparison
+            result.fail = condition
+            return result
+
+        @parsy.generate
+        def successes():
+            result = yield drop
+            comparison.operator_required = True
+            condition = yield comparison | parsy.success('')
+            if not condition:
+                return result
+            result.succ = condition
+            return result
+
+        @parsy.generate
+        def drop():
+            result = yield keep
+            modifier = yield parsy.test_item(lambda x: x in ['d', 'dh', 'dl'], 'd[h|l]') | parsy.success('')
+            if not modifier:
+                return result
+            count = yield number
+            result.drop = (count, modifier == 'dh')
+            return result
+
+        @parsy.generate
+        def keep():
+            result = yield modifiers
+            modifier = yield parsy.test_item(lambda x: x in ['k', 'kh', 'kl'], 'k[h|l]') | parsy.success('')
+            if not modifier:
+                return result
+            count = yield number
+            result.keep = (count, modifier == 'kl')
+            return result
+
+        @parsy.generate
+        def modifiers():
+            result, dice = yield basic_roll
+            modifications = {}
+            while True:
+                prio, modify = yield exploding | compounding | penetrating | reroll | sort | parsy.success((None, None))
+                if not prio:
+                    break
+                if not prio in modifications:
+                    modifications[prio] = []
+                modifications[prio].append(modify)
+            for prio, m in sorted(modifications.items()):
+                if prio == 4: # reroll
+                    changed = True
+                    while changed:
+                        for modify in m:
+                            new_result, dice = modify(result, dice)
+                            changed = new_result != result
+                            result = new_result
+                            if changed:
+                                break
+                else:
+                    for modify in m:
+                        result, dice = modify(result, dice)
+            return Roll(result)
+
+        @parsy.generate
+        def sort():
+            modifier = yield parsy.test_item(lambda x: x in ['s', 'sa', 'sd'], 's[a|d]')
+            def modify(result, dice):
+                result = sorted(result, reverse=(modifier == 'sd'))
+                return result, dice
+            return 5, modify
+
+        @parsy.generate
+        def reroll():
+            modifier = yield parsy.test_item(lambda x: x in ['r', 'ro'], 'r[o]')
+            comparison.operator_required = False
+            condition = yield comparison
+            def modify(result, dice):
+                result = result[:]
+                i = 0
+                while i < len(result):
+                    while condition(result[i], dice[0]):
+                        result[i] = random.choice(dice)
+                        if modifier == 'ro':
+                            break
+                    i += 1
+                return result, dice
+            return 4, modify
+
+        @parsy.generate
+        def penetrating():
+            yield parsy.match_item('!p')
+            comparison.operator_required = False
+            condition = yield comparison
+            def modify(result, dice):
+                result = result[:]
+                i = 0
+                while i < len(result):
+                    sub = [result[i]]
+                    while condition(sub[-1], dice[-1]):
+                        sub.append(random.choice(dice))
+                    result[i+1 : i+1] = [x - 1 for x in sub][1:]
+                    i += len(sub)
+                return result, dice
+            return 3, modify
+
+        @parsy.generate
+        def compounding():
+            yield parsy.match_item('!!')
+            comparison.operator_required = False
+            condition = yield comparison
+            def modify(result, dice):
+                result = result[:]
+                i = 0
+                while i < len(result):
+                    sub = [result[i]]
+                    while condition(sub[-1], dice[-1]):
+                        sub.append(random.choice(dice))
+                    result[i] = sum(sub)
+                    i += 1
+                return result, dice
+            return 2, modify
+
+        @parsy.generate
+        def exploding():
+            yield parsy.match_item('!')
+            comparison.operator_required = False
+            condition = yield comparison
+            def modify(result, dice):
+                result = result[:]
+                i = 0
+                while i < len(result):
+                    sub = [result[i]]
+                    while condition(sub[-1], dice[-1]):
+                        sub.append(random.choice(dice))
+                    result[i: i+1] = sub
+                    i += len(sub)
+                return result, dice
+            return 1, modify
+
+        @parsy.generate
+        def comparison():
+            operator = parsy.match_item('<') | parsy.match_item('>') | parsy.match_item('=')
+            if comparison.operator_required:
+                operation = yield operator
+            else:
+                operation = yield operator | parsy.success('')
+            if operation:
+                operand = yield number
+            else:
+                operand = yield number | parsy.success('')
+            def condition(value, default=None):
+                if operation == '<':
+                    return value <= operand
+                elif operation == '>':
+                    return value >= operand
+                elif operation == '=':
+                    return value == operand
+                elif operand:
+                    return value == operand
+                else:
+                    return value == default
+            return condition
+
+        @parsy.generate
+        def basic_roll():
+            count = yield computed_simple
+            yield parsy.match_item('d')
+            sides = yield computed_simple | parsy.match_item('F')
+            dice = [-1, 0, 1] if sides == 'F' else [x + 1 for x in range(round(sides))]
+            result = [random.choice(dice) for _ in range(round(count))]
+            return result, dice
+
+        @parsy.generate
+        def computed_additive():
+            result = yield computed_multiplicative
+            sign = parsy.match_item('+') | parsy.match_item('-')
+            while True:
+                operation = yield sign | parsy.success('')
+                if not operation:
+                    break
+                operand = yield computed_multiplicative
+                if operation == '+':
+                    result += operand
+                elif operation == '-':
+                    result -= operand
+            return result
+
+        @parsy.generate
+        def computed_multiplicative():
+            result = yield computed_exponential
+            operator = parsy.match_item('*') | parsy.match_item('/') | parsy.match_item('%')
+            while True:
+                operation = yield operator | parsy.success('')
+                if not operation:
+                    break
+                operand = yield computed_exponential
+                if operation == '*':
+                    result *= operand
+                elif operation == '/':
+                    result /= operand
+                elif operation == '%':
+                    result %= operand
+            return result
+
+        @parsy.generate
+        def computed_exponential():
+            result = yield computed_simple
+            operator = parsy.match_item('**')
+            while True:
+                operation = yield operator | parsy.success('')
+                if not operation:
+                    break
+                operand = yield computed_simple
+                if operation == '**':
+                    result **= operand
+            return result
+
+        @parsy.generate
+        def number():
+            sign = yield parsy.match_item('+') | parsy.match_item('-') | parsy.success('+')
+            value = yield parsy.test_item(lambda x: isinstance(x, float), 'number')
+            return num(value if sign == '+' else -value)
+
+        computed_simple = (parsy.match_item('(') >> computed_additive << parsy.match_item(')')) | number
+        expression_simple = roll | (parsy.match_item('(') >> expression_additive << parsy.match_item(')')) | number
+        group_simple = group_failures | group_successes | expression_additive
+        return group_simple.parse(tokens)
+
+
+class Roll20Error(Exception):
+    pass
+
+
+class Roll20(object):
+    @classmethod
+    def execute(cls, formula):
+        try:
+            tokens = Parser.tokenize(formula)
+            result = Parser.parse(tokens)
+        except parsy.ParseError as e:
+            raise Roll20Error(str(e))
+        try:
+            calculated = result.calc()
+        except AttributeError:
+            calculated = num(result)
+        return '{0} = __**{1}**__'.format(str(result), calculated)