class DropKeepExpression(IDiceExpression):
TOKEN_RULE = """expression : DROP_KEEP_DICE"""
@staticmethod
def add_production_function(
parser_generator: rply.ParserGenerator,
probability_distribution_factory: IProbabilityDistributionFactory
) -> Callable:
@parser_generator.production(DropKeepExpression.TOKEN_RULE)
def drop_keep(_, tokens) -> IDiceExpression:
return DropKeepExpression(tokens[0].value,
probability_distribution_factory)
return drop_keep
def __init__(
self, string_form: str,
probability_distribution_factory: IProbabilityDistributionFactory):
self._string_form = string_form
self._number_of_dice = self._get_number_of_dice()
self._single_dice_outcome_map = get_single_dice_outcome_map(
re.split(r"[dk]", self._string_form)[1])
self._number_to_keep_or_drop = self._get_number_to_keep_or_drop()
self._simplified_form = None
self._probability_distribution_factory = probability_distribution_factory
if self._is_keep():
if self._number_of_dice <= self._number_to_keep_or_drop:
self._simplified_form = DiceExpression(
f"{self._number_of_dice}d{self._get_number_of_sides_string()}",
self._probability_distribution_factory,
)
elif self._number_to_keep_or_drop == 0:
self._simplified_form = ConstantIntegerExpression(
"0", self._probability_distribution_factory)
else:
if self._number_of_dice <= self._number_to_keep_or_drop:
self._simplified_form = ConstantIntegerExpression(
"0", self._probability_distribution_factory)
elif self._number_to_keep_or_drop == 0:
self._simplified_form = DiceExpression(
f"{self._number_of_dice}d{self._get_number_of_sides_string()}",
self._probability_distribution_factory,
)
def _get_number_of_dice(self) -> int:
string_num = re.split(r"[dk]", self._string_form)[0]
return 1 if string_num == "" else int(string_num)
def _get_number_to_keep_or_drop(self) -> int:
return int(re.split(r"[dk]", self._string_form)[2])
def _get_number_of_sides_string(self) -> str:
return re.split(r"[dk]", self._string_form)[1]
def _is_keep(self) -> bool:
return "k" in self._string_form
def roll(self) -> int:
if self._simplified_form is not None:
return self._simplified_form.roll()
dice_rolls = random.choices(
list(self._single_dice_outcome_map.keys()),
weights=list(self._single_dice_outcome_map.values()),
k=self._number_of_dice,
)
dice_rolls.sort()
if self._is_keep():
return sum(dice_rolls[-self._number_to_keep_or_drop:])
return sum(dice_rolls[:self._number_of_dice -
self._number_to_keep_or_drop])
def max(self) -> int:
if self._simplified_form is not None:
return self._simplified_form.max()
if self._is_keep():
return self._number_to_keep_or_drop * max(
self._single_dice_outcome_map.keys())
return (self._number_of_dice - self._number_to_keep_or_drop) * max(
self._single_dice_outcome_map.keys())
def min(self) -> int:
if self._simplified_form is not None:
return self._simplified_form.min()
if self._is_keep():
return self._number_to_keep_or_drop * min(
self._single_dice_outcome_map.keys())
return (self._number_of_dice - self._number_to_keep_or_drop) * min(
self._single_dice_outcome_map.keys())
def __str__(self) -> str:
return f"{self._string_form}"
def estimated_cost(self) -> int:
if self._simplified_form is not None:
return self._simplified_form.estimated_cost()
if self._is_keep():
return (self._number_of_dice * self._number_to_keep_or_drop) * len(
self._single_dice_outcome_map.values())
return (self._number_of_dice *
(self._number_of_dice - self._number_to_keep_or_drop)) * len(
self._single_dice_outcome_map.values())
def get_probability_distribution(self) -> IProbabilityDistribution:
if self._simplified_form is not None:
return self._simplified_form.get_probability_distribution()
is_keep = self._is_keep()
number_of_dice_to_select = self._number_of_dice - self._number_to_keep_or_drop
number_of_dice_remaining = self._number_to_keep_or_drop
if is_keep:
number_of_dice_to_select = self._number_to_keep_or_drop
number_of_dice_remaining = self._number_of_dice - self._number_to_keep_or_drop
current = DropKeepExpression._build_dice_dict(
number_of_dice_to_select, self._single_dice_outcome_map)
current = DropKeepExpression._compute_iterations(
current, number_of_dice_remaining, self._single_dice_outcome_map,
is_keep)
out_come_map = DropKeepExpression._collapse_outcomes(current)
return self._probability_distribution_factory.create(out_come_map)
@staticmethod
def _build_dice_dict(number_of_dice: int,
dice_outcome_map: Dict[int, int]) -> Dict[str, int]:
def safe_add_to_dict(dictionary: Dict[str, int], key: str,
value: int) -> None:
if key not in dictionary:
dictionary[key] = 0
dictionary[key] += value
current_dict = {"": 1}
for _ in range(number_of_dice):
new_dict: Dict[str, int] = {}
for dice_outcome, count in dice_outcome_map.items():
for old_key, old_value in current_dict.items():
new_key = DropKeepExpression._string_key_to_list(old_key)
new_key.append(dice_outcome)
new_key.sort()
safe_add_to_dict(
new_dict,
DropKeepExpression._int_list_to_string(new_key),
old_value * count,
)
current_dict = new_dict
return current_dict
@staticmethod
def _string_key_to_list(string: str) -> List[int]:
if string == "":
return []
return [int(n) for n in string.split("|")]
@staticmethod
def _int_list_to_string(values: List[int]) -> str:
values.sort()
return "|".join([str(n) for n in values])
@staticmethod
def _update_key_list(old_key_string: str, new_value: int,
is_keep: bool) -> str:
old_value = DropKeepExpression._string_key_to_list(old_key_string)
old_value.append(new_value)
old_value.sort()
if is_keep:
old_value = old_value[1:]
else:
old_value = old_value[:-1]
return DropKeepExpression._int_list_to_string(old_value)
@staticmethod
def _compute_iterations(
current: Dict[str, int],
number_of_dice: int,
dice_outcome_map: Dict[int, int],
is_keep: bool,
) -> Dict[str, int]:
def safe_add_to_dict(dictionary: Dict[str, int], key: str,
value: int) -> None:
if key not in dictionary:
dictionary[key] = 0
dictionary[key] += value
current_dict = current
for _ in range(number_of_dice):
new_dict: Dict[str, int] = {}
for dice_outcome, count in dice_outcome_map.items():
for old_key, old_value in current_dict.items():
new_key = DropKeepExpression._update_key_list(
old_key, dice_outcome, is_keep)
safe_add_to_dict(new_dict, new_key, old_value * count)
current_dict = new_dict
return current_dict
@staticmethod
def _collapse_outcomes(outcomes: Dict[str, int]) -> Dict[int, int]:
def safe_add_to_dict(dictionary: Dict[int, int], key: int,
value: int) -> None:
if key not in dictionary:
dictionary[key] = 0
dictionary[key] += value
new_dict: Dict[int, int] = {}
for current_key, current_value in outcomes.items():
total = sum(DropKeepExpression._string_key_to_list(current_key))
safe_add_to_dict(new_dict, total, current_value)
return new_dict
def get_contained_variables(self, ) -> Set[str]:
return set()
class TestDiceExpression(TestCase):
def setUp(self):
self._probability_distribution_factory = ProbabilityDistributionFactory(
)
self._test_dice = DiceExpression(
"4d6", self._probability_distribution_factory)
self._mock_parser_gen = create_autospec(rply.ParserGenerator)
def test_dice_add_production_function(self):
DiceExpression.add_production_function(
self._mock_parser_gen, self._probability_distribution_factory)
self._mock_parser_gen.production.assert_called_once_with(
"""expression : DICE""")
def test_dice_roll(self):
self._test_dice = DiceExpression(
"2d10", self._probability_distribution_factory)
roll_set = set()
for _ in range(1000):
roll_set.add(self._test_dice.roll())
self.assertEqual(19, len(roll_set))
self.assertEqual(20, max(roll_set))
self.assertEqual(2, min(roll_set))
def test_dice_roll_missing_dice_amount(self):
self._test_dice = DiceExpression(
"d10", self._probability_distribution_factory)
roll_set = set()
for _ in range(1000):
roll_set.add(self._test_dice.roll())
self.assertEqual(10, len(roll_set))
self.assertEqual(10, max(roll_set))
self.assertEqual(1, min(roll_set))
def test_dice_roll_percentile_dice(self):
self._test_dice = DiceExpression(
"1d%", self._probability_distribution_factory)
roll_set = set()
for _ in range(10000):
roll_set.add(self._test_dice.roll())
self.assertEqual(100, len(roll_set))
self.assertEqual(100, max(roll_set))
self.assertEqual(1, min(roll_set))
def test_dice_roll_fate_dice(self):
self._test_dice = DiceExpression(
"2dF", self._probability_distribution_factory)
roll_set = set()
for _ in range(10000):
roll_set.add(self._test_dice.roll())
self.assertSetEqual(roll_set, {-2, -1, 0, 1, 2})
def test_dice_roll_custom_dice_negative(self):
self._test_dice = DiceExpression(
"2d[-2,2,100]", self._probability_distribution_factory)
roll_set = set()
for _ in range(10000):
roll_set.add(self._test_dice.roll())
self.assertSetEqual(roll_set, {-4, 0, 4, 98, 102, 200})
def test_dice_roll_custom_dice_large_set(self):
self._test_dice = DiceExpression(
"2d[-2,0,2,4,6,31]", self._probability_distribution_factory)
roll_set = set()
for _ in range(10000):
roll_set.add(self._test_dice.roll())
self.assertSetEqual(
roll_set, {-4, -2, 0, 2, 4, 6, 8, 10, 12, 29, 31, 33, 35, 37, 62})
def test_dice_max(self):
self.assertEqual(24, self._test_dice.max())
def test_dice_max_missing_dice_amount(self):
self._test_dice = DiceExpression(
"d10", self._probability_distribution_factory)
self.assertEqual(10, self._test_dice.max())
def test_dice_max_percentile_dice(self):
self._test_dice = DiceExpression(
"1d%", self._probability_distribution_factory)
self.assertEqual(100, self._test_dice.max())
def test_dice_max_fate_dice(self):
self._test_dice = DiceExpression(
"2dF", self._probability_distribution_factory)
self.assertEqual(2, self._test_dice.max())
def test_dice_max_custom_dice_negative(self):
self._test_dice = DiceExpression(
"21d[-2,2,100]", self._probability_distribution_factory)
self.assertEqual(2100, self._test_dice.max())
def test_dice_max_custom_dice_large_set(self):
self._test_dice = DiceExpression(
"76d[-2,0,2,4,6,31]", self._probability_distribution_factory)
self.assertEqual(2356, self._test_dice.max())
def test_dice_min(self):
self.assertEqual(4, self._test_dice.min())
def test_dice_min_missing_dice_amount(self):
self._test_dice = DiceExpression(
"d10", self._probability_distribution_factory)
self.assertEqual(1, self._test_dice.min())
def test_dice_min_percentile_dice(self):
self._test_dice = DiceExpression(
"1d%", self._probability_distribution_factory)
self.assertEqual(1, self._test_dice.min())
def test_dice_min_fate_dice(self):
self._test_dice = DiceExpression(
"4dF", self._probability_distribution_factory)
self.assertEqual(-4, self._test_dice.min())
def test_dice_min_custom_dice_negative(self):
self._test_dice = DiceExpression(
"21d[-2,2,100]", self._probability_distribution_factory)
self.assertEqual(-42, self._test_dice.min())
def test_dice_min_custom_dice_large_set(self):
self._test_dice = DiceExpression(
"d[-2,0,2,4,6,31,-2,-24]", self._probability_distribution_factory)
self.assertEqual(-24, self._test_dice.min())
def test_dice_str(self):
self.assertEqual("4d6", str(self._test_dice))
def test_dice_str_missing_dice_amount(self):
self._test_dice = DiceExpression(
"d10", self._probability_distribution_factory)
self.assertEqual("d10", str(self._test_dice))
def test_dice_str_percentile_dice(self):
self._test_dice = DiceExpression(
"1d%", self._probability_distribution_factory)
self.assertEqual("1d%", str(self._test_dice))
def test_dice_str_fate_dice(self):
self._test_dice = DiceExpression(
"4dF", self._probability_distribution_factory)
self.assertEqual("4dF", str(self._test_dice))
def test_dice_str_custom_dice_negative(self):
self._test_dice = DiceExpression(
"21d[-2,2,100]", self._probability_distribution_factory)
self.assertEqual("21d[-2,2,100]", str(self._test_dice))
def test_dice_str_custom_dice_large_set(self):
self._test_dice = DiceExpression(
"d[-2,0,2,4,6,31,-2,-24]", self._probability_distribution_factory)
self.assertEqual("d[-2,0,2,4,6,31,-2,-24]", str(self._test_dice))
def test_dice_estimated_cost(self):
self.assertEqual(4 * 6, self._test_dice.estimated_cost())
def test_dice_estimated_cost_missing_dice_amount(self):
self._test_dice = DiceExpression(
"d10", self._probability_distribution_factory)
self.assertEqual(10, self._test_dice.estimated_cost())
def test_dice_estimated_cost_percentile_dice(self):
self._test_dice = DiceExpression(
"1d%", self._probability_distribution_factory)
self.assertEqual(100, self._test_dice.estimated_cost())
def test_dice_estimated_cost_fate_dice(self):
self._test_dice = DiceExpression(
"10dF", self._probability_distribution_factory)
self.assertEqual(30, self._test_dice.estimated_cost())
def test_dice_estimated_custom_dice_negative(self):
self._test_dice = DiceExpression(
"21d[-2,2,100]", self._probability_distribution_factory)
self.assertEqual(63, self._test_dice.estimated_cost())
def test_dice_estimated_custom_dice_large_set(self):
self._test_dice = DiceExpression(
"2d[-2,0,2,4,6*7,31,-2,-24]",
self._probability_distribution_factory)
self.assertEqual(14, self._test_dice.estimated_cost())
def test_dice_get_probability_distribution(self):
self._test_dice = DiceExpression(
"4d6", self._probability_distribution_factory)
possible_rolls = itertools.product(range(1, 7), repeat=4)
results = [sum(t) for t in possible_rolls]
self.assertEqual(
dict(collections.Counter(results)),
self._test_dice.get_probability_distribution().get_result_map(),
)
def test_dice_get_probability_distribution_missing_dice_amount(self):
self._test_dice = DiceExpression(
"d10", self._probability_distribution_factory)
possible_rolls = itertools.product(range(1, 11), repeat=1)
results = [sum(t) for t in possible_rolls]
self.assertEqual(
dict(collections.Counter(results)),
self._test_dice.get_probability_distribution().get_result_map(),
)
def test_dice_get_probability_distribution_percentile_dice(self):
self._test_dice = DiceExpression(
"2d%", self._probability_distribution_factory)
possible_rolls = itertools.product(range(1, 101), repeat=2)
results = [sum(t) for t in possible_rolls]
self.assertEqual(
dict(collections.Counter(results)),
self._test_dice.get_probability_distribution().get_result_map(),
)
def test_dice_get_probability_distribution_fate_dice(self):
self._test_dice = DiceExpression(
"4dF", self._probability_distribution_factory)
self.assertEqual(
{
-4: 1,
-3: 4,
-2: 10,
-1: 16,
0: 19,
1: 16,
2: 10,
3: 4,
4: 1
},
self._test_dice.get_probability_distribution().get_result_map(),
)
def test_dice_get_probability_distribution_custom_dice_negative(self):
self._test_dice = DiceExpression(
"2d[-2,2,100]", self._probability_distribution_factory)
self.assertEqual(
{
-4: 1,
0: 2,
4: 1,
98: 2,
102: 2,
200: 1
},
self._test_dice.get_probability_distribution().get_result_map(),
)
def test_dice_get_probability_distribution_custom_dice_large_set(self):
self._test_dice = DiceExpression(
"d[-2,0,2,4,6,31,-2,-24]", self._probability_distribution_factory)
self.assertEqual(
{
-24: 1,
-2: 2,
0: 1,
2: 1,
4: 1,
6: 1,
31: 1
},
self._test_dice.get_probability_distribution().get_result_map(),
)
def test_dice_get_probability_distribution_custom_dice_multiplier(self):
self._test_dice = DiceExpression(
"d[-2*3,0,2]", self._probability_distribution_factory)
self.assertEqual(
{
-2: 3,
0: 1,
2: 1
},
self._test_dice.get_probability_distribution().get_result_map(),
)
def test_dice_get_probability_distribution_custom_dice_range(self):
self._test_dice = DiceExpression(
"d[1-2*3,0,2]", self._probability_distribution_factory)
self.assertEqual(
{
0: 1,
1: 3,
2: 4
},
self._test_dice.get_probability_distribution().get_result_map(),
)
def test_dice_get_probability_distribution_custom_dice_multiplier_range(
self):
self._test_dice = DiceExpression(
"d[1-2*3,-4--9*10]", self._probability_distribution_factory)
self.assertEqual(
{
-9: 10,
-8: 10,
-7: 10,
-6: 10,
-5: 10,
-4: 10,
1: 3,
2: 3
},
self._test_dice.get_probability_distribution().get_result_map(),
)
def test_dice_get_contained_variables(self):
self.assertSetEqual(set(), self._test_dice.get_contained_variables())