Complex Dice Expression Parser and Evaluator [Python]

dice.py

import re
import random

import parsec as p
import dataclasses as dc


###########
### AST ###
###########


### Modifiers ###

@dc.dataclass(frozen=True)
class Modifier:
  def match(self, if_take):
    if isinstance(self, TakeModifier):
      return if_take(self)

@dc.dataclass(frozen=True)
class TakeModifier(Modifier):
  amount: int

### Rolls ###

@dc.dataclass(frozen=True)
class Roll:
  def match(self, if_basic, if_modified):
    if isinstance(self, BasicRoll):
      return if_basic(self)
    elif isinstance(self, ModifiedRoll):
      return if_modified(self)

@dc.dataclass(frozen=True)
class BasicRoll(Roll):
  count: int
  sides: int

@dc.dataclass(frozen=True)
class ModifiedRoll(Roll):
  roll: Roll
  modifier: Modifier

### Expr ###

@dc.dataclass(frozen=True)
class RollExpr:
  def match(self, if_roll, if_int, if_binop):
    if isinstance(self, RollValue):
      return if_roll(self)
    elif isinstance(self, IntValue):
      return if_int(self)
    elif isinstance(self, BinOp):
      return if_binop(self)

@dc.dataclass(frozen=True)
class RollValue(RollExpr):
  roll: Roll

@dc.dataclass(frozen=True)
class IntValue(RollExpr):
  value: int

@dc.dataclass(frozen=True)
class BinOp(RollExpr):
  op: str
  left: RollExpr
  right: RollExpr


##############
### Parser ###
##############


whitespace = p.regex(r'\s*')
lexeme = lambda p: p << whitespace
integer = lexeme(p.regex(r"\d+").parsecmap(int))

def optional(par, default_value=None):
  @p.Parser
  def optional_parser(text, index):
    res = par(text, index)
    if res.status:
      return p.Value.success(res.index, res.value)
    else:
      return p.Value.success(res.index, default_value)
  return optional_parser

def binop(name):
  @p.generate
  def parser():
    yield lexeme(p.string(name))
    return lambda a, b: BinOp(name, a, b)
  return parser

def try_all(parser, *parsers):
  for par in parsers:
    parser = p.try_choice(parser, par)
  return parser

def between(o, par, c):
  @p.generate
  def parser():
    yield lexeme(p.string(o))
    result = yield par
    yield lexeme(p.string(c))
    return result
  return parser

def chainl1(term, f, init=None):
  @p.generate
  def parser():
    if init is None:
      a = yield term
    else:
      a = init
    while True:
      fun = yield optional(f)
      if fun is None:
        break
      else:
        v = yield term
        a = fun(a, v)
    return a
  return parser

@lexeme
@p.generate
def parse_roll():
  count = yield integer
  yield lexeme(p.regex("d", re.I))
  faces = yield integer
  return BasicRoll(count, faces)

@lexeme
@p.generate
def parse_take_modifier():
  yield lexeme(p.regex("t", re.I))
  amount = yield integer
  return TakeModifier(amount)

@p.generate
def parse_modifiers():
  mod = yield parse_take_modifier # | other_modifier
  return mod

@lexeme
@p.generate
def parse_modified_roll():
  roll = yield parse_roll
  mods = yield p.many(parse_modifiers)
  for mod in mods:
    roll = ModifiedRoll(roll, mod)
  return roll

rollval = parse_modified_roll.parsecmap(RollValue)
intval = integer.parsecmap(IntValue)

@lexeme
@p.generate
def term():
  t = yield try_all(rollval, intval, between("(", expression, ")"))
  return t

@lexeme
@p.generate
def expression():
  yield whitespace
  roll = yield rollval
  expr = yield chainl1(term, binop("+") | binop("-"), roll)
  return expr


##################
### Evaluation ###
##################


def generate_roll(count, sides):
  return sorted([random.randint(1, sides+1) for _ in range(count)])

def evaluate_modifier(modifier, roll):
  return modifier.match(
    lambda take: roll[:take.amount]
  )

def evaluate_roll(roll):
  return roll.match(
    lambda basic: generate_roll(roll.count, roll.sides),
    lambda mod: evaluate_modifier(mod.modifier, evaluate_roll(mod.roll))
  )

def evaluate_binop(op, left, right):
  if op == "+":
    return left + right
  elif op == "-":
    return left - right
  else:
    raise ValueError(f"Invalid op: {roll.op}")

def evaluate(roll):
  return roll.match(
    lambda roll: sum(evaluate_roll(roll.roll)),
    lambda intv: intv.value,
    lambda binop: evaluate_binop(
      roll.op,
      evaluate(binop.left),
      evaluate(binop.right))
  )


###############
### Example ###
###############


expr = expression.parse_strict("  10d3 t 5 + 1 d 3 - (4d3 + 8)  ")
print(expr)
print(evaluate(expr))