"""
operator
~~~~~~~~
Operators for handling the dice part of dice notation.
"""
import operator
import random
import secrets
from collections.abc import Callable, Sequence
# Result types for annotation.
Options = tuple[str, str]
Pool = Sequence[int]
# Operation types for annotation.
CompOp = Callable[[int, int], bool]
OptionsOp = Callable[[str, str], Options]
ChoiceOp = Callable[[bool, Options], str]
PoolGenOp = Callable[[int, int], Pool]
DiceOp = Callable[[int, int], int]
MathOp = Callable[[int, int], int]
PoolOp = Callable[[Pool, int], Pool]
PoolDegenOp = Callable[[Pool, int], int]
UPoolDegenOp = Callable[[Pool], int]
Operation = Callable
# Registration.
ops: dict[str, Operation] = {
'^': operator.pow,
'*': operator.mul,
'/': operator.floordiv,
'%': operator.mod,
'+': operator.add,
'-': operator.sub,
'>': operator.gt,
'<': operator.le,
'>=': operator.ge,
'<=': operator.le,
'==': operator.eq,
'!=': operator.ne,
}
class operation:
"""A registration decorator for operations.
:param symbol: The string used to refer to the operation.
:return: None.
:rtype: NoneType
"""
def __init__(self, symbol: str) -> None:
self.symbol = symbol
def __call__(self, fn: Operation) -> Operation:
"""Register the operation.
:param fn: The decorated function. It's sent automatically
when :class:`operator.operation` is used as a decorator.
:return: The decorated :class:`collections.abc.Callable`.
:rtype: Callable
"""
ops[self.symbol] = fn
return fn
# Choice operators.
[docs]
@operation(':')
def choice_options(a: str, b: str) -> Options:
"""Create the options for a choice.
:ref:`YADN` reference: :ref:`choice_options`
:param a: The qualifier for the true condition of a choice.
:param b: The qualifier for the false condition of a choice.
:return: The qualifiers as a :class:`tuple`.
:rtype: tuple
Usage::
>>> choice_options('success', 'failure')
('success', 'failure')
"""
return (a, b)
[docs]
@operation('?')
def choice(boolean: bool, options: Options) -> str:
"""Make a choice.
:ref:`YADN` reference: :ref:`choice`
:param boolean: The decision as a :class:`bool`.
:param options: The two options to pick from.
:return: The chosen option as a :class:`str`.
:rtype: str
Usage::
>>> choice(False, ('spam', 'eggs'))
'eggs'
"""
result = options[0]
if not boolean:
result = options[1]
return result
# Dice operators.
[docs]
@operation('dc')
def concat(num: int, size: int) -> int:
"""Concatenate the least significant digits.
:ref:`YADN` reference: :ref:`concat`
:param num: The number of dice to roll.
:param size: The highest number that can be rolled on a die.
:return: The concatenated least significant digits as an
:class:`int`.
:rtype: int
Usage::
>>> # This line is to ensure predictability for testing.
>>> # Do not use outside of test cases.
>>> _seed('spam')
>>>
>>> # Roll 2d10 as percentile dice. It's not quite right,
>>> # since 00 should be 100, but the physical dice have
>>> # that problem, too.
>>> concat(2, 10)
21
"""
base = 10
pool = dice_pool(num, size)
pool = pool_modulo(pool, base)
return pool_concatenate(pool)
[docs]
@operation('d')
def die(num: int, size: int) -> int:
"""Roll a number of same-sized dice and return the result.
:ref:`YADN` reference: :ref:`die`
:param num: The number of dice to roll.
:param size: The highest number that can be rolled on a die.
:return: The sum of the result of each die as an :class:`int`.
:rtype: int
Usage::
>>> # This line is to ensure predictability for testing.
>>> # Do not use outside of test cases.
>>> _seed('spam')
>>>
>>> # Roll 3d6.
>>> die(3, 6)
5
"""
pool = dice_pool(num, size)
return sum(pool)
[docs]
@operation('d!')
def exploding_die(num: int, size: int) -> int:
"""Roll a number of exploding same-sized dice.
:ref:`YADN` reference: :ref:`explode`
:param num: The number of dice to roll.
:param size: The highest number that can be rolled on a die.
:return: The sum of the result of each die as an :class:`int`.
:rtype: int
Usage::
>>> # This line is to ensure predictability for testing.
>>> # Do not use outside of test cases.
>>> _seed('spam')
>>>
>>> # Roll 5d!6.
>>> exploding_die(5, 6)
15
"""
return sum(exploding_pool(num, size))
[docs]
@operation('dh')
def keep_high_die(num: int, size: int) -> int:
"""Roll a number of dice and keep the highest.
:ref:`YADN` reference: :ref:`keep_high`
:param num: The number of dice to roll.
:param size: The highest number that can be rolled on a die.
:return: The highest value as an :class:`int`.
:rtype: int
Usage::
>>> # This line is to ensure predictability for testing.
>>> # Do not use outside of test cases.
>>> _seed('spam')
>>>
>>> # Roll 5dh6.
>>> keep_high_die(5, 6)
5
"""
pool = dice_pool(num, size)
return max(pool)
[docs]
@operation('dl')
def keep_low_die(num: int, size: int) -> int:
"""Roll a number of dice and keep the lowest.
:ref:`YADN` reference: :ref:`keep_low`
:param num: The number of dice to roll.
:param size: The highest number that can be rolled on a die.
:return: The lowest value as an :class:`int`.
:rtype: int
Usage::
>>> # This line is to ensure predictability for testing.
>>> # Do not use outside of test cases.
>>> _seed('spam')
>>>
>>> # Roll 5dh6.
>>> keep_low_die(5, 6)
1
"""
pool = dice_pool(num, size)
return min(pool)
[docs]
@operation('dw')
def wild_die(num: int, size: int) -> int:
"""Roll a number of same-sized dice and return the result, with
one of the dice being the wild die.
:ref:`YADN` reference: :ref:`wild_die`
:param num: The number of dice to roll.
:param size: The highest number that can be rolled on a die.
:return: The sum of the values as an :class:`int`.
:rtype: int
Usage::
>>> # This line is to ensure predictability for testing.
>>> # Do not use outside of test cases.
>>> _seed('spam')
>>>
>>> # Roll 5dw6.
>>> wild_die(5, 6)
0
"""
wild = exploding_pool(1, size)
regular = dice_pool(num - 1, size)
if wild[0] == 1:
return 0
return sum((sum(wild), sum(regular)))
# Pool operators.
[docs]
@operation('pc')
def pool_cap(pool: Pool, cap: int) -> Pool:
"""Cap the maximum value in a pool.
:ref:`YADN` reference: :ref:`pool_cap`
:param pool: A sequence of die values.
:param cap: The maximum value of a die roll.
:return: The resulting values as a :class:`tuple`.
:rtype: tuple
Usage::
>>> # Roll [4, 10, 3, 5, 1, 9]pc6.
>>> pool_cap([4, 10, 3, 5, 1, 9], 6)
(4, 6, 3, 5, 1, 6)
"""
result = []
for value in pool:
if value > cap:
value = cap
result.append(value)
return tuple(result)
[docs]
@operation('pf')
def pool_floor(pool: Pool, floor: int) -> Pool:
"""Floor the minimum value in a pool.
:ref:`YADN` reference: :ref:`pool_floor`
:param pool: A sequence of die values.
:param floor: The minimum value of a die roll.
:return: The resulting values as a :class:`tuple`.
:rtype: tuple
Usage::
>>> # Roll [4, 10, 3, 5, 1, 9]pf6.
>>> pool_floor([4, 10, 3, 5, 1, 9], 6)
(6, 10, 6, 6, 6, 9)
"""
result = []
for value in pool:
if value < floor:
value = floor
result.append(value)
return tuple(result)
[docs]
@operation('pa')
def pool_keep_above(pool: Pool, floor: int) -> Pool:
"""Discard all values in a pool below a given value.
:ref:`YADN` reference: :ref:`pool_keep_above`
:param pool: A sequence of die values.
:param floor: The minimum value to keep.
:return: The resulting values as a :class:`tuple`.
:rtype: tuple
Usage::
>>> # Roll [4, 10, 3, 5, 1, 9]pa6.
>>> pool_keep_above([4, 10, 3, 5, 1, 9], 6)
(10, 9)
"""
return tuple(n for n in pool if n >= floor)
[docs]
@operation('pb')
def pool_keep_below(pool: Pool, ceiling: int) -> Pool:
"""Discard all values in a pool above a given value.
:ref:`YADN` reference: :ref:`pool_keep_below`
:param pool: A sequence of die values.
:param ceiling: The maximum value to keep.
:return: The resulting values as a :class:`tuple`.
:rtype: tuple
Usage::
>>> # Roll [4, 10, 3, 5, 1, 9]pb6.
>>> pool_keep_below([4, 10, 3, 5, 1, 9], 6)
(4, 3, 5, 1)
"""
return tuple(n for n in pool if n <= ceiling)
[docs]
@operation('ph')
def pool_keep_high(pool: Pool, keep: int) -> Pool:
"""Keep a number of the highest dice.
:ref:`YADN` reference: :ref:`pool_keep_high`
:param pool: A sequence of die values.
:param keep: The maximum value to keep.
:return: The resulting values as a :class:`tuple`.
:rtype: tuple
Usage::
>>> # Roll [4, 10, 3, 5, 1, 9]ph3.
>>> pool_keep_high([4, 10, 3, 5, 1, 9], 3)
(10, 5, 9)
"""
pool = list(pool)
remove = len(pool) - keep
for _ in range(remove):
low_value = max(pool)
low_index = 0
for i, n in enumerate(pool):
if n < low_value:
low_value = n
low_index = i
pool.pop(low_index)
return tuple(pool)
[docs]
@operation('pl')
def pool_keep_low(pool: Pool, keep: int) -> Pool:
"""Keep a number of the lowest dice.
:ref:`YADN` reference: :ref:`pool_keep_low`
:param pool: A sequence of die values.
:param keep: The maximum value to keep.
:return: The resulting values as a :class:`tuple`.
:rtype: tuple
Usage::
>>> # Roll [4, 10, 3, 5, 1, 9]pl3.
>>> pool_keep_low([4, 10, 3, 5, 1, 9], 3)
(4, 3, 1)
"""
pool = list(pool)
remove = len(pool) - keep
for _ in range(remove):
high_value = min(pool)
high_index = 0
for i, n in enumerate(pool):
if n > high_value:
high_value = n
high_index = i
pool.pop(high_index)
return tuple(pool)
[docs]
@operation('p%')
def pool_modulo(pool: Pool, divisor: int) -> Pool:
"""Perform a modulo operation of each member.
:ref:`YADN` reference: :ref:`pool_mod`
:param pool: A sequence of die values.
:param divisor: The maximum value to keep.
:return: The resulting values as a :class:`tuple`.
:rtype: tuple
Usage::
>>> # Roll [4, 10, 3, 5, 1, 9]p%3.
>>> pool_modulo([4, 10, 3, 5, 1, 9], 3)
(1, 1, 0, 2, 1, 0)
"""
return tuple(n % divisor for n in pool)
[docs]
@operation('pr')
def pool_remove(pool: Pool, cut: int) -> Pool:
"""Remove members of a pool of the given value.
:ref:`YADN` reference: :ref:`pool_remove`
:param pool: A sequence of die values.
:param cut: The maximum value to keep.
:return: The resulting values as a :class:`tuple`.
:rtype: tuple
Usage::
>>> # Roll [4, 10, 3, 5, 1, 9]pr5.
>>> pool_remove([4, 10, 3, 5, 1, 9], 5)
(4, 10, 3, 1, 9)
"""
return tuple(n for n in pool if n != cut)
# Pool degeneration operators.
[docs]
@operation('C')
def pool_concatenate(pool: Pool) -> int:
"""Concatenate the dice in the pool.
:ref:`YADN` reference: :ref:`pool_concat`
:param pool: A sequence of die values.
:return: The resulting value as an :class:`int`.
:rtype: int
Usage::
>>> # Roll C[4, 10, 3, 5, 1, 9].
>>> pool_concatenate([4, 10, 3, 5, 1, 9])
4103519
"""
str_value = ''.join((str(m) for m in pool))
return int(str_value)
[docs]
@operation('N')
def pool_count(pool: Pool) -> int:
"""Count the dice in the pool.
:ref:`YADN` reference: :ref:`pool_count`
:param pool: A sequence of die values.
:return: The resulting value as an :class:`int`.
:rtype: int
Usage::
>>> # Roll N[4, 10, 3, 5, 1, 9].
>>> pool_count([4, 10, 3, 5, 1, 9])
6
"""
return len(pool)
[docs]
@operation('S')
def pool_sum(pool: Pool) -> int:
"""Sum the dice in the pool.
:ref:`YADN` reference: :ref:`pool_sum`
:param pool: A sequence of die values.
:return: The resulting value as an :class:`int`.
:rtype: int
Usage::
>>> # Roll S[4, 10, 3, 5, 1, 9].
>>> pool_sum([4, 10, 3, 5, 1, 9])
32
"""
return sum(pool)
[docs]
@operation('ns')
def count_successes(pool: Pool, target: int) -> int:
"""Count the number of successes in the pool.
:ref:`YADN` reference: :ref:`count_successes`
:param pool: A sequence of die values.
:param target: The target number for success.
:return: The resulting value as an :class:`int`.
:rtype: int
Usage::
>>> # Roll [4, 10, 3, 5, 1, 9]ns6.
>>> count_successes([4, 10, 3, 5, 1, 9], 6)
2
"""
pool = pool_keep_above(pool, target)
return len(pool)
[docs]
@operation('nb')
def count_successes_with_botch(pool: Pool, target: int) -> int:
"""Count the number of successes in the pool. Then remove a success
for each botch.
:ref:`YADN` reference: :ref:`count_botch`
:param pool: A sequence of die values.
:param target: The target number for success.
:return: The resulting value as an :class:`int`.
:rtype: int
Usage::
>>> # Roll [4, 10, 3, 5, 1, 9]nb6.
>>> count_successes_with_botch([4, 10, 3, 5, 1, 9], 6)
1
"""
botches = len([n for n in pool if n == 1])
pool = pool_keep_above(pool, target)
return len(pool) - botches
# Pool generation operator.
[docs]
@operation('g')
def dice_pool(num: int, size: int) -> Pool:
"""Roll a dice pool.
:ref:`YADN` reference: :ref:`dice_pool`
:param num: The number of dice to roll.
:param size: The highest number that can be rolled on a die.
:return: The the values as a :class:`tuple`.
:rtype: tuple
Usage::
>>> # This line is to ensure predictability for testing.
>>> # Do not use outside of test cases.
>>> _seed('spam')
>>>
>>> # Roll 5g6.
>>> dice_pool(5, 6)
(1, 1, 3, 5, 5)
"""
return tuple(roll(size) for _ in range(num))
[docs]
@operation('g!')
def exploding_pool(num: int, size: int) -> Pool:
"""Roll an exploding dice pool.
:ref:`YADN` reference: :ref:`exploding_pool`
:param num: The number of dice to roll.
:param size: The highest number that can be rolled on a die.
:return: The the values as a :class:`tuple`.
:rtype: tuple
Usage::
>>> # This line is to ensure predictability for testing.
>>> # Do not use outside of test cases.
>>> _seed('spam')
>>>
>>> # Roll 5g!6.
>>> exploding_pool(5, 6)
(1, 1, 3, 5, 5)
"""
return tuple(_explode(size) for n in range(num))
# Generation operators that use `secrets` instead of `random`.
def dice_pool_secrets(num: int, size: int) -> Pool:
"""Roll a dice pool using :mod:`secrets` instead of :mod:`random`.
The extra security provided by :mod:`secrets` is unnecessary for
the intended usage of :mod:`yadr`, but it's here if you want it.
:ref:`YADN` reference: :ref:`dice_pool`
:param num: The number of dice to roll.
:param size: The highest number that can be rolled on a die.
:return: The the values as a :class:`tuple`.
:rtype: tuple
.. warning::
This function is experimental and may be removed from future
versions of :mod:`yadr` without notice. Use at your own risk.
"""
return tuple(_roll_secrets(size) for _ in range(num))
# Utility functions.
def _explode(size: int) -> int:
"""Explode the value of a die.
:param size: The highest number that can be rolled on a die.
:return: The the values as an :class:`int`.
:rtype: int
"""
result = roll(size)
if result == size:
result += _explode(size)
return result
def _seed(seed: int | str | bytes) -> None:
"""Seed the random number generator for testing purposes.
:param seed: A seed value for the random number generator.
:return: None.
:rtype: NoneType
"""
if isinstance(seed, str):
seed = bytes(seed, encoding='utf_8')
if isinstance(seed, bytes):
seed = int.from_bytes(seed, 'little')
random.seed(seed)
# Die rolling.
def _roll_random(size: int) -> int:
"""Roll a die.
:param size: The size of the die to roll.
:returns: An :class:'int' object.
:rtype: int
"""
return random.randint(1, size)
def _roll_secrets(size: int) -> int:
"""Roll a die using :mod:`secrets`.
:param size: The size of the die to roll.
:returns: An :class:'int' object.
:rtype: int
"""
return secrets.randbelow(size) + 1
# This sets the function used by the generation operations to roll dice.
# If you want to use `secrets` instead of `random`, change this to
# :func:`_roll_secrets`. Since this is happening at the global level,
# you could, theoretically, run into thread safety issues when changing it,
# So be cautious.
roll = _roll_random