124 lines
3.4 KiB
Python
124 lines
3.4 KiB
Python
from random import randint
|
|
|
|
from pyparsing import Combine, OpAssoc, ParseResults, Word, infix_notation, nums, one_of
|
|
|
|
SEPARATOR = "|"
|
|
|
|
|
|
class Realizable:
|
|
pass
|
|
|
|
|
|
class Integer(Realizable):
|
|
def __init__(self, tokens):
|
|
self.value = int(tokens[0])
|
|
|
|
def __str__(self):
|
|
return str(self.value)
|
|
|
|
def realize(self):
|
|
return [self.value]
|
|
|
|
|
|
class Dice(Realizable):
|
|
def __init__(self, tokens):
|
|
split_tokens = tokens[0].split(SEPARATOR)
|
|
self.dice, self.faces = int(split_tokens[0]), int(split_tokens[2])
|
|
|
|
def __str__(self):
|
|
return f"{self.dice}d{self.faces}"
|
|
|
|
def realize(self):
|
|
rolls = [randint(1, self.faces) for _ in range(self.dice)]
|
|
return rolls
|
|
|
|
|
|
def roll_parse(roll):
|
|
integer = Word(nums).set_parse_action(Integer)
|
|
dice = Combine(integer + "d" + integer, join_string="|")
|
|
dice.set_parse_action(Dice)
|
|
|
|
expr = infix_notation(
|
|
dice | integer,
|
|
[
|
|
("-", 1, OpAssoc.RIGHT),
|
|
(one_of("* /"), 2, OpAssoc.LEFT),
|
|
(one_of("+ -"), 2, OpAssoc.LEFT),
|
|
],
|
|
)
|
|
|
|
return expr.parse_string(roll, parse_all=True)
|
|
|
|
|
|
def realize(expr):
|
|
if isinstance(expr, ParseResults):
|
|
return list(map(realize, expr))
|
|
elif isinstance(expr, Realizable):
|
|
return expr.realize()
|
|
elif isinstance(expr, str):
|
|
return expr
|
|
else:
|
|
raise ValueError(f"{expr} has an unexpected type in realization")
|
|
|
|
|
|
def pstr_expr(expr):
|
|
if isinstance(expr, Realizable):
|
|
return str(expr)
|
|
elif isinstance(expr, list) and all(map(lambda e: isinstance(e, int), expr)):
|
|
return "/".join(map(str, expr))
|
|
elif isinstance(expr, str):
|
|
return expr
|
|
elif isinstance(expr, list) and len(expr) == 1:
|
|
return pstr_expr(expr[0])
|
|
elif isinstance(expr, ParseResults) or isinstance(expr, list):
|
|
return " ".join(map(pstr_expr, expr))
|
|
else:
|
|
raise ValueError(f"{expr} has an unexpected type in pstr_expr")
|
|
|
|
|
|
def eval_op(op):
|
|
if op == "+":
|
|
return int.__add__
|
|
elif op == "*":
|
|
return int.__mul__
|
|
elif op == "-":
|
|
return int.__sub__
|
|
elif op == "/":
|
|
return int.__floordiv__
|
|
elif op == "%":
|
|
return int.__mod__
|
|
raise ValueError(f"Unknown {op} operator in eval_op")
|
|
|
|
|
|
def compute(rexpr):
|
|
if isinstance(rexpr, Realizable):
|
|
raise ValueError(f"{rexpr} should already be realized")
|
|
elif isinstance(rexpr, int):
|
|
return rexpr
|
|
elif isinstance(rexpr, list) and all(map(lambda e: isinstance(e, int), rexpr)):
|
|
return sum(rexpr)
|
|
elif isinstance(rexpr, list) and len(rexpr) == 1:
|
|
return compute(rexpr[0])
|
|
elif isinstance(rexpr, list):
|
|
if len(rexpr) % 2 == 0:
|
|
raise ValueError(
|
|
f"{rexpr} should be an infix expression with an odd number of elements"
|
|
)
|
|
total = compute(rexpr[0])
|
|
for i in range(1, len(rexpr), 2):
|
|
op = eval_op(rexpr[i])
|
|
total = op(total, compute(rexpr[i + 1]))
|
|
return total
|
|
else:
|
|
raise ValueError(f"{rexpr} has an unexpected type in compute")
|
|
|
|
|
|
async def roll(bot, mtch, room, message):
|
|
expr = roll_parse(" ".join(mtch.args()))
|
|
rexpr = realize(expr)
|
|
value = compute(rexpr)
|
|
await bot.api.send_markdown_message(
|
|
room.room_id,
|
|
"> " + pstr_expr(rexpr) + "\n\n" + str(value),
|
|
reply_to=message.event_id,
|
|
)
|