Switch to roll implementation based on pyparsing

roll.py

@@ -1,11 +1,14 @@
 #!/usr/bin/env python
 
+import attr
 import logging
 import re
 import sys
 import readline
+import operator
+from numbers import Number
 from random import randint
-
+from pyparsing import Regex, oneOf, Optional, Group, Combine, Literal, CaselessLiteral, ZeroOrMore, StringStart, StringEnd, opAssoc, infixNotation, ParseException, Empty, pyparsing_common, ParseResults
 logFormatter = logging.Formatter('%(asctime)s %(levelname)s: %(message)s')
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.INFO)
@@ -14,108 +17,548 @@ logger.addHandler(logging.StreamHandler())
 for handler in logger.handlers:
     handler.setFormatter(logFormatter)
 
-def format_roll(s, n=1, mod=0, drop_low=0, drop_high=0):
-    fmt = ''
-    if n > 1:
-        fmt += str(n)
-    fmt += 'd' + str(s)
-    if drop_low > 0:
-        fmt += "-L"
-        if drop_low > 1:
-            fmt += str(drop_low)
-    if drop_high > 0:
-        fmt += "-H"
-        if drop_high > 1:
-            fmt += str(drop_high)
-    if mod != 0:
-        fmt += '%+i' % (mod,)
-    return fmt
-
-def roll_die(s, mod=0, adv=0):
-    '''Roll a single S-sided die. '''
-    roll = randint(1, s) + mod
-    logger.debug('%s rolled: %s', format_roll(s, 1, mod), roll)
-    return roll
-
-def roll_dice(s, n=1, mod=0, drop_low=0, drop_high=0):
-    '''Roll n s-sided dice, then add modifier.
-
-    (The modifier feature of this function currently isn't used by the
-    script.)
+@attr.s
+class IntegerValidator(object):
+    min_val = attr.ib(default='-inf', convert=float)
+    max_val = attr.ib(default='+inf', convert=float)
+    handle_float = attr.ib(default='exception')
+    @handle_float.validator
+    def validate_handle_float(self, attribute, value):
+        assert value in ('exception', 'truncate', 'round')
+
+    def __call__(self, value):
+        try:
+            xf = float(value)
+        except ValueError:
+            raise ValueError('{} does not look like a number'.format(value))
+        if not xf.is_integer():
+            if self.handle_float == 'exception':
+                raise ValueError('{} is not an integer'.format(value))
+            elif self.handle_float == 'truncate':
+                x = int(xf)
+            else:
+                x = round(xf)
+        else:
+            x = int(xf)
+        if self.min_val is not None:
+            assert x >= self.min_val
+        if self.max_val is not None:
+            assert x <= self.max_val
+        return x
+
+def normalize_die_type(x):
+    if x == 'F':
+        return x
+    elif x == '%':
+        return 100
+    else:
+        try:
+            return IntegerValidator(min_val=2, handle_float='exception')(x)
+        except Exception:
+            raise ValueError('Invalid die type: d{}'.format(x))
+
+def normalize_dice_count(x):
+    xf = float(x)
+    x = int(x)
+    if not xf.is_integer():
+        raise ValueError('dice count must be an integer, not {}'.format(xf))
+    if x < 1:
+        raise ValueError("dice count must be positive; {} is invalid".format(x))
+    return x
+
+def ImplicitToken(x):
+    '''Like pyparsing.Empty, but yields one or more tokens instead of nothing.'''
+    return Empty().setParseAction(lambda toks: x)
+
+# TODO: Look at http://infohost.nmt.edu/tcc/help/pubs/pyparsing/web/classes.html#class-ParserElement
+
+# Implementing the syntax described here: https://www.critdice.com/roll-advanced-dice
+# https://stackoverflow.com/a/23956778/125921
+
+# https://stackoverflow.com/a/46583691/125921
+
+var_name = pyparsing_common.identifier.copy()
+real_num = pyparsing_common.fnumber.copy()
+positive_int = pyparsing_common.integer.copy().setParseAction(lambda toks: [ IntegerValidator(min_val=1)(toks[0]) ])
+
+drop_type = oneOf('K k X x -H -L')
+drop_spec = Group(drop_type.setResultsName('type') +
+                    (positive_int | ImplicitToken(1)).setResultsName('count')
+).setResultsName('drop')
+
+pos_int_implicit_one = (positive_int | ImplicitToken(1))
+
+comparator_type = oneOf('<= < >= > ≤ ≥ =')
+
+reroll_type = Combine(oneOf('R r') | ( oneOf('! !!') + Optional('p')))
+reroll_spec = Group(
+    reroll_type.setResultsName('type') +
+    Optional(
+        (comparator_type | ImplicitToken('=')).setResultsName('operator') + \
+        positive_int.setResultsName('value')
+    )
+).setResultsName('reroll')
+
+count_spec = Group(
+    Group(
+        comparator_type.setResultsName('operator') + \
+        positive_int.setResultsName('value')
+    ).setResultsName('success_condition') +
+    Optional(
+        Literal('f') +
+        Group(
+            comparator_type.setResultsName('operator') + \
+            positive_int.setResultsName('value')
+        ).setResultsName('failure_condition')
+    )
+).setResultsName('count_successes')
+
+roll_spec = Group(
+    (positive_int | ImplicitToken(1)).setResultsName('dice_count') +
+    CaselessLiteral('d') +
+    (positive_int | oneOf('% F')).setResultsName('die_type') +
+    Optional(reroll_spec | drop_spec) +
+    Optional(count_spec)
+).setResultsName('roll')
+
+expr_parser = (StringStart() + infixNotation(
+    baseExpr=(roll_spec | positive_int | real_num | var_name),
+    opList=[
+        (oneOf('**').setResultsName('operator', True), 2, opAssoc.RIGHT),
+        (oneOf('* / × ÷').setResultsName('operator', True), 2, opAssoc.LEFT),
+        (oneOf('+ -').setResultsName('operator', True), 2, opAssoc.LEFT),
+    ]
+) + StringEnd())
+
+
+def roll_die(sides=6):
+    '''Roll a single die.
+
+Supports any valid integer number of sides as well as 'F' for a fate
+die, which can return -1, 0, or 1 with equal probability.
+
+    '''
+    if sides == 'F':
+        # Fate die = 1d3-2
+        return roll_die(3) - 2
+    else:
+        return randint(1, int(sides))
+
+class DieRolled(int):
+    '''Subclass of int that allows a string suffix.
+
+    This is meant for recording the result of rolling a die. The
+    suffix is purely cosmetic, for the purposes of string conversion.
+    It can be used to indicate a die roll that has been re-rolled or
+    exploded, or to indicate a critical hit/miss.
 
     '''
-    s = int(s)
-    n = int(n)
-    mod = int(mod)
-    drop_low = int(drop_low)
-    drop_high = int(drop_high)
-    if n < 1:
-        raise ValueError('Must roll at least one die.')
-    if s < 2:
-        raise ValueError('Dice must have at least 2 sides')
-    if drop_low < 0 or drop_high < 0:
-        raise ValueError('Cannot drop negative number of dice.')
-    if drop_low + drop_high >= n:
-        raise ValueError('Dropping too many dice; must keep at least one die.')
-    n = int(n)
-    rolls = [ roll_die(s) for i in range(n) ]
-    dropped_low_rolls = sorted(rolls)[:drop_low]
-    dropped_high_rolls = sorted(rolls, reverse=True)[:drop_high]
-    kept_rolls = list(rolls)
-    # Cannot use setdiff because of possible repeated values
-    for drop in dropped_low_rolls + dropped_high_rolls:
-        kept_rolls.remove(drop)
-    total = mod + sum(kept_rolls)
-    # TODO: Special reporting for natural 1 and natural 20 (only when a single d20 roll is returned)
-    natural = ''
-    if len(kept_rolls) == 1 and s == 20:
-        if kept_rolls[0] == 1:
-            natural = " (NATURAL 1)"
-        elif kept_rolls[0] == 20:
-            natural = " (NATURAL 20)"
-    if n > 1:
-        paren_stmts = [ ('Individual rolls', kept_rolls), ]
-        if dropped_low_rolls:
-            paren_stmts.append( ('Dropped low', dropped_low_rolls) )
-        if dropped_high_rolls:
-            paren_stmts.append( ('Dropped high', dropped_high_rolls) )
-        if dropped_low_rolls or dropped_high_rolls:
-            paren_stmts.append( ('Original rolls', rolls) )
-        paren_stmt = "; ".join("%s: %s" % (k, repr(v)) for k,v in paren_stmts)
-        logger.info('%s rolled: %s%s\n(%s)', format_roll(s, n, mod, drop_low, drop_high), total, natural, paren_stmt)
+    def __new__(cls, value, formatter='{}'):
+        newval = super(DieRolled, cls).__new__(cls, value)
+        newval.formatter = formatter
+        return newval
+
+    def __str__(self):
+        return self.formatter.format(super().__str__())
+
+    def __repr__(self):
+        if self.formatter != '{}':
+            return 'DieRolled(value={value!r}, formatter={formatter!r})'.format(
+                value=int(self),
+                formatter=self.formatter,
+            )
+        else:
+            return 'DieRolled({value!r})'.format(value=int(self))
+
+def validate_dice_roll_list(instance, attribute, value):
+    for x in value:
+        # Not using positive_int here because 0 is a valid roll for
+        # penetrating dice
+        pyparsing_common.integer.parseString(str(x))
+def format_dice_roll_list(rolls, always_list=False):
+    if len(rolls) == 0:
+        raise ValueError('Need at least one die rolled')
+    elif len(rolls) == 1 and not always_list:
+        return str(rolls[0])
     else:
-        logger.info('%s rolled: %s%s', format_roll(s, n, mod, drop_low, drop_high), total, natural)
-    return total
-
-def _roll_matchgroup(m):
-    sides = int(m.group(2))
-    n = int(m.group(1) or 1)
-    dropspec = m.group(3) or ''
-    mod = int(m.group(4) or 0)
-    drop_low, drop_high = 0,0
-    drops = re.findall('-([HL])(\\d*)', dropspec)
-    for (dtype, dnum) in drops:
-        if dnum == '':
-            dnum = 1
+        return '[' + ",".join(map(str, rolls)) + ']'
+
+@attr.s
+class DiceRolled(object):
+    '''Class representing the result of rolling one or more similar dice.'''
+    dice_results = attr.ib(convert=list, validator = validate_dice_roll_list)
+    dropped_results = attr.ib(default=attr.Factory(list), convert=list,
+                              validator = validate_dice_roll_list)
+    roll_desc = attr.ib(default='', convert=str)
+    success_count = attr.ib(default=None)
+    @success_count.validator
+    def validate_success_count(self, attribute, value):
+        if value is not None:
+            self.success_count = int(value)
+
+    def __attrs_post_init__(self):
+        if len(self.dice_results) < 1:
+            raise ValueError('Need at least one non-dropped roll')
+
+    def total(self):
+        if self.success_count is not None:
+            return int(self.success_count)
+        else:
+            return sum(self.dice_results)
+
+    def __str__(self):
+        if self.roll_desc:
+            prefix = '{roll} rolled'.format(roll=self.roll_desc)
+        else:
+            prefix = 'Rolled'
+        if self.dropped_results:
+            drop = ' (dropped {dropped})'.format(dropped = format_dice_roll_list(self.dropped_results))
         else:
-            dnum = int(dnum)
-        if dtype == 'L':
-            drop_low += dnum
+            drop = ''
+        if self.success_count is not None:
+            tot = ', Total successes: ' + str(self.total())
+        elif len(self.dice_results) > 1:
+            tot = ', Total: ' + str(self.total())
         else:
-            drop_high += dnum
-    return str(roll_dice(sides, n, mod, drop_low, drop_high))
-
-def roll(expr):
-    try:
-        logger.info("Rolling %s", expr)
-        # Multi dice rolls, e.g. "4d6"
-        subbed = re.sub('(\\d+)?\\s*d\\s*(\\d+)((?:-[HL]\\d*)+)?([+-]\\d+)?',
-                           lambda m: _roll_matchgroup(m),
-                           expr)
-        logger.debug("Expression with dice rolls substituted: %s", repr(subbed))
-        return int(eval(subbed))
-    except Exception as exc:
-        raise Exception("Cannot parse expression %s: %s" % (repr(expr), exc))
+            tot = ''
+        return '{prefix}: {results}{drop}{tot}'.format(
+            prefix=prefix,
+            results=format_dice_roll_list(self.dice_results),
+            drop=drop,
+            tot=tot,
+        )
+
+    def __int__(self):
+        return self.total()
+
+def validate_by_parser(parser):
+    def private_validator(instance, attribute, value):
+        parser.parseString(str(value), True)
+    return private_validator
+
+@attr.s
+class Comparator(object):
+    cmp_dict = {
+        '<=': operator.le,
+        '<': operator.lt,
+        '>=': operator.ge,
+        '>': operator.gt,
+        '≤': operator.le,
+        '≥': operator.ge,
+        '=': operator.eq,
+    }
+    operator = attr.ib(convert=str, validator=validate_by_parser(comparator_type))
+    value = attr.ib(convert=int, validator=validate_by_parser(positive_int))
+
+    def __str__(self):
+        return '{op}{val}'.format(op=self.operator, val=self.value)
+
+    def compare(self, x):
+        '''Return True if x satisfies the comparator.
+
+        In other words, x is placed on the left-hand side of the
+        comparison, the Comparator's value is placed on the right hand
+        side, and the truth value of the resulting test is returned.
+
+        '''
+        return self.cmp_dict[self.operator](x, self.value)
+
+@attr.s
+class RerollSpec(object):
+    type = attr.ib(convert=str, validator=validate_by_parser(reroll_type))
+    operator = attr.ib(default=None)
+    value = attr.ib(default=None)
+    comparator = attr.ib(default=None, repr=False)
+
+    def __attrs_post_init__(self):
+        if self.comparator is None:
+            self.comparator = Comparator(self.operator, self.value)
+        else:
+            if self.operator is not None or self.value is not None:
+                raise ValueError('Do not provide opeartor or value if providing a pre-build comparator')
+            self.operator = self.comparator.operator
+            self.value = self.comparator.value
+
+    def __str__(self):
+        return '{typ}{cmp}'.format(typ=self.type, cmp=self.comparator)
+
+    def compare(self, x):
+        return self.comparator.compare(x)
+
+    def roll_die(self, sides):
+        '''Roll a single die, following specified re-rolling rules.
+
+        Returns a list of rolls, since some types of re-rolling
+        collect the result of multiple die rolls.
+
+        '''
+        if sides == 'F':
+            raise ValueError("Re-rolling/exploding is incompatible with Fate")
+        if self.type == 'r':
+            # Single reroll
+            roll = roll_die(sides)
+            if self.compare(roll):
+                roll = DieRolled(roll_die(sides), '{}' + self.type)
+            return [ roll ]
+        elif self.type == 'R':
+            # Indefinite reroll
+            roll = roll_die(sides)
+            while self.compare(roll):
+                roll = DieRolled(roll_die(sides), '{}' + self.type)
+            return [ roll ]
+        elif self.type in ['!', '!!', '!p', '!!p']:
+            # Explode/penetrate/compound
+            all_rolls = [ roll_die(sides) ]
+            while self.compare(all_rolls[-1]):
+                all_rolls.append(roll_die(sides))
+            # If we never re-rolled, no need to do anything special
+            if len(all_rolls) == 1:
+                return all_rolls
+            # For penetration, subtract 1 from all rolls except the first
+            if self.type.endswith('p'):
+                for i in range(1, len(all_rolls)):
+                    all_rolls[i] -= 1
+            # For compounding, return the sum
+            if self.type.startswith('!!'):
+                total = sum(all_rolls)
+                return [ DieRolled(total, '{}' + self.type) ]
+            else:
+                for i in range(0, len(all_rolls)-1):
+                    all_rolls[i] = DieRolled(all_rolls[i], '{}' + self.type)
+                return all_rolls
+
+@attr.s
+class DropSpec(object):
+    type = attr.ib(convert=str, validator=validate_by_parser(drop_type))
+    count = attr.ib(default=1, convert=int, validator=validate_by_parser(positive_int))
+
+    def __str__(self):
+        return self.type + str(self.count)
+
+    def drop_rolls(self, rolls):
+        '''Drop the appripriate rolls from a list of rolls.
+
+        Returns a 2-tuple of roll lists. The first list is the kept
+        rolls, and the second list is the dropped rolls.
+
+        The order of the rolls is not preserved.
+
+        '''
+        if not isinstance(rolls, list):
+            rolls = list(rolls)
+        keeping = self.type in ('K', 'k')
+        # if keeping:
+        #     num_to_keep = self.count
+        # else:
+        #     num_to_keep = len(rolls) - self.count
+        # if num_to_keep == 0:
+        #     raise ValueError('Not enough rolls: would drop all rolls')
+        # elif num_to_keep == len(rolls):
+        #     raise ValueError('Keeping too many rolls: would not drop any rolls')
+        rolls.sort()
+        if self.type in ('K', 'X', '-H'):
+            rolls.reverse()
+        (head, tail) = rolls[:self.count], rolls[self.count:]
+        if keeping:
+            (kept, dropped) = (head, tail)
+        else:
+            (kept, dropped) = (tail, head)
+        return (kept, dropped)
+
+@attr.s
+class DiceRoller(object):
+    die_type = attr.ib(convert = normalize_die_type)
+    dice_count = attr.ib(default=1, convert=normalize_dice_count)
+    reroll_spec = attr.ib(default=None)
+    @reroll_spec.validator
+    def validate_reroll_spec(self, attribute, value):
+        if value is not None:
+            assert isinstance(value, RerollSpec)
+    drop_spec = attr.ib(default=None)
+    @drop_spec.validator
+    def validate_drop_spec(self, attribute, value):
+        if value is not None:
+            assert isinstance(value, DropSpec)
+    success_comparator = attr.ib(default=None)
+    failure_comparator = attr.ib(default=None)
+    @success_comparator.validator
+    @failure_comparator.validator
+    def validate_comparator(self, attribute, value):
+        if value is not None:
+            assert isinstance(value, Comparator)
+
+    def __attrs_post_init__(self):
+        if self.reroll_spec is not None and self.drop_spec is not None:
+            raise ValueError('Reroll and drop specs are mutually exclusive')
+        if self.success_comparator is None and self.failure_comparator is not None:
+            raise ValueError('Cannot use a failure condition without a success condition')
+
+    def __str__(self):
+        return '{count}d{type}{reroll}{drop}{success}{fail}'.format(
+            count = self.dice_count,
+            type = self.die_type,
+            reroll = self.reroll_spec or '',
+            drop = self.drop_spec or '',
+            success = self.success_comparator or '',
+            fail = 'f' + str(self.success_comparator) if self.success_comparator else '',
+        )
+
+    def roll(self):
+        '''Roll dice according to specifications. Returns a DiceRolled object.'''
+        all_rolls = []
+        if self.reroll_spec:
+            for i in range(self.dice_count):
+                all_rolls.extend(self.reroll_spec.roll_die(self.die_type))
+        else:
+            for i in range(self.dice_count):
+                all_rolls.append(roll_die(self.die_type))
+        if self.drop_spec:
+            (dice_results, dropped_results) = self.drop_spec.drop_rolls(all_rolls)
+        else:
+            (dice_results, dropped_results) = (all_rolls, [])
+        if self.success_comparator is not None:
+            success_count = 0
+            for roll in dice_results:
+                if self.success_comparator.compare(roll):
+                    success_count += 1
+            if self.failure_comparator is not None:
+                for roll in dice_results:
+                    if self.failure_comparator.compare(roll):
+                        success_count -= 1
+        else:
+            success_count = None
+        return DiceRolled(
+            dice_results=dice_results,
+            dropped_results=dropped_results,
+            roll_desc=str(self),
+            success_count=success_count,
+        )
+
+def make_dice_roller(expr):
+    if isinstance(expr, str):
+        expr = roll_spec.parseString(expr, True)[0]
+    assert expr.getName() == 'roll'
+    expr = expr.asDict()
+
+    dtype = normalize_die_type(expr['die_type'])
+    dcount = normalize_dice_count(expr['dice_count'])
+    constructor_args = {
+        'die_type': dtype,
+        'dice_count': dcount,
+        'reroll_spec': None,
+        'drop_spec': None,
+        'success_comparator': None,
+        'failure_comparator': None,
+    }
+
+    rrdict = None
+    if 'reroll' in expr:
+        rrdict = expr['reroll']
+        if 'value' not in rrdict:
+            rrdict['operator'] = '='
+            # Default value is 1 for rerollers, dtype for exploders
+            if rrdict['type'] in ('R', 'r'):
+                rrdict['value'] = 1
+            else:
+                rrdict['value'] = dtype
+        constructor_args['reroll_spec'] = RerollSpec(**rrdict)
+
+    if 'drop' in expr:
+        ddict = expr['drop']
+        constructor_args['drop_spec'] = DropSpec(**ddict)
+
+    if 'count_successes' in expr:
+        csdict = expr['count_successes']
+        constructor_args['success_comparator'] = Comparator(**csdict['success_condition'])
+        if 'failure_condition' in csdict:
+            constructor_args['failure_comparator'] = Comparator(**csdict['failure_condition'])
+    return DiceRoller(**constructor_args)
+
+# examples = [
+#     '1+1',
+#     '1 + 1 + x',
+#     '3d8',
+#     '2e3 * 4d6 + 2',
+#     '2d20k',
+#     '3d20x2',
+#     '4d4rK3',
+#     '4d4R4',
+#     '4d4R>=3',
+#     '4d4r>=3',
+#     '4d4!1',
+#     '4d4!<3',
+#     '4d4!p',
+#     '2D20K+10',
+#     '2D20k+10',
+#     '10d6X4',
+#     '4d8r + 6',
+#     '20d6R≤2',
+#     '6d10!≥8+6',
+#     '10d4!p',
+#     '20d6≥6',
+#     '8d12≥10f≤2',
+#     # '4d20R<=2!>=19Xx21>=20f<=5*2+3',  # Pretty much every possible feature
+# ]
+
+# example_results = {}
+# for x in examples:
+#     try:
+#         example_results[x] = parse_roll(x)
+#     except ParseException as ex:
+#         example_results[x] = ex
+# example_results
+
+# rs = RerollSpec('!!p', '=', 6)
+# rs.roll_die(6)
+
+# ds = DropSpec('K', 2)
+# ds.drop_rolls([1,2,3,4,5])
+# ds = DropSpec('x', 2)
+# ds.drop_rolls([1,2,3,4,5])
+
+# parse_roll = lambda x: expr_parser.parseString(x)[0]
+# exprstring = 'x + 1 + (2 + (3 + 4))'
+# expr = parse_roll(exprstring)
+
+# r = parse_roll('x + 1 - 2 * y * 4d4 + 2d20K1>=20f<=5')[0]
+
+op_dict = {
+    '+': operator.add,
+    '-': operator.sub,
+    '*': operator.mul,
+    '×': operator.mul,
+    '/': operator.truediv,
+    '÷': operator.truediv,
+    '**': operator.pow,
+}
+
+def eval_expr(expr, env={}, print_rolls=True):
+    if isinstance(expr, str):
+        expr = expr_parser.parseString(expr)[0]
+    if isinstance(expr, Number):
+        # Numeric literal
+        return expr
+    elif isinstance(expr, str):
+        # variable name
+        raise NotImplementedError("Variables are not implemented yet")
+    if 'operator' in expr:
+        # Compound expression
+        operands = expr[::2]
+        operators = expr[1::2]
+        assert len(operands) == len(operators) + 1
+        values = [ eval_expr(x) for x in operands ]
+        result = values[0]
+        for (op, nextval) in zip(operators, values[1:]):
+            opfun = op_dict[op]
+            result = opfun(result, nextval)
+        return result
+    else:
+        # roll specification
+        roller = make_dice_roller(expr)
+        result = roller.roll()
+        if print_rolls:
+            print(result)
+        return int(result)
 
 def read_roll(handle=sys.stdin):
     return input("Enter roll> ")
@@ -138,8 +581,9 @@ if __name__ == '__main__':
                         break
                     if re.search("\\S", expr):
                         try:
-                            result = roll(expr)
+                            result = eval_expr(expr)
                             logger.info("Total roll: %s", result)
+                            print('', file=sys.stderr)
                         except Exception as exc:
                             logger.error("Error while rolling: %s", repr(exc))
                 except KeyboardInterrupt: