Implement mathematical functions

services/roll20.py

@@ -1,9 +1,9 @@
+import math
 import random
 
 import parsy
 
 
-# TODO: functions
 # TODO: nested groups
 
 
@@ -52,6 +52,9 @@ class Group(object):
             try:
                 traverse(node.left, subrolls)
                 traverse(node.right, subrolls)
+            except AttributeError:
+                try:
+                    traverse(node.operand, subrolls)
                 except AttributeError:
                     try:
                         node.result
@@ -66,7 +69,10 @@ class Group(object):
         def exp(x):
             def copy(node):
                 try:
-                    return Operation(node.op, node.func, copy(node.left), copy(node.right))
+                    return Operation2(node.op, node.func, copy(node.left), copy(node.right))
+                except AttributeError:
+                    try:
+                        return Operation1(node.op, node.func, copy(node.operand))
                     except AttributeError:
                         try:
                             node.result
@@ -81,7 +87,7 @@ class Group(object):
         return exp
 
     def _update_subrolls(self, subrolls):
-        results = [(r, i, j, s) for i, s in enumerate(subrolls) for j, r in enumerate(s.result) if j in s.kept()]
+        results = [(r, i, j, s) for i, s in enumerate(subrolls) for j, r in enumerate(s.result) if j in s.kept(True)]
         results = sorted(results, key=lambda x: (x[0], len(subrolls) - x[1], len(x[3].result) - x[2]))
         if self.keep:
             if self.keep[1]:
@@ -153,7 +159,27 @@ class Group(object):
         return sum(filtered)
 
 
-class Operation(object):
+class Operation1(object):
+    def __init__(self, op, func, operand):
+        self.op = op
+        self.func = func
+        self.operand = operand
+
+    def __repr__(self):
+        return '<{0} {1}>'.format(self.op, repr(self.operand))
+
+    def __str__(self):
+        return '{0}( {1} )'.format(self.op, str(self.operand))
+
+    def calc(self):
+        try:
+            operand = self.operand.calc()
+        except AttributeError:
+            operand = num(self.operand)
+        return self.func(operand)
+
+
+class Operation2(object):
     def __init__(self, op, func, left, right):
         self.op = op
         self.func = func
@@ -211,8 +237,8 @@ class Roll(object):
                 result.append('~~*{0}*~~'.format(x))
         return '**(** {0} **)**'.format(' + '.join(result))
 
-    def kept(self):
+    def kept(self, ignore_group=False):
-        if self.group_kept is not None:
+        if not ignore_group and self.group_kept is not None:
             return self.group_kept
         result = sorted(enumerate(self.result), key=lambda x: (x[1], len(self.result) - x[0]))
         if self.keep:
@@ -225,7 +251,10 @@ class Roll(object):
                 result = result[:-round(self.drop[0])]
             else:
                 result = result[round(self.drop[0]):]
+        try:
             return list(list(zip(*sorted(result)))[0])
+        except IndexError:
+            return []
 
     def filtered(self):
         kept = self.kept()
@@ -248,7 +277,8 @@ class Parser(object):
         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.string('floor') | parsy.string('ceil') | parsy.string('round') | parsy.string('abs') |
+             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)
@@ -298,7 +328,7 @@ class Parser(object):
         def group():
             yield parsy.match_item('{')
             #result = yield group_simple
-            result = yield expression_additive
+            result = yield function | expression_additive
             result = Group([result])
             while True:
                 end = yield parsy.match_item('}') | parsy.success('')
@@ -306,10 +336,26 @@ class Parser(object):
                     break
                 yield parsy.match_item(',')
                 #other = yield group_simple
-                other = yield expression_additive
+                other = yield function | expression_additive
                 result.items.append(other)
             return result
 
+        @parsy.generate
+        def function():
+            func = yield parsy.test_item(lambda x: x in ['floor', 'ceil', 'round', 'abs'], 'floor|ceil|round|abs')
+            yield parsy.match_item('(')
+            operand = yield function | expression_additive
+            yield parsy.match_item(')')
+            if func == 'floor':
+                result = Operation1(func, lambda x: math.floor(x), operand)
+            elif func == 'ceil':
+                result = Operation1(func, lambda x: math.ceil(x), operand)
+            elif func == 'round':
+                result = Operation1(func, lambda x: round(x), operand)
+            elif func == 'abs':
+                result = Operation1(func, lambda x: abs(x), operand)
+            return result
+
         @parsy.generate
         def expression_additive():
             result = yield expression_multiplicative
@@ -320,9 +366,9 @@ class Parser(object):
                     break
                 operand = yield expression_multiplicative
                 if operation == '+':
-                    result = Operation(operation, lambda x, y: x + y, result, operand)
+                    result = Operation2(operation, lambda x, y: x + y, result, operand)
                 elif operation == '-':
-                    result = Operation(operation, lambda x, y: x - y, result, operand)
+                    result = Operation2(operation, lambda x, y: x - y, result, operand)
             return result
 
         @parsy.generate
@@ -335,11 +381,11 @@ class Parser(object):
                     break
                 operand = yield expression_exponential
                 if operation == '*':
-                    result = Operation(operation, lambda x, y: x * y, result, operand)
+                    result = Operation2(operation, lambda x, y: x * y, result, operand)
                 elif operation == '/':
-                    result = Operation(operation, lambda x, y: x / y, result, operand)
+                    result = Operation2(operation, lambda x, y: x / y, result, operand)
                 elif operation == '%':
-                    result = Operation(operation, lambda x, y: x % y, result, operand)
+                    result = Operation2(operation, lambda x, y: x % y, result, operand)
             return result
 
         @parsy.generate
@@ -352,7 +398,7 @@ class Parser(object):
                     break
                 operand = yield expression_simple
                 if operation == '**':
-                    result = Operation(operation, lambda x, y: x ** y, result, operand)
+                    result = Operation2(operation, lambda x, y: x ** y, result, operand)
             return result
 
         @parsy.generate
@@ -551,8 +597,11 @@ class Parser(object):
             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))]
+            dice = [-1, -1, 0, 0, 1, 1] if sides == 'F' else [x + 1 for x in range(round(sides))]
+            try:
                 result = [random.choice(dice) for _ in range(round(count))]
+            except IndexError:
+                raise RuntimeError('Dice has to have at least one side!')
             return result, dice
 
         @parsy.generate
@@ -607,8 +656,8 @@ class Parser(object):
             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
+        expression_simple = roll | function | (parsy.match_item('(') >> expression_additive << parsy.match_item(')')) | number
-        group_simple = group_failures | group_successes | expression_additive
+        group_simple = group_failures | group_successes | function | expression_additive
         return group_simple.parse(tokens)