roll.py 30 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845
  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 pyparsing import ParserElement, Token, Regex, oneOf, Optional, Group, Combine, Literal, CaselessLiteral, ZeroOrMore, StringStart, StringEnd, opAssoc, infixNotation, ParseException, Empty, pyparsing_common, ParseResults, White, Suppress
  12. from typing import Union, List, Any, Tuple, Sequence, 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. sysrand = SystemRandom()
  35. randint = sysrand.randint
  36. def int_limit_converter(x: Any) -> OptionalType[int]:
  37. if x is None:
  38. return None
  39. else:
  40. return int(x)
  41. @attr.s
  42. class IntegerValidator(object):
  43. min_val: OptionalType[int] = attr.ib(default = None, converter = int_limit_converter)
  44. max_val: OptionalType[int] = attr.ib(default = None, converter = int_limit_converter)
  45. handle_float: str = attr.ib(default = 'exception')
  46. @handle_float.validator
  47. def validate_handle_float(self, attribute, value):
  48. assert value in ('exception', 'truncate', 'round')
  49. value_name: str = attr.ib(default = "value")
  50. def __call__(self, value: Any) -> int:
  51. xf: float
  52. x: int
  53. try:
  54. xf = float(value)
  55. except ValueError:
  56. raise ValueError('{} {} does not look like a number'.format(self.value_name, value))
  57. if not xf.is_integer():
  58. if self.handle_float == 'exception':
  59. raise ValueError('{} {} is not an integer'.format(self.value_name, value))
  60. elif self.handle_float == 'truncate':
  61. x = int(xf)
  62. else:
  63. x = round(xf)
  64. else:
  65. x = int(xf)
  66. if self.min_val is not None and x < self.min_val:
  67. raise ValueError('{} {} is too small; must be at least {}'.format(self.value_name, value, self.min_val))
  68. if self.max_val is not None and x > self.max_val:
  69. raise ValueError('{} {} is too large; must be at most {}'.format(self.value_name, value, self.max_val))
  70. return x
  71. die_face_num_validator = IntegerValidator(
  72. min_val = 2, handle_float = 'exception',
  73. value_name = 'die type',
  74. )
  75. DieFaceType = Union[int, str]
  76. def is_fate_face(x: DieFaceType) -> bool:
  77. if isinstance(x, int):
  78. return False
  79. else:
  80. x = str(x).upper()
  81. return x in ('F', 'F.1', 'F.2')
  82. def normalize_die_type(x: DieFaceType) -> DieFaceType:
  83. if is_fate_face(x):
  84. return str(x).upper()
  85. elif x == '%':
  86. return 100
  87. else:
  88. return die_face_num_validator(x)
  89. dice_count_validator = IntegerValidator(
  90. min_val = 1, handle_float = 'exception',
  91. value_name = 'dice count'
  92. )
  93. # Just a named function wrapper for dice_count_validator
  94. def normalize_dice_count(x: Any) -> int:
  95. return dice_count_validator(x)
  96. def ImplicitToken(x) -> ParserElement:
  97. '''Like pyparsing.Empty, but yields one or more tokens instead of nothing.'''
  98. return Empty().setParseAction(lambda toks: x)
  99. # TODO: Look at http://infohost.nmt.edu/tcc/help/pubs/pyparsing/web/classes.html#class-ParserElement
  100. # Implementing the syntax described here: https://www.critdice.com/roll-advanced-dice
  101. # https://stackoverflow.com/a/23956778/125921
  102. # https://stackoverflow.com/a/46583691/125921
  103. var_name: ParserElement = pyparsing_common.identifier.copy().setResultsName('varname')
  104. real_num: ParserElement = pyparsing_common.fnumber.copy()
  105. positive_int: ParserElement = pyparsing_common.integer.copy().setParseAction(lambda toks: [ IntegerValidator(min_val=1)(toks[0]) ])
  106. drop_type: ParserElement = oneOf('K k X x -H -L')
  107. drop_spec: ParserElement = Group(
  108. drop_type.setResultsName('type') +
  109. (positive_int | ImplicitToken(1)).setResultsName('count')
  110. ).setResultsName('drop')
  111. pos_int_implicit_one: ParserElement = (positive_int | ImplicitToken(1))
  112. comparator_type: ParserElement = oneOf('<= < >= > ≤ ≥ =')
  113. reroll_type: ParserElement = Combine(oneOf('R r') ^ ( oneOf('! !!') + Optional('p')))
  114. reroll_spec: ParserElement = Group(
  115. reroll_type.setResultsName('type') +
  116. Optional(
  117. (comparator_type | ImplicitToken('=')).setResultsName('operator') + \
  118. positive_int.setResultsName('value')
  119. )
  120. ).setResultsName('reroll')
  121. count_spec: ParserElement = Group(
  122. Group(
  123. comparator_type.setResultsName('operator') + \
  124. positive_int.setResultsName('value')
  125. ).setResultsName('success_condition') +
  126. Optional(
  127. Literal('f') +
  128. Group(
  129. comparator_type.setResultsName('operator') + \
  130. positive_int.setResultsName('value')
  131. ).setResultsName('failure_condition')
  132. )
  133. ).setResultsName('count_successes')
  134. roll_spec: ParserElement = Group(
  135. (positive_int | ImplicitToken(1)).setResultsName('dice_count') +
  136. CaselessLiteral('d') +
  137. (positive_int | oneOf('% F F.1 F.2')).setResultsName('die_type') +
  138. Optional(reroll_spec ^ drop_spec) +
  139. Optional(count_spec)
  140. ).setResultsName('roll')
  141. expr_parser: ParserElement = infixNotation(
  142. baseExpr=(roll_spec ^ positive_int ^ real_num ^ var_name),
  143. opList=[
  144. (oneOf('** ^').setResultsName('operator', True), 2, opAssoc.RIGHT),
  145. (oneOf('* / × ÷').setResultsName('operator', True), 2, opAssoc.LEFT),
  146. (oneOf('+ -').setResultsName('operator', True), 2, opAssoc.LEFT),
  147. ]
  148. ).setResultsName('expr')
  149. assignment_parser: ParserElement = var_name + Literal('=').setResultsName('assignment') + expr_parser
  150. def roll_die(sides: DieFaceType = 6) -> int:
  151. '''Roll a single die.
  152. Supports any valid integer number of sides as well as 'F' for a fate
  153. die, which can return -1, 0, or 1 with equal probability.
  154. '''
  155. if sides in ('F', 'F.2'):
  156. # Fate die = 1d3-2
  157. return roll_die(3) - 2
  158. elif sides == 'F.1':
  159. d6 = roll_die(6)
  160. if d6 == 1:
  161. return -1
  162. elif d6 == 6:
  163. return 1
  164. else:
  165. return 0
  166. else:
  167. return randint(1, int(sides))
  168. class DieRolled(int):
  169. '''Subclass of int that allows a string suffix.
  170. This is meant for recording the result of rolling a die. The
  171. suffix is purely cosmetic, for the purposes of string conversion.
  172. It can be used to indicate a die roll that has been re-rolled or
  173. exploded, or to indicate a critical hit/miss.
  174. '''
  175. formatter: str
  176. def __new__(cls: type, value: int, formatter: str = '{}') -> 'DieRolled':
  177. # https://github.com/python/typeshed/issues/2686
  178. newval = super(DieRolled, cls).__new__(cls, value) # type: ignore
  179. newval.formatter = formatter
  180. return newval
  181. def __str__(self) -> str:
  182. return self.formatter.format(super().__str__())
  183. def __repr__(self) -> str:
  184. if self.formatter != '{}':
  185. return 'DieRolled(value={value!r}, formatter={formatter!r})'.format(
  186. value=int(self),
  187. formatter=self.formatter,
  188. )
  189. else:
  190. return 'DieRolled({value!r})'.format(value=int(self))
  191. def normalize_dice_roll_list(value: List[Any]) -> List[int]:
  192. result = []
  193. for x in value:
  194. if isinstance(x, int):
  195. result.append(x)
  196. else:
  197. result.append(int(x))
  198. return result
  199. def format_dice_roll_list(rolls: List[int], always_list: bool = False) -> str:
  200. if len(rolls) == 0:
  201. raise ValueError('Need at least one die rolled')
  202. elif len(rolls) == 1 and not always_list:
  203. return color(str(rolls[0]), DETAIL_COLOR)
  204. else:
  205. return '[' + color(" ".join(map(str, rolls)), DETAIL_COLOR) + ']'
  206. def int_or_none(x: OptionalType[Any]) -> OptionalType[int]:
  207. if x is None:
  208. return None
  209. else:
  210. return int(x)
  211. @attr.s
  212. class DiceRolled(object):
  213. '''Class representing the result of rolling one or more similar dice.'''
  214. dice_results: List[int] = attr.ib(converter = normalize_dice_roll_list)
  215. @dice_results.validator
  216. def validate_dice_results(self, attribute, value):
  217. if len(value) == 0:
  218. raise ValueError('Need at least one non-dropped roll')
  219. dropped_results: List[int] = attr.ib(
  220. default = attr.Factory(list),
  221. converter = normalize_dice_roll_list)
  222. roll_desc: str = attr.ib(default = '', converter = str)
  223. success_count: OptionalType[int] = attr.ib(default = None, converter = int_or_none)
  224. def total(self) -> int:
  225. if self.success_count is not None:
  226. return int(self.success_count)
  227. else:
  228. return sum(self.dice_results)
  229. def __str__(self) -> str:
  230. if self.roll_desc:
  231. prefix = '{roll} rolled'.format(roll=color(self.roll_desc, EXPR_COLOR))
  232. else:
  233. prefix = 'Rolled'
  234. if self.dropped_results:
  235. drop = ' (dropped {dropped})'.format(dropped = format_dice_roll_list(self.dropped_results))
  236. else:
  237. drop = ''
  238. if self.success_count is not None:
  239. tot = ', Total successes: ' + color(str(self.total()), DETAIL_COLOR)
  240. elif len(self.dice_results) > 1:
  241. tot = ', Total: ' + color(str(self.total()), DETAIL_COLOR)
  242. else:
  243. tot = ''
  244. return '{prefix}: {results}{drop}{tot}'.format(
  245. prefix=prefix,
  246. results=format_dice_roll_list(self.dice_results),
  247. drop=drop,
  248. tot=tot,
  249. )
  250. def __int__(self) -> int:
  251. return self.total()
  252. def __float__(self) -> float:
  253. return float(self.total())
  254. def validate_by_parser(parser):
  255. '''Return a validator that validates anything parser can parse.'''
  256. def private_validator(instance, attribute, value):
  257. parser.parseString(str(value), True)
  258. return private_validator
  259. @attr.s
  260. class Comparator(object):
  261. cmp_dict = {
  262. '<=': operator.le,
  263. '<': operator.lt,
  264. '>=': operator.ge,
  265. '>': operator.gt,
  266. '≤': operator.le,
  267. '≥': operator.ge,
  268. '=': operator.eq,
  269. }
  270. operator: str = attr.ib(converter = str,
  271. validator = validate_by_parser(comparator_type))
  272. value: int = attr.ib(converter = int,
  273. validator = validate_by_parser(positive_int))
  274. def __str__(self) -> str:
  275. return '{op}{val}'.format(op=self.operator, val=self.value)
  276. def compare(self, x) -> bool:
  277. '''Return True if x satisfies the comparator.
  278. In other words, x is placed on the left-hand side of the
  279. comparison, the Comparator's value is placed on the right hand
  280. side, and the truth value of the resulting test is returned.
  281. '''
  282. return self.cmp_dict[self.operator](x, self.value)
  283. @attr.s
  284. class RerollSpec(object):
  285. # Yes, it has to be called type
  286. type: str = attr.ib(converter = str, validator=validate_by_parser(reroll_type))
  287. operator: OptionalType[str] = attr.ib(default = None)
  288. value: OptionalType[int] = attr.ib(default = None)
  289. def __attrs_post_init__(self):
  290. if (self.operator is None) != (self.value is None):
  291. raise ValueError('Operator and value must be provided together')
  292. def __str__(self) -> str:
  293. result = self.type
  294. if self.operator is not None:
  295. result += self.operator + str(self.value)
  296. return result
  297. def roll_die(self, sides: DieFaceType) -> List[int]:
  298. '''Roll a single die, following specified re-rolling rules.
  299. Returns a list of rolls, since some types of re-rolling
  300. collect the result of multiple die rolls.
  301. '''
  302. if is_fate_face(sides):
  303. raise ValueError("Re-rolling/exploding is incompatible with Fate dice")
  304. sides = int(sides)
  305. cmpr: Comparator
  306. if self.value is None:
  307. if self.type in ('R', 'r'):
  308. cmpr = Comparator('=', 1)
  309. else:
  310. cmpr = Comparator('=', sides)
  311. else:
  312. cmpr = Comparator(self.operator, self.value)
  313. if self.type == 'r':
  314. # Single reroll
  315. roll = roll_die(sides)
  316. if cmpr.compare(roll):
  317. roll = DieRolled(roll_die(sides), '{}' + self.type)
  318. return [ roll ]
  319. elif self.type == 'R':
  320. # Indefinite reroll
  321. roll = roll_die(sides)
  322. while cmpr.compare(roll):
  323. roll = DieRolled(roll_die(sides), '{}' + self.type)
  324. return [ roll ]
  325. elif self.type in ['!', '!!', '!p', '!!p']:
  326. # Explode/penetrate/compound
  327. all_rolls: List[int] = [ roll_die(sides) ]
  328. while cmpr.compare(all_rolls[-1]):
  329. all_rolls.append(roll_die(sides))
  330. # If we never re-rolled, no need to do anything special
  331. if len(all_rolls) == 1:
  332. return all_rolls
  333. # For penetration, subtract 1 from all rolls except the first
  334. if self.type.endswith('p'):
  335. for i in range(1, len(all_rolls)):
  336. all_rolls[i] -= 1
  337. # For compounding, return the sum
  338. if self.type.startswith('!!'):
  339. total = sum(all_rolls)
  340. return [ DieRolled(total, '{}' + self.type) ]
  341. else:
  342. for i in range(0, len(all_rolls)-1):
  343. all_rolls[i] = DieRolled(all_rolls[i], '{}' + self.type)
  344. return all_rolls
  345. else:
  346. raise Exception('Unknown reroll type: {}'.format(self.type))
  347. @attr.s
  348. class DropSpec(object):
  349. # Yes, it has to be called type
  350. type: str = attr.ib(converter = str, validator=validate_by_parser(drop_type))
  351. count: int = attr.ib(default = 1, converter = int, validator=validate_by_parser(positive_int))
  352. def __str__(self) -> str:
  353. if self.count > 1:
  354. return self.type + str(self.count)
  355. else:
  356. return self.type
  357. def drop_rolls(self, rolls: List[int]) -> Tuple[List[int], List[int]]:
  358. '''Drop the appripriate rolls from a list of rolls.
  359. Returns a 2-tuple of roll lists. The first list is the kept
  360. rolls, and the second list is the dropped rolls.
  361. The order of the rolls is not preserved. (TODO FIX THIS)
  362. '''
  363. if not isinstance(rolls, list):
  364. rolls = list(rolls)
  365. keeping = self.type in ('K', 'k')
  366. if keeping:
  367. num_to_keep = self.count
  368. else:
  369. num_to_keep = len(rolls) - self.count
  370. if num_to_keep == 0:
  371. raise ValueError('Not enough rolls: would drop all rolls')
  372. elif num_to_keep == len(rolls):
  373. raise ValueError('Keeping too many rolls: would not drop any rolls')
  374. rolls.sort()
  375. if self.type in ('K', 'X', '-H'):
  376. rolls.reverse()
  377. (head, tail) = rolls[:self.count], rolls[self.count:]
  378. if keeping:
  379. (kept, dropped) = (head, tail)
  380. else:
  381. (kept, dropped) = (tail, head)
  382. return (kept, dropped)
  383. @attr.s
  384. class DiceRoller(object):
  385. die_type: DieFaceType = attr.ib(converter = normalize_die_type)
  386. dice_count: int = attr.ib(default = 1, converter = normalize_dice_count)
  387. reroll_spec: OptionalType[RerollSpec] = attr.ib(default = None)
  388. @reroll_spec.validator
  389. def validate_reroll_spec(self, attribute, value):
  390. if value is not None:
  391. assert isinstance(value, RerollSpec)
  392. drop_spec: OptionalType[DropSpec] = attr.ib(default = None)
  393. @drop_spec.validator
  394. def validate_drop_spec(self, attribute, value):
  395. if value is not None:
  396. assert isinstance(value, DropSpec)
  397. success_comparator: OptionalType[Comparator] = attr.ib(default = None)
  398. failure_comparator: OptionalType[Comparator] = attr.ib(default = None)
  399. @success_comparator.validator
  400. @failure_comparator.validator
  401. def validate_comparator(self, attribute, value):
  402. if value is not None:
  403. assert isinstance(value, Comparator)
  404. def __attrs_post_init__(self):
  405. if self.reroll_spec is not None and self.drop_spec is not None:
  406. raise ValueError('Reroll and drop specs are mutually exclusive')
  407. if self.success_comparator is None and self.failure_comparator is not None:
  408. raise ValueError('Cannot use a failure condition without a success condition')
  409. def __str__(self) -> str:
  410. return '{count}d{type}{reroll}{drop}{success}{fail}'.format(
  411. count = self.dice_count if self.dice_count > 1 else '',
  412. type = self.die_type,
  413. reroll = self.reroll_spec or '',
  414. drop = self.drop_spec or '',
  415. success = self.success_comparator or '',
  416. fail = ('f' + str(self.failure_comparator)) if self.failure_comparator else '',
  417. )
  418. def roll(self) -> DiceRolled:
  419. '''Roll dice according to specifications. Returns a DiceRolled object.'''
  420. all_rolls = []
  421. if self.reroll_spec:
  422. for i in range(self.dice_count):
  423. all_rolls.extend(self.reroll_spec.roll_die(self.die_type))
  424. else:
  425. for i in range(self.dice_count):
  426. all_rolls.append(roll_die(self.die_type))
  427. if self.drop_spec:
  428. (dice_results, dropped_results) = self.drop_spec.drop_rolls(all_rolls)
  429. else:
  430. (dice_results, dropped_results) = (all_rolls, [])
  431. success_count: OptionalType[int]
  432. if self.success_comparator is not None:
  433. success_count = 0
  434. for roll in dice_results:
  435. if self.success_comparator.compare(roll):
  436. success_count += 1
  437. if self.failure_comparator is not None:
  438. for roll in dice_results:
  439. if self.failure_comparator.compare(roll):
  440. success_count -= 1
  441. else:
  442. success_count = None
  443. return DiceRolled(
  444. dice_results=dice_results,
  445. dropped_results=dropped_results,
  446. roll_desc=str(self),
  447. success_count=success_count,
  448. )
  449. def make_dice_roller(expr: Union[str,ParseResults]) -> DiceRoller:
  450. if isinstance(expr, str):
  451. expr = roll_spec.parseString(expr, True)['roll']
  452. assert expr.getName() == 'roll'
  453. expr = expr.asDict()
  454. dtype = normalize_die_type(expr['die_type'])
  455. dcount = normalize_dice_count(expr['dice_count'])
  456. constructor_args: Dict[str, Any] = {
  457. 'die_type': dtype,
  458. 'dice_count': dcount,
  459. 'reroll_spec': None,
  460. 'drop_spec': None,
  461. 'success_comparator': None,
  462. 'failure_comparator': None,
  463. }
  464. rrdict = None
  465. if 'reroll' in expr:
  466. rrdict = expr['reroll']
  467. constructor_args['reroll_spec'] = RerollSpec(**rrdict)
  468. if 'drop' in expr:
  469. ddict = expr['drop']
  470. constructor_args['drop_spec'] = DropSpec(**ddict)
  471. if 'count_successes' in expr:
  472. csdict = expr['count_successes']
  473. constructor_args['success_comparator'] = Comparator(**csdict['success_condition'])
  474. if 'failure_condition' in csdict:
  475. constructor_args['failure_comparator'] = Comparator(**csdict['failure_condition'])
  476. return DiceRoller(**constructor_args)
  477. # examples = [
  478. # '1+1',
  479. # '1 + 1 + x',
  480. # '3d8',
  481. # '2e3 * 4d6 + 2',
  482. # '2d20k',
  483. # '3d20x2',
  484. # '4d4rK3',
  485. # '4d4R4',
  486. # '4d4R>=3',
  487. # '4d4r>=3',
  488. # '4d4!1',
  489. # '4d4!<3',
  490. # '4d4!p',
  491. # '2D20K+10',
  492. # '2D20k+10',
  493. # '10d6X4',
  494. # '4d8r + 6',
  495. # '20d6R≤2',
  496. # '6d10!≥8+6',
  497. # '10d4!p',
  498. # '20d6≥6',
  499. # '8d12≥10f≤2',
  500. # # '4d20R<=2!>=19Xx21>=20f<=5*2+3', # Pretty much every possible feature
  501. # ]
  502. # example_results = {}
  503. # for x in examples:
  504. # try:
  505. # example_results[x] = parse_roll(x)
  506. # except ParseException as ex:
  507. # example_results[x] = ex
  508. # example_results
  509. # rs = RerollSpec('!!p', '=', 6)
  510. # rs.roll_die(6)
  511. # ds = DropSpec('K', 2)
  512. # ds.drop_rolls([1,2,3,4,5])
  513. # ds = DropSpec('x', 2)
  514. # ds.drop_rolls([1,2,3,4,5])
  515. # parse_roll = lambda x: expr_parser.parseString(x)[0]
  516. # exprstring = 'x + 1 + (2 + (3 + 4))'
  517. # expr = parse_roll(exprstring)
  518. # r = parse_roll('x + 1 - 2 * y * 4d4 + 2d20K1>=20f<=5')[0]
  519. op_dict: Dict[str, Callable] = {
  520. '+': operator.add,
  521. '-': operator.sub,
  522. '*': operator.mul,
  523. '×': operator.mul,
  524. '/': operator.truediv,
  525. '÷': operator.truediv,
  526. '**': operator.pow,
  527. '^': operator.pow,
  528. }
  529. NumericType = Union[float,int]
  530. ExprType = Union[NumericType, str, ParseResults]
  531. def normalize_expr(expr: ExprType) -> ParseResults:
  532. if isinstance(expr, str):
  533. return expr_parser.parseString(expr)['expr']
  534. elif isinstance(expr, Number):
  535. return expr
  536. else:
  537. assert isinstance(expr, ParseResults)
  538. return expr['expr']
  539. def _eval_expr_internal(
  540. expr: ExprType,
  541. env: Dict[str, str] = {},
  542. print_rolls: bool = True,
  543. recursed_vars: Set[str] = set()) -> NumericType:
  544. if isinstance(expr, float) or isinstance(expr, int):
  545. # Numeric literal
  546. return expr
  547. elif isinstance(expr, str):
  548. # variable name
  549. if expr in recursed_vars:
  550. raise ValueError('Recursive variable definition detected for {!r}'.format(expr))
  551. elif expr in env:
  552. var_value = env[expr]
  553. parsed = normalize_expr(var_value)
  554. return _eval_expr_internal(parsed, env, print_rolls,
  555. recursed_vars = recursed_vars.union([expr]))
  556. else:
  557. raise ValueError('Expression referenced undefined variable {!r}'.format(expr))
  558. else:
  559. assert isinstance(expr, ParseResults)
  560. if 'operator' in expr:
  561. # Compound expression
  562. operands = expr[::2]
  563. operators = expr[1::2]
  564. assert len(operands) == len(operators) + 1
  565. values = [ _eval_expr_internal(x, env, print_rolls, recursed_vars)
  566. for x in operands ]
  567. result = values[0]
  568. for (op, nextval) in zip(operators, values[1:]):
  569. opfun = op_dict[op]
  570. result = opfun(result, nextval)
  571. return result
  572. else:
  573. # roll specification
  574. roller = make_dice_roller(expr)
  575. rolled = roller.roll()
  576. if print_rolls:
  577. print(rolled)
  578. return int(rolled)
  579. def eval_expr(expr: ExprType,
  580. env: Dict[str,str] = {},
  581. print_rolls: bool = True) -> NumericType:
  582. expr = normalize_expr(expr)
  583. return _eval_expr_internal(expr, env, print_rolls)
  584. def _expr_as_str_internal(expr: ExprType,
  585. env: Dict[str,str] = {},
  586. recursed_vars: Set[str] = set()) -> str:
  587. if isinstance(expr, float) or isinstance(expr, int):
  588. return '{:g}'.format(expr)
  589. elif isinstance(expr, str):
  590. # variable name
  591. if expr in recursed_vars:
  592. raise ValueError('Recursive variable definition detected for {!r}'.format(expr))
  593. elif expr in env:
  594. var_value = env[expr]
  595. parsed = normalize_expr(var_value)
  596. return _expr_as_str_internal(parsed, env, recursed_vars = recursed_vars.union([expr]))
  597. # Not a variable name, just a string
  598. else:
  599. return expr
  600. else:
  601. assert isinstance(expr, ParseResults)
  602. if 'operator' in expr:
  603. # Compound expression
  604. operands = expr[::2]
  605. operators = expr[1::2]
  606. assert len(operands) == len(operators) + 1
  607. values = [ _expr_as_str_internal(x, env, recursed_vars)
  608. for x in operands ]
  609. result = str(values[0])
  610. for (op, nextval) in zip(operators, values[1:]):
  611. result += ' {} {}'.format(op, nextval)
  612. return '(' + result + ')'
  613. else:
  614. # roll specification
  615. return str(make_dice_roller(expr))
  616. def expr_as_str(expr: ExprType, env: Dict[str,str] = {}) -> str:
  617. expr = normalize_expr(expr)
  618. expr = _expr_as_str_internal(expr, env)
  619. if expr.startswith('(') and expr.endswith(')'):
  620. expr = expr[1:-1]
  621. return expr
  622. def read_roll(handle: TextIO = sys.stdin) -> str:
  623. if handle == sys.stdin:
  624. return input("Enter roll> ")
  625. else:
  626. return handle.readline()[:-1]
  627. special_command_parser: ParserElement = (
  628. oneOf('h help ?').setResultsName('help') |
  629. oneOf('q quit exit').setResultsName('quit') |
  630. oneOf('v vars').setResultsName('vars') |
  631. (oneOf('d del delete').setResultsName('delete').leaveWhitespace() + Suppress(White()) + var_name)
  632. )
  633. def var_name_allowed(vname: str) -> bool:
  634. '''Disallow variable names like 'help' and 'quit'.'''
  635. parsers = [ special_command_parser, roll_spec ]
  636. for parser in [ special_command_parser, roll_spec ]:
  637. try:
  638. parser.parseString(vname, True)
  639. return False
  640. except ParseException:
  641. pass
  642. # If the variable name didn't parse as anything else, it's valid
  643. return True
  644. line_parser: ParserElement = (
  645. special_command_parser ^
  646. (assignment_parser | expr_parser)
  647. )
  648. def print_interactive_help() -> None:
  649. print('\n' + '''
  650. To make a roll, type in the roll in dice notation, e.g. '4d4 + 4'.
  651. Nearly all dice notation forms listed in the following references should be supported:
  652. - http://rpg.greenimp.co.uk/dice-roller/
  653. - https://www.critdice.com/roll-advanced-dice
  654. Expressions can include numeric constants, addition, subtraction,
  655. multiplication, division, and exponentiation.
  656. To assign a variable, use 'VAR = VALUE'. For example 'health_potion =
  657. 2d4+2'. Subsequent roll expressions (and other variables) can refer to
  658. this variable, whose value will be substituted in to the expression.
  659. If a variable's value includes any dice rolls, those dice will be
  660. rolled (and produce a different value) every time the variable is
  661. used.
  662. Special commands:
  663. - To show the values of all currently assigned variables, type 'vars'.
  664. - To delete a previously defined variable, type 'del VAR'.
  665. - To show this help text, type 'help'.
  666. - To quit, type 'quit'.
  667. '''.strip() + '\n', file=sys.stdout)
  668. def print_vars(env: Dict[str,str]) -> None:
  669. if len(env):
  670. print('Currently defined variables:')
  671. for k in sorted(env.keys()):
  672. print('{} = {!r}'.format(k, env[k]))
  673. else:
  674. print('No variables are currently defined.')
  675. if __name__ == '__main__':
  676. expr_string = " ".join(sys.argv[1:])
  677. if re.search("\\S", expr_string):
  678. try:
  679. # Note: using expr_parser instead of line_parser, because
  680. # on the command line only roll expressions are valid.
  681. expr = expr_parser.parseString(expr_string, True)
  682. result = eval_expr(expr)
  683. print('Result: {result} (rolled {expr})'.format(
  684. expr=color(expr_as_str(expr), EXPR_COLOR),
  685. result=color("{:g}".format(result), RESULT_COLOR),
  686. ))
  687. except Exception as exc:
  688. logger.error("Error while rolling: %s", repr(exc))
  689. raise exc
  690. sys.exit(1)
  691. else:
  692. env: Dict[str, str] = {}
  693. while True:
  694. try:
  695. expr_string = read_roll()
  696. if not re.search("\\S", expr_string):
  697. continue
  698. parsed = line_parser.parseString(expr_string, True)
  699. if 'help' in parsed:
  700. print_interactive_help()
  701. elif 'quit' in parsed:
  702. logger.info('Quitting.')
  703. break
  704. elif 'vars' in parsed:
  705. print_vars(env)
  706. elif 'delete' in parsed:
  707. vname = parsed['varname']
  708. if vname in env:
  709. print('Deleting saved value for "{var}".'.format(var=color(vname, RESULT_COLOR)))
  710. del env[vname]
  711. else:
  712. logger.error('Variable "{var}" is not defined.'.format(var=color(vname, RESULT_COLOR)))
  713. elif re.search("\\S", expr_string):
  714. if 'assignment' in parsed:
  715. # We have an assignment operation
  716. vname = parsed['varname']
  717. if var_name_allowed(vname):
  718. env[vname] = expr_as_str(parsed['expr'])
  719. print('Saving "{var}" as "{expr}"'.format(
  720. var=color(vname, RESULT_COLOR),
  721. expr=color(env[vname], EXPR_COLOR),
  722. ))
  723. else:
  724. logger.error('You cannot use {!r} as a variable name.'.format(vname))
  725. else:
  726. # Just an expression to evaluate
  727. result = eval_expr(parsed['expr'], env)
  728. print('Result: {result} (rolled {expr})'.format(
  729. expr=color(expr_as_str(parsed, env), EXPR_COLOR),
  730. result=color("{:g}".format(result), RESULT_COLOR),
  731. ))
  732. print('')
  733. except KeyboardInterrupt:
  734. print('')
  735. except EOFError:
  736. print('')
  737. logger.info('Quitting.')
  738. break
  739. except Exception as exc:
  740. logger.error('Error while evaluating {expr!r}:\n{tb}'.format(
  741. expr=expr_string,
  742. tb=traceback.format_exc(),
  743. ))