dr.botzo/dice/roller.py

"""Dice rollers used by the views, bots, etc."""
import logging
import random

import ply.lex as lex
import ply.yacc as yacc

logger = logging.getLogger(__name__)


class DiceRoller(object):

    tokens = ['NUMBER', 'TEXT', 'ROLLSEP']
    literals = ['#', '/', '+', '-', 'd']

    t_TEXT = r'\s+[^;]+'
    t_ROLLSEP = r';\s*'

    def build(self):
        lex.lex(module=self)
        yacc.yacc(module=self)

    def t_NUMBER(self, t):
        r'\d+'
        t.value = int(t.value)
        return t

    def t_error(self, t):
        t.lexer.skip(1)

    precedence = (
        ('left', 'ROLLSEP'),
        ('left', '+', '-'),
        ('right', 'd'),
        ('left', '#'),
        ('left', '/')
    )

    output = ""

    def roll_dice(self, keep, dice, size):
        """Takes the parsed dice string for a single roll (eg 3/4d20) and performs
        the actual roll. Returns a string representing the result.
        """

        a = list(range(dice))
        for i in range(dice):
            a[i] = random.randint(1, size)
        if keep != dice:
            b = sorted(a, reverse=True)
            b = b[0:keep]
        else:
            b = a
        total = sum(b)
        outstr = "[" + ",".join(str(i) for i in a) + "]"

        return (total, outstr)

    def process_roll(self, trials, mods, comment):
        """Processes rolls coming from the parser.

        This generates the inputs for the roll_dice() command, and returns
        the full string representing the whole current dice string (the part
        up to a semicolon or end of line).
        """
        rolls = mods[0][0][1] if mods[0][0][1] else 1
        if trials:
            assert trials <= 10, "Too many rolls (max: 10)."
        assert rolls <= 50, "Too many dice (max: 50) in roll."

        output = ""
        repeat = 1
        if trials is not None:
            repeat = trials
        for i in range(repeat):
            mode = 1
            total = 0
            curr_str = ""
            if i > 0:
                output += ", "
            for m in mods:
                keep = 0
                dice = 1
                res = 0
                # if m is a tuple, then it is a die roll
                # m[0] = (keep, num dice)
                # m[1] = num faces on the die
                if type(m) == tuple:
                    if m[0] is not None:
                        if m[0][0] is not None:
                            keep = m[0][0]
                        dice = m[0][1]
                    size = m[1]
                    if keep > dice or keep == 0:
                        keep = dice
                    if size < 1:
                        output = "# of sides for die is incorrect: %d" % size
                        return output
                    if dice < 1:
                        output = "# of dice is incorrect: %d" % dice
                        return output
                    res = self.roll_dice(keep, dice, size)
                    curr_str += "%d%s" % (res[0], res[1])
                    res = res[0]
                elif m == "+":
                    mode = 1
                    curr_str += "+"
                elif m == "-":
                    mode = -1
                    curr_str += "-"
                else:
                    res = m
                    curr_str += str(m)
                total += mode * res
            if repeat == 1:
                if comment is not None:
                    output = "%d %s (%s)" % (total, comment.strip(), curr_str)
                else:
                    output = "%d (%s)" % (total, curr_str)
            else:
                output += "%d (%s)" % (total, curr_str)
                if i == repeat - 1:
                    if comment is not None:
                        output += " (%s)" % (comment.strip())
        return output

    def p_roll_r(self, p):
        # Chain rolls together.

        # General idea I had when creating this grammar: A roll string is a chain
        # of modifiers, which may be repeated for a certain number of trials. It can
        # have a comment that describes the roll
        # Multiple roll strings can be chained with semicolon

        'roll : roll ROLLSEP roll'
        global output
        p[0] = p[1] + "; " + p[3]
        output = p[0]

    def p_roll(self, p):
        # Parse a basic roll string.

        'roll : trial modifier comment'
        global output
        mods = []
        if type(p[2]) == list:
            mods = p[2]
        else:
            mods = [p[2]]
        p[0] = self.process_roll(p[1], mods, p[3])
        output = p[0]

    def p_roll_no_trials(self, p):
        # Parse a roll string without trials.

        'roll : modifier comment'
        global output
        mods = []
        if type(p[1]) == list:
            mods = p[1]
        else:
            mods = [p[1]]
        p[0] = self.process_roll(None, mods, p[2])
        output = p[0]

    def p_comment(self, p):
        # Parse a comment.

        '''comment : TEXT
                   |'''
        if len(p) == 2:
            p[0] = p[1]
        else:
            p[0] = None

    def p_modifier(self, p):
        # Parse a modifier on a roll string.

        '''modifier : modifier "+" modifier
                    | modifier "-" modifier'''
        # Use append to prevent nested lists (makes dealing with this easier)
        if type(p[1]) == list:
            p[1].append(p[2])
            p[1].append(p[3])
            p[0] = p[1]
        elif type(p[3]) == list:
            p[3].insert(0, p[2])
            p[3].insert(0, p[1])
            p[0] = p[3]
        else:
            p[0] = [p[1], p[2], p[3]]

    def p_die(self, p):
        # Return the left side before the "d", and the number of faces.

        'modifier : left NUMBER'
        p[0] = (p[1], p[2])

    def p_die_num(self, p):
        'modifier : NUMBER'
        p[0] = p[1]

    def p_left(self, p):
        # Parse the number of dice we are rolling, and how many we are keeping.
        'left : keep dice'
        if p[1] == None:
            p[0] = [None, p[2]]
        else:
            p[0] = [p[1], p[2]]

    def p_left_all(self, p):
        'left : dice'
        p[0] = [None, p[1]]

    def p_left_e(self, p):
        'left :'
        p[0] = None

    def p_total(self, p):
        'trial : NUMBER "#"'
        if len(p) > 1:
            p[0] = p[1]
        else:
            p[0] = None

    def p_keep(self, p):
        'keep : NUMBER "/"'
        if p[1] != None:
            p[0] = p[1]
        else:
            p[0] = None

    def p_dice(self, p):
        'dice : NUMBER "d"'
        p[0] = p[1]

    def p_dice_one(self, p):
        'dice : "d"'
        p[0] = 1

    def p_error(self, p):
        """Raise ValueError on unparseable strings."""
        raise ValueError("Error occurred in parser")

    def get_result(self):
        global output
        return output

    def do_roll(self, dicestr):
        """
        Roll some dice and get the result (with broken out rolls).

        Keyword arguments:
        dicestr - format:
        N#X/YdS+M label
            N#: do the following roll N times (optional)
            X/: take the top X rolls of the Y times rolled (optional)
            Y : roll the die specified Y times (optional, defaults to 1)
            dS: roll a S-sided die
            +M: add M to the result (-M for subtraction) (optional)
        """
        self.build()
        yacc.parse(dicestr)
        return self.get_result()