ctan
2a58826bd4
Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
123 lines
3.5 KiB
Python
123 lines
3.5 KiB
Python
"""Dice notation parser and expression evaluator.
|
|
|
|
Design: `DiceExpr` is a frozen Pydantic model representing a dice expression.
|
|
Normal dice have count >= 1, sides >= 2. Flat bonuses (e.g. parse_dice("7"))
|
|
are represented with count=0, sides=0 as a sentinel — this is the ONLY path
|
|
where count=0 is valid, produced exclusively by `DiceExpr.flat()` or
|
|
`parse_dice()` on a plain integer.
|
|
|
|
The `roll` method returns the raw sum + bonus (no minimum floor).
|
|
Damage floors are the caller's concern.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import re
|
|
from typing import TYPE_CHECKING
|
|
|
|
from pydantic import BaseModel, ConfigDict, model_validator
|
|
|
|
if TYPE_CHECKING:
|
|
from pf1e_simulator.rng import Rng
|
|
|
|
_MIN_SIDES = 2
|
|
|
|
_DICE_RE = re.compile(r"^(\d*)d(\d+)([+-]\d+)?$")
|
|
|
|
|
|
class DiceParseError(Exception):
|
|
"""Raised when dice notation cannot be parsed.
|
|
|
|
Carries the offending text in its message.
|
|
"""
|
|
|
|
def __init__(self, text: str) -> None:
|
|
super().__init__(f"invalid dice notation: {text!r}")
|
|
self.text = text
|
|
|
|
|
|
class DiceExpr(BaseModel):
|
|
"""A parsed dice expression.
|
|
|
|
Invariants (enforced by model_validator):
|
|
- Normal dice: count >= 1, sides >= 2
|
|
- Flat bonus sentinel: count == 0, sides == 0 (via .flat() or parse_dice)
|
|
"""
|
|
|
|
model_config = ConfigDict(frozen=True)
|
|
|
|
count: int
|
|
sides: int
|
|
bonus: int = 0
|
|
|
|
@model_validator(mode="after")
|
|
def _validate_invariants(self) -> DiceExpr:
|
|
if self.count == 0 and self.sides == 0:
|
|
return self
|
|
if self.count < 1:
|
|
msg = f"count must be >= 1 for dice expressions, got {self.count}"
|
|
raise ValueError(msg)
|
|
if self.sides < _MIN_SIDES:
|
|
msg = f"sides must be >= {_MIN_SIDES} for dice expressions, got {self.sides}"
|
|
raise ValueError(msg)
|
|
return self
|
|
|
|
@classmethod
|
|
def flat(cls, n: int) -> DiceExpr:
|
|
"""Create a flat-bonus expression (no dice rolled)."""
|
|
return cls(count=0, sides=0, bonus=n)
|
|
|
|
def roll(self, rng: Rng) -> int:
|
|
"""Roll the expression using the given RNG.
|
|
|
|
Returns sum of `count` rolls of `sides` + `bonus`.
|
|
For flat expressions (count=0), returns `bonus` without calling rng.
|
|
"""
|
|
if self.count == 0:
|
|
return self.bonus
|
|
total = sum(rng.roll_die(self.sides) for _ in range(self.count))
|
|
return total + self.bonus
|
|
|
|
def mean(self) -> float:
|
|
"""Expected value: count * (sides + 1) / 2 + bonus."""
|
|
if self.count == 0:
|
|
return float(self.bonus)
|
|
return self.count * (self.sides + 1) / 2 + self.bonus
|
|
|
|
|
|
def parse_dice(text: str) -> DiceExpr:
|
|
"""Parse a dice notation string into a DiceExpr.
|
|
|
|
Accepted forms:
|
|
- "2d6+3", "2d6-1", "1d4", "d20" (standard notation)
|
|
- "7", "-5", "0" (flat integers)
|
|
|
|
Raises DiceParseError on invalid input.
|
|
"""
|
|
stripped = text.strip()
|
|
if not stripped:
|
|
raise DiceParseError(text)
|
|
|
|
m = _DICE_RE.match(stripped)
|
|
if m:
|
|
count_str, sides_str, bonus_str = m.groups()
|
|
if not sides_str:
|
|
raise DiceParseError(text)
|
|
count = int(count_str) if count_str else 1
|
|
sides = int(sides_str)
|
|
bonus = int(bonus_str) if bonus_str else 0
|
|
|
|
if count == 0:
|
|
raise DiceParseError(text)
|
|
if sides < 1:
|
|
raise DiceParseError(text)
|
|
|
|
return DiceExpr(count=count, sides=sides, bonus=bonus)
|
|
|
|
try:
|
|
n = int(stripped)
|
|
except ValueError:
|
|
raise DiceParseError(text) from None
|
|
|
|
return DiceExpr.flat(n)
|