roll.py 32 KB

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