roll.py 32 KB

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