roll.py 32 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867
  1. #!/usr/bin/env python
  2. import attr
  3. import logging
  4. import re
  5. import sys
  6. import readline
  7. import operator
  8. import traceback
  9. from numbers import Number
  10. from random import SystemRandom
  11. from copy import copy
  12. from arpeggio import ParserPython, RegExMatch, Optional, ZeroOrMore, OneOrMore, OrderedChoice, Sequence, Combine, Not, EOF, PTNodeVisitor, visit_parse_tree, ParseTreeNode, SemanticActionResults
  13. from typing import Union, List, Any, Tuple, Dict, Callable, Set, TextIO
  14. from typing import Optional as OptionalType
  15. try:
  16. import colorama
  17. colorama.init()
  18. from colors import color
  19. except ImportError:
  20. # Fall back to no color
  21. def color(s: str, *args, **kwargs):
  22. '''Fake color function that does nothing.
  23. Used when the colors module cannot be imported.'''
  24. return s
  25. EXPR_COLOR = "green"
  26. RESULT_COLOR = "red"
  27. DETAIL_COLOR = "yellow"
  28. logFormatter = logging.Formatter('%(asctime)s %(levelname)s: %(message)s', '%Y-%m-%d %H:%M:%S')
  29. logger = logging.getLogger(__name__)
  30. logger.setLevel(logging.INFO)
  31. logger.handlers = []
  32. logger.addHandler(logging.StreamHandler())
  33. for handler in logger.handlers:
  34. handler.setFormatter(logFormatter)
  35. sysrand = SystemRandom()
  36. randint = sysrand.randint
  37. # Implementing the syntax described here: https://www.critdice.com/roll-advanced-dice
  38. # https://stackoverflow.com/a/23956778/125921
  39. # Whitespace parsing
  40. def Whitespace(): return RegExMatch(r'\s+')
  41. def OpWS(): return Optional(Whitespace)
  42. # Number parsing
  43. def Digits(): return RegExMatch('[0-9]+')
  44. def NonzeroDigits():
  45. '''Digits with at least one nonzero number.'''
  46. return RegExMatch('0*[1-9][0-9]*')
  47. def Sign(): return ['+', '-']
  48. def Integer(): return Optional(Sign), Digits
  49. def PositiveInteger(): return Optional('+'), Digits
  50. def FloatingPoint():
  51. return (
  52. Optional(Sign),
  53. [
  54. # e.g. '1.', '1.0'
  55. (Digits, '.', Optional(Digits)),
  56. # e.g. '.1'
  57. ('.', Digits),
  58. ]
  59. )
  60. def Scientific():
  61. return ([FloatingPoint, Integer], RegExMatch('[eE]'), Integer)
  62. def Number(): return Combine([Scientific, FloatingPoint, Integer])
  63. def ReservedWord():
  64. '''Matches identifiers that aren't allowed as variable names.'''
  65. command_word_parsers = []
  66. for cmd_type in Command():
  67. cmd_parser = cmd_type()
  68. if isinstance(cmd_parser, tuple):
  69. command_word_parsers.append(cmd_parser[0])
  70. else:
  71. command_word_parsers.append(cmd_parser)
  72. return([
  73. # Starts with a roll expression
  74. RollExpr,
  75. # Matches a command word exactly
  76. (command_word_parsers, [RegExMatch('[^A-Za-z0-9_]'), EOF]),
  77. ])
  78. # Valid variable name parser (should disallow names like 'help', 'quit', or 'd4r')
  79. def Identifier(): return (
  80. Not(ReservedWord),
  81. RegExMatch(r'[A-Za-z_][A-Za-z0-9_]*')
  82. )
  83. def MyNum(): return (
  84. Not('0'),
  85. RegExMatch('[0-9]+'),
  86. )
  87. # Roll expression parsing
  88. def PercentileFace(): return '%'
  89. def DieFace(): return [NonzeroDigits, PercentileFace, RegExMatch(r'F(\.[12])?')]
  90. def BasicRollExpr():
  91. return (
  92. Optional(NonzeroDigits),
  93. RegExMatch('[dD]'),
  94. DieFace,
  95. )
  96. def DropSpec(): return 'K k X x -H -L'.split(' '), Optional(NonzeroDigits)
  97. def CompareOp(): return '<= < >= > ≤ ≥ ='.split(' ')
  98. def Comparison(): return CompareOp, Integer
  99. def RerollType(): return Combine(['r', 'R', ('!', Optional('!'), Optional('p'))])
  100. def RerollSpec():
  101. return (
  102. RerollType,
  103. Optional(
  104. Optional(CompareOp),
  105. Integer,
  106. ),
  107. )
  108. def CountSpec():
  109. return (
  110. Comparison,
  111. Optional('f', Comparison),
  112. )
  113. def RollExpr():
  114. return (
  115. BasicRollExpr,
  116. Optional([DropSpec, RerollSpec]),
  117. Optional(CountSpec),
  118. )
  119. # Arithmetic expression parsing
  120. def PrimaryTerm(): return [RollExpr, Number, Identifier]
  121. def TermOrGroup(): return [PrimaryTerm, ParenExpr]
  122. def Exponent(): return ['**', '^'], OpWS, TermOrGroup
  123. def ExponentExpr(): return TermOrGroup, ZeroOrMore(OpWS, Exponent)
  124. def Mul(): return ['*', '×'], OpWS, ExponentExpr
  125. def Div(): return ['/', '÷'], OpWS, ExponentExpr
  126. def ProductExpr(): return ExponentExpr, ZeroOrMore(OpWS, [Mul, Div])
  127. def Add(): return '+', OpWS, ProductExpr
  128. def Sub(): return '-', OpWS, ProductExpr
  129. def SumExpr(): return ProductExpr, ZeroOrMore(OpWS, [Add, Sub])
  130. def ParenExpr(): return Optional(Sign), '(', OpWS, SumExpr, OpWS, ')'
  131. def Expression():
  132. # Wrapped in a tuple to force a separate entry in the parse tree
  133. return (SumExpr,)
  134. # For parsing vars/expressions without evaluating them. The Combine()
  135. # hides the child nodes from a visitor.
  136. def UnevaluatedExpression(): return Combine(Expression)
  137. def UnevaluatedVariable(): return Combine(Identifier)
  138. # Variable assignment
  139. def VarAssignment(): return (
  140. UnevaluatedVariable,
  141. OpWS, '=', OpWS,
  142. UnevaluatedExpression
  143. )
  144. # Commands
  145. def DeleteCommand(): return (
  146. Combine(['delete', 'del', 'd']),
  147. Whitespace,
  148. UnevaluatedVariable,
  149. )
  150. def HelpCommand(): return Combine(['help', 'h', '?'], Not(Identifier))
  151. def QuitCommand(): return Combine(['quit', 'exit', 'q'], Not(Identifier))
  152. def ListVarsCommand(): return Combine(['variables', 'vars', 'v'], Not(Identifier))
  153. def Command(): return [ ListVarsCommand, DeleteCommand, HelpCommand, QuitCommand, ]
  154. def InputParser(): return Optional([Command, VarAssignment, Expression, Whitespace])
  155. def FullParserPython(language_def, *args, **kwargs):
  156. '''Like ParserPython, but auto-adds EOF to the end of the parser.'''
  157. def TempFullParser(): return (language_def, EOF)
  158. return ParserPython(TempFullParser, *args, **kwargs)
  159. expr_parser = FullParserPython(Expression, skipws = False, memoization = True)
  160. input_parser = FullParserPython(InputParser, skipws = False, memoization = True)
  161. def test_parse(rule, text):
  162. if isinstance(text, str):
  163. return FullParserPython(rule, skipws=False, memoization = True).parse(text)
  164. else:
  165. return [ test_parse(rule, x) for x in text ]
  166. def eval_infix(terms: List[float],
  167. operators: List[Callable[[float,float],float]],
  168. associativity: str = 'l') -> float:
  169. '''Evaluate an infix expression with N terms and N-1 operators.'''
  170. assert associativity in ['l', 'r']
  171. assert len(terms) == len(operators) + 1, 'Need one more term than operator'
  172. if len(terms) == 1:
  173. return terms[0]
  174. elif associativity == 'l':
  175. value = terms[0]
  176. for op, term in zip(operators, terms[1:]):
  177. value = op(value, term)
  178. return value
  179. elif associativity == 'r':
  180. value = terms[-1]
  181. for op, term in zip(reversed(operators), reversed(terms[:-1])):
  182. value = op(term, value)
  183. return value
  184. else:
  185. raise ValueError(f'Invalid associativity: {associativity!r}')
  186. class UndefinedVariableError(KeyError):
  187. pass
  188. def print_vars(env: Dict[str,str]) -> None:
  189. if len(env):
  190. print('Currently defined variables:')
  191. for k in sorted(env.keys()):
  192. print('{} = {!r}'.format(k, env[k]))
  193. else:
  194. print('No variables are currently defined.')
  195. def print_interactive_help() -> None:
  196. print('\n' + '''
  197. To make a roll, type in the roll in dice notation, e.g. '4d4 + 4'.
  198. Nearly all dice notation forms listed in the following references should be supported:
  199. - http://rpg.greenimp.co.uk/dice-roller/
  200. - https://www.critdice.com/roll-advanced-dice
  201. Expressions can include numeric constants, addition, subtraction,
  202. multiplication, division, and exponentiation.
  203. To assign a variable, use 'VAR = VALUE'. For example 'health_potion =
  204. 2d4+2'. Subsequent roll expressions (and other variables) can refer to
  205. this variable, whose value will be substituted in to the expression.
  206. If a variable's value includes any dice rolls, those dice will be
  207. rolled (and produce a different value) every time the variable is
  208. used.
  209. Special commands:
  210. - To show the values of all currently assigned variables, type 'vars'.
  211. - To delete a previously defined variable, type 'del VAR'.
  212. - To show this help text, type 'help'.
  213. - To quit, type 'quit'.
  214. '''.strip() + '\n', file=sys.stdout)
  215. DieFaceType = Union[int, str]
  216. def roll_die(face: DieFaceType = 6) -> int:
  217. '''Roll a single die.
  218. Supports any valid integer number of sides as well as 'F', 'F.1', and
  219. 'F.2' for a Face die, which can return -1, 0, or 1.
  220. '''
  221. if face in ('F', 'F.2'):
  222. # Fate die = 1d3-2
  223. return roll_die(3) - 2
  224. elif face == 'F.1':
  225. d6 = roll_die(6)
  226. if d6 == 1:
  227. return -1
  228. elif d6 == 6:
  229. return 1
  230. else:
  231. return 0
  232. else:
  233. face = int(face)
  234. if face < 2:
  235. raise ValueError(f"Can't roll a {face}-sided die")
  236. return randint(1, face)
  237. def roll_die_with_rerolls(face: int, reroll_condition: Callable, reroll_limit = None) -> List[int]:
  238. '''Roll a single die, and maybe reroll it several times.
  239. After each roll, 'reroll_condition' is called on the result, and
  240. if it returns True, the die is rolled again. All rolls are
  241. collected in a list, and the list is returned as soon as the
  242. condition returns False.
  243. If reroll_limit is provided, it limits the maximum number of
  244. rerolls. Note that the total number of rolls can be one more than
  245. the reroll limit, since the first roll is not considered a reroll.
  246. '''
  247. all_rolls = [roll_die(face)]
  248. while reroll_condition(all_rolls[-1]):
  249. if reroll_limit is not None and len(all_rolls) > reroll_limit:
  250. break
  251. all_rolls.append(roll_die(face))
  252. return all_rolls
  253. class DieRolled(int):
  254. '''Subclass of int that allows an alternate string representation.
  255. This is meant for recording the result of rolling a die. The
  256. formatter argument should include '{}' anywhere that the integer
  257. value should be substituted into the string representation.
  258. (However, it can also override the string representation entirely
  259. by not including '{}'.) The string representation has no effect on
  260. the numeric value. It can be used to indicate a die roll that has
  261. been re-rolled or exploded, or to indicate a critical hit/miss.
  262. '''
  263. formatter: str
  264. def __new__(cls: type, value: int, formatter: str = '{}') -> 'DieRolled':
  265. # https://github.com/python/typeshed/issues/2686
  266. newval = super(DieRolled, cls).__new__(cls, value) # type: ignore
  267. newval.formatter = formatter
  268. return newval
  269. def __str__(self) -> str:
  270. return self.formatter.format(super().__str__())
  271. def __repr__(self) -> str:
  272. if self.formatter != '{}':
  273. return f'DieRolled(value={int(self)!r}, formatter={self.formatter!r})'
  274. else:
  275. return f'DieRolled({int(self)!r})'
  276. def format_dice_roll_list(rolls: List[int], always_list: bool = False) -> str:
  277. if len(rolls) == 0:
  278. raise ValueError('Need at least one die rolled')
  279. elif len(rolls) == 1 and not always_list:
  280. return color(str(rolls[0]), DETAIL_COLOR)
  281. else:
  282. return '[' + color(" ".join(map(str, rolls)), DETAIL_COLOR) + ']'
  283. def int_or_none(x: OptionalType[Any]) -> OptionalType[int]:
  284. if x is None:
  285. return None
  286. else:
  287. return int(x)
  288. def str_or_none(x: OptionalType[Any]) -> OptionalType[str]:
  289. if x is None:
  290. return None
  291. else:
  292. return str(x)
  293. @attr.s
  294. class DiceRolled(object):
  295. '''Class representing the result of rolling one or more similar dice.'''
  296. dice_results: List[int] = attr.ib()
  297. @dice_results.validator
  298. def validate_dice_results(self, attribute, value):
  299. if len(value) == 0:
  300. raise ValueError('Need at least one non-dropped roll')
  301. dropped_results: OptionalType[List[int]] = attr.ib(default = None)
  302. roll_text: OptionalType[str] = attr.ib(
  303. default = None, converter = str_or_none)
  304. success_count: OptionalType[int] = attr.ib(
  305. default = None, converter = int_or_none)
  306. def total(self) -> int:
  307. if self.success_count is not None:
  308. return int(self.success_count)
  309. else:
  310. return sum(self.dice_results)
  311. def __str__(self) -> str:
  312. results = format_dice_roll_list(self.dice_results)
  313. if self.roll_text:
  314. prefix = '{text} rolled'.format(text=color(self.roll_text, EXPR_COLOR))
  315. else:
  316. prefix = 'Rolled'
  317. if self.dropped_results:
  318. drop = ' (dropped {dropped})'.format(dropped = format_dice_roll_list(self.dropped_results))
  319. else:
  320. drop = ''
  321. if self.success_count is not None:
  322. tot = ', Total successes: ' + color(str(self.total()), DETAIL_COLOR)
  323. elif len(self.dice_results) > 1:
  324. tot = ', Total: ' + color(str(self.total()), DETAIL_COLOR)
  325. else:
  326. tot = ''
  327. return f'{prefix}: {results}{drop}{tot}'
  328. def __int__(self) -> int:
  329. return self.total()
  330. def __float__(self) -> float:
  331. return float(self.total())
  332. cmp_dict = {
  333. '<=': operator.le,
  334. '<': operator.lt,
  335. '>=': operator.ge,
  336. '>': operator.gt,
  337. '≤': operator.le,
  338. '≥': operator.ge,
  339. '=': operator.eq,
  340. }
  341. @attr.s
  342. class Comparator(object):
  343. cmp_dict = {
  344. '<=': operator.le,
  345. '<': operator.lt,
  346. '>=': operator.ge,
  347. '>': operator.gt,
  348. '≤': operator.le,
  349. '≥': operator.ge,
  350. '=': operator.eq,
  351. }
  352. operator: str = attr.ib(converter = str)
  353. @operator.validator
  354. def validate_operator(self, attribute, value):
  355. if not value in self.cmp_dict:
  356. raise ValueError(f'Unknown comparison operator: {value!r}')
  357. value: int = attr.ib(converter = int)
  358. def __str__(self) -> str:
  359. return '{op}{val}'.format(op=self.operator, val=self.value)
  360. def compare(self, x: float) -> bool:
  361. '''Return True if x satisfies the comparator.
  362. In other words, x is placed on the left-hand side of the
  363. comparison, the Comparator's value is placed on the right hand
  364. side, and the truth value of the resulting test is returned.
  365. '''
  366. return self.cmp_dict[self.operator](x, self.value)
  367. def __call__(self, x: float) -> bool:
  368. '''Calls Comparator.compare() on x.
  369. This allows the Comparator to be used as a callable.'''
  370. return self.compare(x)
  371. def validate_comparators(face: DieFaceType, *comparators):
  372. '''Validate one or more comparators for a die face type.
  373. This will test every possible die value, making sure that each
  374. test succeeds on at least one value and fails on at least one
  375. value, and it will make sure that at most one test succeeds on any
  376. given value.
  377. '''
  378. if isinstance(face, str):
  379. assert face.startswith('F')
  380. values = range(-1, 2)
  381. else:
  382. values = range(1, face+1)
  383. comparators = list(comparators)
  384. # Validate individual comparators
  385. for comp in comparators:
  386. results = [ comp(val) for val in values ]
  387. if all(results):
  388. raise ValueError(f"Test {str(comp)!r} can never fail for d{face}")
  389. if not any(results):
  390. raise ValueError(f"Test {str(comp)!r} can never succeed for d{face}")
  391. # Check for comparator overlap
  392. for val in values:
  393. passing_comps = [ comp for comp in comparators if comp(val) ]
  394. if len(passing_comps) > 1:
  395. raise ValueError(f"Can't have overlapping tests: {str(passing_comps[0])!r} and {str(passing_comps[1])!r}")
  396. def roll_dice(roll_desc: Dict) -> DiceRolled:
  397. '''Roll dice based on a roll description.
  398. See InputHandler.visit_RollExpr(), which generates roll
  399. descriptions. This function assumes the roll description is
  400. already validated.
  401. Returns a tuple of two lists. The first list is the kept rolls,
  402. and the second list is the dropped rolls.
  403. '''
  404. die_face: DieFaceType = roll_desc['die_face']
  405. dice_count: int = roll_desc['dice_count']
  406. kept_rolls: List[int] = []
  407. dropped_rolls: OptionalType[List[int]] = None
  408. success_count: Optional[int] = None
  409. if 'reroll_type' in roll_desc:
  410. die_face = int(die_face) # No fate dice
  411. reroll_type: str = roll_desc['reroll_type']
  412. reroll_limit = 1 if reroll_type == 'r' else None
  413. reroll_desc: Dict = roll_desc['reroll_desc']
  414. reroll_comparator = Comparator(
  415. operator = reroll_desc['comparator'],
  416. value = reroll_desc['target'],
  417. )
  418. validate_comparators(die_face, reroll_comparator)
  419. for i in range(dice_count):
  420. current_rolls = roll_die_with_rerolls(die_face, reroll_comparator, reroll_limit)
  421. if len(current_rolls) == 1:
  422. # If no rerolls happened, then just add the single
  423. # roll as is.
  424. kept_rolls.append(current_rolls[0])
  425. elif reroll_type in ['r', 'R']:
  426. # Keep only the last roll, and mark it as rerolled
  427. kept_rolls.append(DieRolled(current_rolls[-1], '{}' + reroll_type))
  428. elif reroll_type in ['!', '!!', '!p', '!!p']:
  429. if reroll_type.endswith('p'):
  430. # For penetration, subtract 1 from all rolls
  431. # except the first
  432. for i in range(1, len(current_rolls)):
  433. current_rolls[i] -= 1
  434. if reroll_type.startswith('!!'):
  435. # For compounding, return the sum, marked as a
  436. # compounded roll.
  437. kept_rolls.append(DieRolled(sum(current_rolls),
  438. '{}' + reroll_type))
  439. else:
  440. # For exploding, add each individual roll to the
  441. # list. Mark each roll except the last as
  442. # rerolled.
  443. for i in range(0, len(current_rolls) - 1):
  444. current_rolls[i] = DieRolled(current_rolls[i], '{}' + reroll_type)
  445. kept_rolls.extend(current_rolls)
  446. else:
  447. raise ValueError(f'Unknown reroll type: {reroll_type}')
  448. else:
  449. # Roll the requested number of dice
  450. all_rolls = [ roll_die(die_face) for i in range(dice_count) ]
  451. if 'drop_type' in roll_desc:
  452. keep_count: int = roll_desc['keep_count']
  453. keep_high: bool = roll_desc['keep_high']
  454. # We just need to keep the highest/lowest N rolls. The
  455. # extra complexity here is just to preserve the original
  456. # order of those rolls.
  457. rolls_to_keep = sorted(all_rolls, reverse = keep_high)[:keep_count]
  458. kept_rolls = []
  459. for roll in rolls_to_keep:
  460. kept_rolls.append(all_rolls.pop(all_rolls.index(roll)))
  461. # Remaining rolls are dropped
  462. dropped_rolls = all_rolls
  463. else:
  464. kept_rolls = all_rolls
  465. if 'count_success' in roll_desc:
  466. die_face = int(die_face) # No fate dice
  467. success_desc = roll_desc['count_success']
  468. success_test = Comparator(
  469. operator = success_desc['comparator'],
  470. value = success_desc['target'],
  471. )
  472. # Sanity check: make sure the success test can be met
  473. if not any(map(success_test, range(1, die_face +1))):
  474. raise ValueError(f"Test {str(success_test)!r} can never succeed for d{die_face}")
  475. validate_comparators(die_face, success_test)
  476. success_count = sum(success_test(x) for x in kept_rolls)
  477. if 'count_failure' in roll_desc:
  478. failure_desc = roll_desc['count_failure']
  479. failure_test = Comparator(
  480. operator = failure_desc['comparator'],
  481. value = failure_desc['target'],
  482. )
  483. validate_comparators(die_face, success_test, failure_test)
  484. success_count -= sum(failure_test(x) for x in kept_rolls)
  485. else:
  486. # TODO: Label crits and critfails here
  487. pass
  488. return DiceRolled(
  489. dice_results = kept_rolls,
  490. dropped_results = dropped_rolls,
  491. success_count = success_count,
  492. roll_text = roll_desc['roll_text'],
  493. )
  494. class ExpressionStringifier(PTNodeVisitor):
  495. def __init__(self, **kwargs):
  496. self.env: Dict[str, str] = kwargs.pop('env', {})
  497. self.recursed_vars: Set[str] = kwargs.pop('recursed_vars', set())
  498. self.expr_parser = kwargs.pop('expr_parser', expr_parser)
  499. super().__init__(**kwargs)
  500. def visit__default__(self, node, children):
  501. if children:
  502. return ''.join(children)
  503. else:
  504. return node.value
  505. def visit_Identifier(self, node, children):
  506. '''Interpolate variable.'''
  507. var_name = node.value
  508. if var_name in self.recursed_vars:
  509. raise ValueError(f'Recursive variable definition detected for {var_name!r}')
  510. try:
  511. var_expression = self.env[var_name]
  512. except KeyError as ex:
  513. raise UndefinedVariableError(*ex.args)
  514. recursive_visitor = copy(self)
  515. recursive_visitor.recursed_vars = self.recursed_vars.union([var_name])
  516. return self.expr_parser.parse(var_expression).visit(recursive_visitor)
  517. class QuitRequested(BaseException):
  518. pass
  519. class InputHandler(PTNodeVisitor):
  520. def __init__(self, **kwargs):
  521. self.expr_stringifier = ExpressionStringifier(**kwargs)
  522. self.env: Dict[str, str] = kwargs.pop('env', {})
  523. self.recursed_vars: Set[str] = kwargs.pop('recursed_vars', set())
  524. self.expr_parser = kwargs.pop('expr_parser', expr_parser)
  525. self.print_results = kwargs.pop('print_results', True)
  526. self.print_rolls = kwargs.pop('print_rolls', True)
  527. super().__init__(**kwargs)
  528. def visit_Whitespace(self, node, children):
  529. '''Remove whitespace nodes'''
  530. return None
  531. def visit_Number(self, node, children):
  532. '''Return the numeric value.
  533. Uses int if possible, otherwise float.'''
  534. try:
  535. return int(node.value)
  536. except ValueError:
  537. return float(node.value)
  538. def visit_NonzeroDigits(self, node, children):
  539. return int(node.flat_str())
  540. def visit_Integer(self, node, children):
  541. return int(node.flat_str())
  542. def visit_PercentileFace(self, node, children):
  543. return 100
  544. def visit_BasicRollExpr(self, node, children):
  545. die_face = children[-1]
  546. if isinstance(die_face, int) and die_face < 2:
  547. raise ValueError(f"Invalid roll: Can't roll a {die_face}-sided die")
  548. return {
  549. 'dice_count': children[0] if len(children) == 3 else 1,
  550. 'die_face': die_face,
  551. }
  552. def visit_DropSpec(self, node, children):
  553. return {
  554. 'drop_type': children[0],
  555. 'drop_or_keep_count': children[1] if len(children) > 1 else 1,
  556. }
  557. def visit_RerollSpec(self, node, children):
  558. if len(children) == 1:
  559. return {
  560. 'reroll_type': children[0],
  561. # The default reroll condition depends on other parts
  562. # of the roll expression, so it will be "filled in"
  563. # later.
  564. }
  565. elif len(children) == 2:
  566. return {
  567. 'reroll_type': children[0],
  568. 'reroll_desc': {
  569. 'comparator': '=',
  570. 'target': children[1],
  571. },
  572. }
  573. elif len(children) == 3:
  574. return {
  575. 'reroll_type': children[0],
  576. 'reroll_desc': {
  577. 'comparator': children[1],
  578. 'target': children[2],
  579. },
  580. }
  581. else:
  582. raise ValueError("Invalid reroll specification")
  583. def visit_Comparison(self, node, children):
  584. return {
  585. 'comparator': children[0],
  586. 'target': children[1],
  587. }
  588. def visit_CountSpec(self, node, children):
  589. result = { 'count_success': children[0], }
  590. if len(children) > 1:
  591. result['count_failure'] = children[1]
  592. return result
  593. def visit_RollExpr(self, node, children):
  594. # Collect all child dicts into one
  595. roll_desc = {
  596. 'roll_text': node.flat_str(),
  597. }
  598. for child in children:
  599. roll_desc.update(child)
  600. logger.debug(f'Initial roll description: {roll_desc!r}')
  601. # Perform some validation that can only be done once the
  602. # entire roll description is collected.
  603. if not isinstance(roll_desc['die_face'], int):
  604. if 'reroll_type' in roll_desc:
  605. raise ValueError('Can only reroll/explode numeric dice, not Fate dice')
  606. if 'count_success' in roll_desc:
  607. raise ValueError('Can only count successes on numeric dice, not Fate dice')
  608. # Fill in implicit reroll type
  609. if 'reroll_type' in roll_desc and not 'reroll_desc' in roll_desc:
  610. rrtype = roll_desc['reroll_type']
  611. if rrtype in ['r', 'R']:
  612. roll_desc['reroll_desc'] = {
  613. 'comparator': '=',
  614. 'target': 1,
  615. }
  616. else:
  617. roll_desc['reroll_desc'] = {
  618. 'comparator': '=',
  619. 'target': roll_desc['die_face'],
  620. }
  621. # Validate drop spec and determine exactly how many dice to
  622. # drop/keep
  623. if 'drop_type' in roll_desc:
  624. dtype = roll_desc['drop_type']
  625. keeping = dtype in ['K', 'k']
  626. if keeping:
  627. roll_desc['keep_count'] = roll_desc['drop_or_keep_count']
  628. else:
  629. roll_desc['keep_count'] = roll_desc['dice_count'] - roll_desc['drop_or_keep_count']
  630. if roll_desc['keep_count'] < 1:
  631. drop_count = roll_desc['dice_count'] - roll_desc['keep_count']
  632. raise ValueError(f"Can't drop {drop_count} dice out of {roll_desc['dice_count']}")
  633. if roll_desc['keep_count'] >= roll_desc['dice_count']:
  634. raise ValueError(f"Can't keep {roll_desc['keep_count']} dice out of {roll_desc['dice_count']}")
  635. # Keeping high rolls is the same as dropping low rolls
  636. roll_desc['keep_high'] = dtype in ['K', 'x', '-L']
  637. # Validate count spec
  638. elif 'count_failure' in roll_desc and not 'count_success' in roll_desc:
  639. # The parser shouldn't allow this, but just in case
  640. raise ValueError("Can't have a failure condition without a success condition")
  641. logger.debug(f'Final roll description: {roll_desc!r}')
  642. result = roll_dice(roll_desc)
  643. if self.print_rolls:
  644. print(str(result))
  645. return int(result)
  646. def visit_Identifier(self, node, children):
  647. '''Interpolate variable.'''
  648. var_name = node.value
  649. if var_name in self.recursed_vars:
  650. raise ValueError(f'Recursive variable definition detected for {var_name!r}')
  651. try:
  652. var_expression = self.env[var_name]
  653. except KeyError as ex:
  654. raise UndefinedVariableError(*ex.args)
  655. recursive_visitor = copy(self)
  656. recursive_visitor.recursed_vars = self.recursed_vars.union([var_name])
  657. # Don't print the results of evaluating variables
  658. recursive_visitor.print_results = False
  659. if self.debug:
  660. self.dprint(f'Evaluating variable {var_name} with expression {var_expression!r}')
  661. return self.expr_parser.parse(var_expression).visit(recursive_visitor)
  662. def visit_Expression(self, node, children):
  663. if self.print_results:
  664. expr_full_text = node.visit(self.expr_stringifier)
  665. print('Result: {result} (rolled {expr})'.format(
  666. expr=color(expr_full_text, EXPR_COLOR),
  667. result=color(f'{children[0]:g}', RESULT_COLOR),
  668. ))
  669. return children[0]
  670. # Each of these returns a tuple of (operator, value)
  671. def visit_Add(self, node, children):
  672. return (operator.add, children[-1])
  673. def visit_Sub(self, node, children):
  674. return (operator.sub, children[-1])
  675. def visit_Mul(self, node, children):
  676. return (operator.mul, children[-1])
  677. def visit_Div(self, node, children):
  678. return (operator.truediv, children[-1])
  679. def visit_Exponent(self, node, children):
  680. return (operator.pow, children[-1])
  681. # Each of these receives a first child that is a number and the
  682. # remaining children are tuples of (operator, number)
  683. def visit_SumExpr(self, node, children):
  684. values = [children[0]]
  685. ops = []
  686. for (op, val) in children[1:]:
  687. values.append(val)
  688. ops.append(op)
  689. if self.debug:
  690. self.dprint(f'Sum: values: {values!r}; ops: {ops!r}')
  691. return eval_infix(values, ops, 'l')
  692. def visit_ProductExpr(self, node, children):
  693. values = [children[0]]
  694. ops = []
  695. for (op, val) in children[1:]:
  696. values.append(val)
  697. ops.append(op)
  698. if self.debug:
  699. self.dprint(f'Product: values: {values!r}; ops: {ops!r}')
  700. return eval_infix(values, ops, 'l')
  701. def visit_ExponentExpr(self, node, children):
  702. values = [children[0]]
  703. ops = []
  704. for (op, val) in children[1:]:
  705. values.append(val)
  706. ops.append(op)
  707. if self.debug:
  708. self.dprint(f'Exponent: values: {values!r}; ops: {ops!r}')
  709. return eval_infix(values, ops, 'l')
  710. def visit_Sign(self, node, children):
  711. if node.value == '-':
  712. return -1
  713. else:
  714. return 1
  715. def visit_ParenExpr(self, node, children):
  716. assert len(children) > 0
  717. # Multiply the sign (if present) and the value inside the
  718. # parens
  719. return functools.reduce(operator.mul, children)
  720. def visit_VarAssignment(self, node, children):
  721. logger.debug(f'Doing variable assignment: {node.flat_str()}')
  722. var_name, var_value = children
  723. print('Saving "{var}" as "{expr}"'.format(
  724. var=color(var_name, RESULT_COLOR),
  725. expr=color(var_value, EXPR_COLOR),
  726. ))
  727. self.env[var_name] = var_value
  728. def visit_ListVarsCommand(self, node, children):
  729. print_vars(self.env)
  730. def visit_DeleteCommand(self, node, children):
  731. var_name = children[-1]
  732. print('Deleting saved value for "{var}".'.format(
  733. var=color(var_name, RESULT_COLOR)))
  734. try:
  735. self.env.pop(var_name)
  736. except KeyError as ex:
  737. raise UndefinedVariableError(*ex.args)
  738. def visit_HelpCommand(self, node, children):
  739. print_interactive_help()
  740. def visit_QuitCommand(self, node, children):
  741. raise QuitRequested()
  742. # def handle_input(expr: str, **kwargs) -> float:
  743. # return input_parser.parse(expr).visit(InputHandler(**kwargs))
  744. # handle_input('help')
  745. # handle_input('2+2 * 2 ** 2')
  746. # env = {}
  747. # handle_input('y = 2 + 2', env = env)
  748. # handle_input('x = y + 2', env = env)
  749. # handle_input('2 + x', env = env)
  750. # handle_input('del x', env = env)
  751. # handle_input('vars', env = env)
  752. # handle_input('2 + x', env = env)
  753. # handle_input('d4 = 5', env = env)
  754. def read_input(handle: TextIO = sys.stdin) -> str:
  755. if handle == sys.stdin:
  756. return input("Enter roll> ")
  757. else:
  758. return handle.readline()[:-1]
  759. if __name__ == '__main__':
  760. expr_string = " ".join(sys.argv[1:])
  761. if re.search("\\S", expr_string):
  762. try:
  763. expr_parser.parse(expr_string).visit(InputHandler())
  764. except Exception as exc:
  765. logger.error("Error while rolling: %s", repr(exc))
  766. raise exc
  767. sys.exit(1)
  768. else:
  769. env: Dict[str, str] = {}
  770. handler = InputHandler(env = env)
  771. while True:
  772. try:
  773. input_string = read_input()
  774. input_parser.parse(input_string).visit(handler)
  775. except KeyboardInterrupt:
  776. print('')
  777. except (EOFError, QuitRequested):
  778. print('')
  779. logger.info('Quitting.')
  780. break
  781. except Exception as exc:
  782. logger.error('Error while evaluating {expr!r}:\n{tb}'.format(
  783. expr=expr_string,
  784. tb=traceback.format_exc(),
  785. ))