From 77e6442f6cb0df2db160b5866c592b05c466b06a Mon Sep 17 00:00:00 2001
From: koaning <vincentwarmerdam@gmail.com>
Date: Wed, 17 Jun 2026 12:08:16 +0200
Subject: [PATCH] Add autocomplete-friendly type hints

---
 .github/workflows/tests.yml |   2 +
 Makefile                    |   3 +
 dicekit/__init__.py         | 124 +++++++++++++++++-----------------
 dicekit/py.typed            |   1 +
 nbs/__init__.py             | 130 +++++++++++++++++++-----------------
 pyproject.toml              |  13 ++++
 tests/test_basics.py        |   5 ++
 uv.lock                     |  34 +++++++++-
 8 files changed, 188 insertions(+), 124 deletions(-)
 create mode 100644 dicekit/py.typed

diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
index 39cc5d0..f6d06f9 100644
--- a/.github/workflows/tests.yml
+++ b/.github/workflows/tests.yml
@@ -27,6 +27,8 @@ jobs:
       run: uv sync --group dev
     - name: Run tests
       run: uv run pytest
+    - name: Run type checks
+      run: uv run basedpyright dicekit
 
   build:
     runs-on: ubuntu-latest
diff --git a/Makefile b/Makefile
index 24edf16..4101f43 100644
--- a/Makefile
+++ b/Makefile
@@ -15,5 +15,8 @@ pypi: clean test
 test:
 	uv run pytest
 
+typecheck:
+	uv run basedpyright dicekit
+
 clean:
 	rm -rf dist nbs/__pycache__ dicekit/__pycache__
diff --git a/dicekit/__init__.py b/dicekit/__init__.py
index 85c40f9..f82bf5d 100644
--- a/dicekit/__init__.py
+++ b/dicekit/__init__.py
@@ -4,13 +4,15 @@
 import marimo as mo
 from hastyplot import qplot
 import random
-from collections.abc import Mapping
+from collections.abc import Callable, Hashable, Mapping, Sequence
 from collections import Counter
-from itertools import product, combinations, permutations, combinations_with_replacement
+from itertools import product, permutations
 from functools import reduce
+from typing import Any, Self, TypeAlias, cast
 
+Weight: TypeAlias = int | float
 
-class Dice:
+class Dice[T: Hashable]:
     """
     A class representing dice with arbitrary probability distributions.
 
@@ -18,7 +20,7 @@ class Dice:
     custom probability distributions for each face value.
     """
 
-    def __init__(self, probs):
+    def __init__(self, probs: Mapping[T, Weight]) -> None:
         """
         Initialize a Dice object with given probabilities.
 
@@ -28,7 +30,7 @@ def __init__(self, probs):
         self.probs = {k: v / sum(probs.values()) for k, v in probs.items()}
 
     @classmethod
-    def from_sides(cls, n=6):
+    def from_sides(cls, n: int = 6) -> "Dice[int]":
         """
         Create a fair dice with n sides.
 
@@ -38,10 +40,10 @@ def from_sides(cls, n=6):
         Returns:
             Dice: A fair dice with n sides
         """
-        return cls({i: 1 / n for i in range(1, n + 1)})
+        return cast(Dice[int], cast(Any, cls)({i: 1 / n for i in range(1, n + 1)}))
 
     @classmethod
-    def from_numbers(cls, *args):
+    def from_numbers[U: Hashable](cls, *args: U) -> "Dice[U]":
         """
         Create a dice from a sequence of numbers with frequencies.
 
@@ -52,9 +54,9 @@ def from_numbers(cls, *args):
             Dice: A dice with probabilities based on the frequency of each number
         """
         c = Counter(args)
-        return cls({k: v / len(args) for k, v in args.items()})
+        return cast(Dice[U], cast(Any, cls)({k: v / len(args) for k, v in c.items()}))
 
-    def roll(self, n=1):
+    def roll(self, n: int = 1) -> list[T]:
         """
         Simulate rolling the dice n times.
 
@@ -66,7 +68,7 @@ def roll(self, n=1):
         """
         return random.choices(list(self.probs.keys()), weights=list(self.probs.values()), k=n)
 
-    def sample(self, n=1):
+    def sample(self, n: int = 1) -> list[T]:
         """
         Simulate rolling the dice n times.
 
@@ -78,7 +80,7 @@ def sample(self, n=1):
         """
         return self.roll(n=1)
 
-    def map(self, transform):
+    def map[U: Hashable](self, transform: Callable[[T], U] | Mapping[T, U]) -> "Dice[U]":
         """
         Transform the outcomes of the dice.
 
@@ -95,7 +97,11 @@ def map(self, transform):
             new_probs[new_outcome] = new_probs.get(new_outcome, 0) + probability
         return Dice(new_probs)
 
-    def operate(self, other, operator):
+    def operate[U: Hashable, V: Hashable](
+        self,
+        other: "Dice[U] | U",
+        operator: Callable[[T, U], V],
+    ) -> "Dice[V]":
         """
         Apply an operation between this dice and another dice or number.
 
@@ -111,7 +117,8 @@ def operate(self, other, operator):
                 other = Dice({other: 1})
             except TypeError as exc:
                 raise TypeError("operations require a Dice or hashable value") from exc
-        new_probs = {}
+        other = cast(Dice[U], other)
+        new_probs: dict[V, float] = {}
         for s1, p1 in self.probs.items():
             for s2, p2 in other.probs.items():
                 new_key = operator(s1, s2)
@@ -120,7 +127,7 @@ def operate(self, other, operator):
                 new_probs[new_key] += p1 * p2
         return Dice(new_probs)
 
-    def filter(self, func):
+    def filter(self, func: Callable[[T], bool]) -> "Dice[T]":
         """
         Create a new dice by filtering face values based on a function.
 
@@ -134,7 +141,7 @@ def filter(self, func):
         total_prob = sum(new_probs.values())
         return Dice({k: v / total_prob for k, v in new_probs.items()})
 
-    def _repr_html_(self):
+    def _repr_html_(self) -> str:
         """
         Return HTML representation of the dice for display in notebooks.
 
@@ -143,7 +150,7 @@ def _repr_html_(self):
         """
         return mo.as_html(self.prob_chart()).text
 
-    def prob_chart(self):
+    def prob_chart(self) -> Any:
         """
         Create a visualization of the dice probability distribution.
 
@@ -163,7 +170,7 @@ def prob_chart(self):
             theme="default",
         )
 
-    def cdf_chart(self):
+    def cdf_chart(self) -> Any:
         """
         Create a visualization of the cumulative distribution function.
 
@@ -188,7 +195,7 @@ def cdf_chart(self):
             theme="default",
         )
 
-    def ordered(self, n, k=None):
+    def ordered(self, n: int, k: int | None = None) -> list["Dice[T]"]:
         """
         Take the dice `n` times and calculate the ordered distributions for it.
 
@@ -202,7 +209,7 @@ def ordered(self, n, k=None):
             return ordered(*items, k=k)
         return ordered(*items)
 
-    def out_of(self, n=2, func=max):
+    def out_of[U: Hashable](self, n: int = 2, func: Callable[[list[T]], U] = max) -> "Dice[U]":
         """
         Create a dice representing the result of applying a function to n rolls.
 
@@ -213,7 +220,7 @@ def out_of(self, n=2, func=max):
         Returns:
             Dice: A new dice representing the distribution of the function's results
         """
-        result = {}
+        result: dict[U, float] = {}
         dice_in = [self] * n
         for _i in product(*[d.probs.items() for d in dice_in]):
             values = [_[0] for _ in _i]
@@ -224,47 +231,47 @@ def out_of(self, n=2, func=max):
             result[outcome] += prob
         return Dice(result)
 
-    def __add__(self, other):
+    def __add__(self, other: Any) -> "Dice[Any] | Self":
         if not isinstance(other, Dice) and other == 0:
             return self
         return self.operate(other, lambda a, b: a + b)
 
-    def __radd__(self, other):
+    def __radd__(self, other: Any) -> "Dice[Any] | Self":
         if not isinstance(other, Dice) and other == 0:
             return self
         return self.operate(other, lambda a, b: a + b)
 
-    def __sub__(self, other):
+    def __sub__(self, other: Any) -> "Dice[Any]":
         return self.operate(other, lambda a, b: a - b)
 
-    def __rsub__(self, other):
+    def __rsub__(self, other: Any) -> "Dice[Any]":
         return self.operate(other, lambda a, b: b - a)
 
-    def __mul__(self, other):
+    def __mul__(self, other: Any) -> "Dice[Any]":
         return self.operate(other, lambda a, b: a * b)
 
-    def __rmul__(self, other):
+    def __rmul__(self, other: Any) -> "Dice[Any]":
         return self.operate(other, lambda a, b: a * b)
 
-    def __le__(self, other):
+    def __le__(self, other: Any) -> "Dice[bool]":
         return self.operate(other, lambda a, b: a <= b)
 
-    def __lt__(self, other):
+    def __lt__(self, other: Any) -> "Dice[bool]":
         return self.operate(other, lambda a, b: a < b)
 
-    def __ge__(self, other):
+    def __ge__(self, other: Any) -> "Dice[bool]":
         return self.operate(other, lambda a, b: a >= b)
 
-    def __gt__(self, other):
+    def __gt__(self, other: Any) -> "Dice[bool]":
         return self.operate(other, lambda a, b: a > b)
 
-    def __eq__(self, other):
+    def __eq__(self, other: Any) -> "Dice[bool]":  # type: ignore[override]
         return self.operate(other, lambda a, b: a == b)
 
-    def __ne__(self, other):
+    def __ne__(self, other: Any) -> "Dice[bool]":  # type: ignore[override]
         return self.operate(other, lambda a, b: a != b)
 
-    def __len__(self):
+    def __len__(self) -> int:
         return len(self.probs)
 
 
@@ -277,7 +284,7 @@ class Vase:
     replacement, where the order of drawn items may optionally matter.
     """
 
-    def __init__(self, contents):
+    def __init__(self, contents: Sequence[str]) -> None:
         """
         Initialize a vase with the items it contains.
 
@@ -287,7 +294,7 @@ def __init__(self, contents):
         self._contents = contents
 
     @classmethod
-    def from_counts(self, **kwargs):
+    def from_counts(cls, **kwargs: int) -> "Vase":
         """
         Create a vase from item names and their quantities.
 
@@ -300,9 +307,9 @@ def from_counts(self, **kwargs):
         contents = []
         for k, v in kwargs.items():
             contents.extend([k]*v)
-        return Vase(contents)
+        return cls(contents)
 
-    def _to_sorted_key(self, tup):
+    def _to_sorted_key(self, tup: Sequence[str]) -> str:
         """
         Convert a collection of items into an order-independent key.
 
@@ -314,7 +321,7 @@ def _to_sorted_key(self, tup):
         """
         return "".join(sorted(tup))
 
-    def take(self, n=1, replace=False, ordered=False):
+    def take(self, n: int = 1, replace: bool = False, ordered: bool = False) -> "Dice[str]":
         """
         Calculate the distribution of drawing items from the vase.
 
@@ -335,7 +342,7 @@ def take(self, n=1, replace=False, ordered=False):
 
 
 
-def p(expression):
+def p(expression: "Dice[bool]") -> float:
     """
     Returns the probability of a True outcome from a dice expression.
 
@@ -348,7 +355,7 @@ def p(expression):
     return expression.probs.get(True, 0)
 
 
-def exp(dice):
+def exp(dice: "Dice[int | float]") -> float:
     """
     Calculates the expected value (mean) of a dice.
 
@@ -361,7 +368,7 @@ def exp(dice):
     return sum(i * p for i, p in dice.probs.items())
 
 
-def var(dice):
+def var(dice: "Dice[int | float]") -> float:
     """
     Calculates the variance of a dice.
 
@@ -373,7 +380,7 @@ def var(dice):
     """
     return sum(p * (i - exp(dice)) ** 2 for i, p in dice.probs.items())
 
-def mix(*dice, weights=None):
+def mix[T: Hashable](*dice: "Dice[T]", weights: Sequence[Weight] | None = None) -> "Dice[T]":
     """
     Create a weighted mixture of dice.
 
@@ -400,7 +407,7 @@ def mix(*dice, weights=None):
     if total_weight == 0:
         raise ValueError("weights must have a positive sum")
 
-    new_probs = {}
+    new_probs: dict[T, float] = {}
     for die, weight in zip(dice, weights):
         for outcome, probability in die.probs.items():
             contribution = probability * weight / total_weight
@@ -408,21 +415,21 @@ def mix(*dice, weights=None):
     return Dice(new_probs)
 
 
-def _slice_limit(length, stop):
+def _slice_limit(length: int, stop: int) -> int:
     return len(range(length)[:stop])
 
 
-def _threshold_probabilities(die, thresholds):
+def _threshold_probabilities[T: Hashable](die: "Dice[T]", thresholds: Sequence[T]) -> dict[T, float]:
     """
     Return P(die >= threshold) for each threshold.
     """
-    probabilities = {}
+    probabilities: dict[T, float] = {}
     cumulative = 0
-    items = sorted(die.probs.items(), reverse=True)
+    items = sorted(cast(list[Any], list(die.probs.items())), reverse=True)
     item_index = 0
 
     for threshold in thresholds:
-        while item_index < len(items) and items[item_index][0] >= threshold:
+        while item_index < len(items) and items[item_index][0] >= cast(Any, threshold):
             cumulative += items[item_index][1]
             item_index += 1
         probabilities[threshold] = cumulative
@@ -430,7 +437,7 @@ def _threshold_probabilities(die, thresholds):
     return probabilities
 
 
-def ordered(*dice_in, k=None):
+def ordered[T: Hashable](*dice_in: "Dice[T]", k: int | None = None) -> list["Dice[T]"]:
     """
     Return dice that represent order statistics. Highest first.
 
@@ -451,25 +458,22 @@ def ordered(*dice_in, k=None):
     if n_outputs == 0:
         return []
 
-    thresholds = sorted(
-        {outcome for die in dice_in for outcome in die.probs},
-        reverse=True,
-    )
+    thresholds = sorted(cast(list[Any], list({outcome for die in dice_in for outcome in die.probs})), reverse=True)
     probabilities = [
         _threshold_probabilities(die, thresholds)
         for die in dice_in
     ]
 
-    dice_out = [{} for _ in range(n_outputs)]
-    higher_threshold_survival = [0] * (n_outputs + 1)
+    dice_out: list[dict[T, float]] = [{} for _ in range(n_outputs)]
+    higher_threshold_survival: list[float] = [0.0] * (n_outputs + 1)
 
     for threshold in thresholds:
-        pass_count_probs = [1] + [0] * n_outputs
+        pass_count_probs: list[float] = [1.0] + [0.0] * n_outputs
 
         for probability in probabilities:
             p_ge = probability[threshold]
             p_lt = 1 - p_ge
-            new_pass_count_probs = [0] * (n_outputs + 1)
+            new_pass_count_probs: list[float] = [0.0] * (n_outputs + 1)
 
             for pass_count, pass_count_probability in enumerate(pass_count_probs):
                 if pass_count_probability == 0:
@@ -482,8 +486,8 @@ def ordered(*dice_in, k=None):
 
             pass_count_probs = new_pass_count_probs
 
-        survival = [0] * (n_outputs + 1)
-        running = 0
+        survival: list[float] = [0.0] * (n_outputs + 1)
+        running = 0.0
         for pass_count in range(n_outputs, -1, -1):
             running += pass_count_probs[pass_count]
             survival[pass_count] = running
@@ -492,7 +496,7 @@ def ordered(*dice_in, k=None):
             order = order_index + 1
             probability = survival[order] - higher_threshold_survival[order]
             if probability < 0 and abs(probability) < 1e-15:
-                probability = 0
+                probability = 0.0
             if probability != 0:
                 dice_out[order_index][threshold] = probability
 
diff --git a/dicekit/py.typed b/dicekit/py.typed
new file mode 100644
index 0000000..8b13789
--- /dev/null
+++ b/dicekit/py.typed
@@ -0,0 +1 @@
+
diff --git a/nbs/__init__.py b/nbs/__init__.py
index 0e77d00..176e414 100644
--- a/nbs/__init__.py
+++ b/nbs/__init__.py
@@ -252,13 +252,15 @@ def _():
     import marimo as mo
     from hastyplot import qplot
     import random
-    from collections.abc import Mapping
+    from collections.abc import Callable, Hashable, Mapping, Sequence
     from collections import Counter
-    from itertools import product, combinations, permutations, combinations_with_replacement
+    from itertools import product, permutations
     from functools import reduce
+    from typing import Any, Self, TypeAlias, cast
 
+    Weight: TypeAlias = int | float
 
-    class Dice:
+    class Dice[T: Hashable]:
         """
         A class representing dice with arbitrary probability distributions.
 
@@ -266,7 +268,7 @@ class Dice:
         custom probability distributions for each face value.
         """
 
-        def __init__(self, probs):
+        def __init__(self, probs: Mapping[T, Weight]) -> None:
             """
             Initialize a Dice object with given probabilities.
 
@@ -276,7 +278,7 @@ def __init__(self, probs):
             self.probs = {k: v / sum(probs.values()) for k, v in probs.items()}
 
         @classmethod
-        def from_sides(cls, n=6):
+        def from_sides(cls, n: int = 6) -> "Dice[int]":
             """
             Create a fair dice with n sides.
 
@@ -286,10 +288,10 @@ def from_sides(cls, n=6):
             Returns:
                 Dice: A fair dice with n sides
             """
-            return cls({i: 1 / n for i in range(1, n + 1)})
+            return cast(Dice[int], cast(Any, cls)({i: 1 / n for i in range(1, n + 1)}))
 
         @classmethod
-        def from_numbers(cls, *args):
+        def from_numbers[U: Hashable](cls, *args: U) -> "Dice[U]":
             """
             Create a dice from a sequence of numbers with frequencies.
 
@@ -300,9 +302,9 @@ def from_numbers(cls, *args):
                 Dice: A dice with probabilities based on the frequency of each number
             """
             c = Counter(args)
-            return cls({k: v / len(args) for k, v in args.items()})
+            return cast(Dice[U], cast(Any, cls)({k: v / len(args) for k, v in c.items()}))
 
-        def roll(self, n=1):
+        def roll(self, n: int = 1) -> list[T]:
             """
             Simulate rolling the dice n times.
 
@@ -314,7 +316,7 @@ def roll(self, n=1):
             """
             return random.choices(list(self.probs.keys()), weights=list(self.probs.values()), k=n)
 
-        def sample(self, n=1):
+        def sample(self, n: int = 1) -> list[T]:
             """
             Simulate rolling the dice n times.
 
@@ -326,7 +328,7 @@ def sample(self, n=1):
             """
             return self.roll(n=1)
 
-        def map(self, transform):
+        def map[U: Hashable](self, transform: Callable[[T], U] | Mapping[T, U]) -> "Dice[U]":
             """
             Transform the outcomes of the dice.
 
@@ -343,7 +345,11 @@ def map(self, transform):
                 new_probs[new_outcome] = new_probs.get(new_outcome, 0) + probability
             return Dice(new_probs)
 
-        def operate(self, other, operator):
+        def operate[U: Hashable, V: Hashable](
+            self,
+            other: "Dice[U] | U",
+            operator: Callable[[T, U], V],
+        ) -> "Dice[V]":
             """
             Apply an operation between this dice and another dice or number.
 
@@ -359,7 +365,8 @@ def operate(self, other, operator):
                     other = Dice({other: 1})
                 except TypeError as exc:
                     raise TypeError("operations require a Dice or hashable value") from exc
-            new_probs = {}
+            other = cast(Dice[U], other)
+            new_probs: dict[V, float] = {}
             for s1, p1 in self.probs.items():
                 for s2, p2 in other.probs.items():
                     new_key = operator(s1, s2)
@@ -368,7 +375,7 @@ def operate(self, other, operator):
                     new_probs[new_key] += p1 * p2
             return Dice(new_probs)
 
-        def filter(self, func):
+        def filter(self, func: Callable[[T], bool]) -> "Dice[T]":
             """
             Create a new dice by filtering face values based on a function.
 
@@ -382,7 +389,7 @@ def filter(self, func):
             total_prob = sum(new_probs.values())
             return Dice({k: v / total_prob for k, v in new_probs.items()})
 
-        def _repr_html_(self):
+        def _repr_html_(self) -> str:
             """
             Return HTML representation of the dice for display in notebooks.
 
@@ -391,7 +398,7 @@ def _repr_html_(self):
             """
             return mo.as_html(self.prob_chart()).text
 
-        def prob_chart(self):
+        def prob_chart(self) -> Any:
             """
             Create a visualization of the dice probability distribution.
 
@@ -411,7 +418,7 @@ def prob_chart(self):
                 theme="default",
             )
 
-        def cdf_chart(self):
+        def cdf_chart(self) -> Any:
             """
             Create a visualization of the cumulative distribution function.
 
@@ -436,7 +443,7 @@ def cdf_chart(self):
                 theme="default",
             )
 
-        def ordered(self, n, k=None):
+        def ordered(self, n: int, k: int | None = None) -> list["Dice[T]"]:
             """
             Take the dice `n` times and calculate the ordered distributions for it.
 
@@ -450,7 +457,7 @@ def ordered(self, n, k=None):
                 return ordered(*items, k=k)
             return ordered(*items)
 
-        def out_of(self, n=2, func=max):
+        def out_of[U: Hashable](self, n: int = 2, func: Callable[[list[T]], U] = max) -> "Dice[U]":
             """
             Create a dice representing the result of applying a function to n rolls.
 
@@ -461,7 +468,7 @@ def out_of(self, n=2, func=max):
             Returns:
                 Dice: A new dice representing the distribution of the function's results
             """
-            result = {}
+            result: dict[U, float] = {}
             dice_in = [self] * n
             for _i in product(*[d.probs.items() for d in dice_in]):
                 values = [_[0] for _ in _i]
@@ -472,54 +479,54 @@ def out_of(self, n=2, func=max):
                 result[outcome] += prob
             return Dice(result)
 
-        def __add__(self, other):
+        def __add__(self, other: Any) -> "Dice[Any] | Self":
             if not isinstance(other, Dice) and other == 0:
                 return self
             return self.operate(other, lambda a, b: a + b)
 
-        def __radd__(self, other):
+        def __radd__(self, other: Any) -> "Dice[Any] | Self":
             if not isinstance(other, Dice) and other == 0:
                 return self
             return self.operate(other, lambda a, b: a + b)
 
-        def __sub__(self, other):
+        def __sub__(self, other: Any) -> "Dice[Any]":
             return self.operate(other, lambda a, b: a - b)
 
-        def __rsub__(self, other):
+        def __rsub__(self, other: Any) -> "Dice[Any]":
             return self.operate(other, lambda a, b: b - a)
 
-        def __mul__(self, other):
+        def __mul__(self, other: Any) -> "Dice[Any]":
             return self.operate(other, lambda a, b: a * b)
 
-        def __rmul__(self, other):
+        def __rmul__(self, other: Any) -> "Dice[Any]":
             return self.operate(other, lambda a, b: a * b)
 
-        def __le__(self, other):
+        def __le__(self, other: Any) -> "Dice[bool]":
             return self.operate(other, lambda a, b: a <= b)
 
-        def __lt__(self, other):
+        def __lt__(self, other: Any) -> "Dice[bool]":
             return self.operate(other, lambda a, b: a < b)
 
-        def __ge__(self, other):
+        def __ge__(self, other: Any) -> "Dice[bool]":
             return self.operate(other, lambda a, b: a >= b)
 
-        def __gt__(self, other):
+        def __gt__(self, other: Any) -> "Dice[bool]":
             return self.operate(other, lambda a, b: a > b)
 
-        def __eq__(self, other):
+        def __eq__(self, other: Any) -> "Dice[bool]":  # type: ignore[override]
             return self.operate(other, lambda a, b: a == b)
 
-        def __ne__(self, other):
+        def __ne__(self, other: Any) -> "Dice[bool]":  # type: ignore[override]
             return self.operate(other, lambda a, b: a != b)
 
-        def __len__(self):
+        def __len__(self) -> int:
             return len(self.probs)
 
-    return Counter, Dice, mo, permutations, product, random, reduce
+    return Any, Callable, Counter, Dice, Hashable, Mapping, Sequence, Weight, cast, mo, permutations, product, random, reduce
 
 
 @app.cell
-def _(Counter, Dice, permutations, product):
+def _(Counter, Dice, Sequence, permutations, product):
     ## Export
 
 
@@ -531,7 +538,7 @@ class Vase:
         replacement, where the order of drawn items may optionally matter.
         """
 
-        def __init__(self, contents):
+        def __init__(self, contents: Sequence[str]) -> None:
             """
             Initialize a vase with the items it contains.
 
@@ -541,7 +548,7 @@ def __init__(self, contents):
             self._contents = contents
 
         @classmethod
-        def from_counts(self, **kwargs):
+        def from_counts(cls, **kwargs: int) -> "Vase":
             """
             Create a vase from item names and their quantities.
 
@@ -554,9 +561,9 @@ def from_counts(self, **kwargs):
             contents = []
             for k, v in kwargs.items():
                 contents.extend([k]*v)
-            return Vase(contents)
+            return cls(contents)
 
-        def _to_sorted_key(self, tup):
+        def _to_sorted_key(self, tup: Sequence[str]) -> str:
             """
             Convert a collection of items into an order-independent key.
 
@@ -568,7 +575,7 @@ def _to_sorted_key(self, tup):
             """
             return "".join(sorted(tup))
 
-        def take(self, n=1, replace=False, ordered=False):
+        def take(self, n: int = 1, replace: bool = False, ordered: bool = False) -> "Dice[str]":
             """
             Calculate the distribution of drawing items from the vase.
 
@@ -591,11 +598,11 @@ def take(self, n=1, replace=False, ordered=False):
 
 
 @app.cell
-def _(Dice, product, reduce):
+def _(Any, Callable, Dice, Hashable, Sequence, Weight, cast, product, reduce):
     ## Export
 
 
-    def p(expression):
+    def p(expression: "Dice[bool]") -> float:
         """
         Returns the probability of a True outcome from a dice expression.
 
@@ -608,7 +615,7 @@ def p(expression):
         return expression.probs.get(True, 0)
 
 
-    def exp(dice):
+    def exp(dice: "Dice[int | float]") -> float:
         """
         Calculates the expected value (mean) of a dice.
 
@@ -621,7 +628,7 @@ def exp(dice):
         return sum(i * p for i, p in dice.probs.items())
 
 
-    def var(dice):
+    def var(dice: "Dice[int | float]") -> float:
         """
         Calculates the variance of a dice.
 
@@ -633,7 +640,7 @@ def var(dice):
         """
         return sum(p * (i - exp(dice)) ** 2 for i, p in dice.probs.items())
 
-    def mix(*dice, weights=None):
+    def mix[T: Hashable](*dice: "Dice[T]", weights: Sequence[Weight] | None = None) -> "Dice[T]":
         """
         Create a weighted mixture of dice.
 
@@ -660,7 +667,7 @@ def mix(*dice, weights=None):
         if total_weight == 0:
             raise ValueError("weights must have a positive sum")
 
-        new_probs = {}
+        new_probs: dict[T, float] = {}
         for die, weight in zip(dice, weights):
             for outcome, probability in die.probs.items():
                 contribution = probability * weight / total_weight
@@ -668,21 +675,21 @@ def mix(*dice, weights=None):
         return Dice(new_probs)
 
 
-    def _slice_limit(length, stop):
+    def _slice_limit(length: int, stop: int) -> int:
         return len(range(length)[:stop])
 
 
-    def _threshold_probabilities(die, thresholds):
+    def _threshold_probabilities[T: Hashable](die: "Dice[T]", thresholds: Sequence[T]) -> dict[T, float]:
         """
         Return P(die >= threshold) for each threshold.
         """
-        probabilities = {}
+        probabilities: dict[T, float] = {}
         cumulative = 0
-        items = sorted(die.probs.items(), reverse=True)
+        items = sorted(cast(list[Any], list(die.probs.items())), reverse=True)
         item_index = 0
 
         for threshold in thresholds:
-            while item_index < len(items) and items[item_index][0] >= threshold:
+            while item_index < len(items) and items[item_index][0] >= cast(Any, threshold):
                 cumulative += items[item_index][1]
                 item_index += 1
             probabilities[threshold] = cumulative
@@ -690,7 +697,7 @@ def _threshold_probabilities(die, thresholds):
         return probabilities
 
 
-    def ordered(*dice_in, k=None):
+    def ordered[T: Hashable](*dice_in: "Dice[T]", k: int | None = None) -> list["Dice[T]"]:
         """
         Return dice that represent order statistics. Highest first.
 
@@ -711,25 +718,22 @@ def ordered(*dice_in, k=None):
         if n_outputs == 0:
             return []
 
-        thresholds = sorted(
-            {outcome for die in dice_in for outcome in die.probs},
-            reverse=True,
-        )
+        thresholds = sorted(cast(list[Any], list({outcome for die in dice_in for outcome in die.probs})), reverse=True)
         probabilities = [
             _threshold_probabilities(die, thresholds)
             for die in dice_in
         ]
 
-        dice_out = [{} for _ in range(n_outputs)]
-        higher_threshold_survival = [0] * (n_outputs + 1)
+        dice_out: list[dict[T, float]] = [{} for _ in range(n_outputs)]
+        higher_threshold_survival: list[float] = [0.0] * (n_outputs + 1)
 
         for threshold in thresholds:
-            pass_count_probs = [1] + [0] * n_outputs
+            pass_count_probs: list[float] = [1.0] + [0.0] * n_outputs
 
             for probability in probabilities:
                 p_ge = probability[threshold]
                 p_lt = 1 - p_ge
-                new_pass_count_probs = [0] * (n_outputs + 1)
+                new_pass_count_probs: list[float] = [0.0] * (n_outputs + 1)
 
                 for pass_count, pass_count_probability in enumerate(pass_count_probs):
                     if pass_count_probability == 0:
@@ -742,8 +746,8 @@ def ordered(*dice_in, k=None):
 
                 pass_count_probs = new_pass_count_probs
 
-            survival = [0] * (n_outputs + 1)
-            running = 0
+            survival: list[float] = [0.0] * (n_outputs + 1)
+            running = 0.0
             for pass_count in range(n_outputs, -1, -1):
                 running += pass_count_probs[pass_count]
                 survival[pass_count] = running
@@ -752,7 +756,7 @@ def ordered(*dice_in, k=None):
                 order = order_index + 1
                 probability = survival[order] - higher_threshold_survival[order]
                 if probability < 0 and abs(probability) < 1e-15:
-                    probability = 0
+                    probability = 0.0
                 if probability != 0:
                     dice_out[order_index][threshold] = probability
 
diff --git a/pyproject.toml b/pyproject.toml
index 9b5f6cf..5ef9b94 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -30,6 +30,7 @@ dependencies = [
 
 [dependency-groups]
 dev = [
+    "basedpyright",
     "pytest",
 ]
 
@@ -46,3 +47,15 @@ exclude = [
     "notebooks",
     "tests",
 ]
+
+[tool.setuptools.package-data]
+dicekit = ["py.typed"]
+
+[tool.basedpyright]
+include = ["dicekit"]
+exclude = ["nbs"]
+reportMissingTypeStubs = false
+typeCheckingMode = "basic"
+
+[tool.marimo.language_servers.basedpyright]
+enabled = true
diff --git a/tests/test_basics.py b/tests/test_basics.py
index 67b79d1..713ab40 100644
--- a/tests/test_basics.py
+++ b/tests/test_basics.py
@@ -54,6 +54,11 @@ def test_dice_creation():
     assert len(custom.probs) == 2
     assert custom.probs[1] == 0.5
 
+def test_dice_from_numbers():
+    dice = Dice.from_numbers(1, 1, 2)
+
+    assert dice.probs == {1: pytest.approx(2 / 3), 2: pytest.approx(1 / 3)}
+
 def test_dice_operations():
     d1 = Dice({1: 0.5, 2: 0.5})
     d2 = Dice({1: 0.5, 2: 0.5})
diff --git a/uv.lock b/uv.lock
index d328009..45fabe2 100644
--- a/uv.lock
+++ b/uv.lock
@@ -40,6 +40,18 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/6a/21/5b6702a7f963e95456c0de2d495f67bf5fd62840ac655dc451586d23d39a/attrs-24.2.0-py3-none-any.whl", hash = "sha256:81921eb96de3191c8258c199618104dd27ac608d9366f5e35d011eae1867ede2", size = 63001, upload-time = "2024-08-06T14:37:36.958Z" },
 ]
 
+[[package]]
+name = "basedpyright"
+version = "1.39.8"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "nodejs-wheel-binaries" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/d2/62/8550c75850b2185df984d1de437b4805b039ba856cacbee2966236203133/basedpyright-1.39.8.tar.gz", hash = "sha256:bb1a86d4d71425d52d1501b317fe23d45527baed06bd5d5e1a07cd4b60d07b55", size = 25488230, upload-time = "2026-06-14T09:05:30.177Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/f5/94/878454aefe94328ba7ad808ecd63da8311aae1198da46cfb29f5cfe130a8/basedpyright-1.39.8-py3-none-any.whl", hash = "sha256:a79d89928064bd9023d429b50c625d87d023bacc2fe3932ef6c7bd13b5426048", size = 13182726, upload-time = "2026-06-14T09:05:26.368Z" },
+]
+
 [[package]]
 name = "click"
 version = "8.1.7"
@@ -73,6 +85,7 @@ dependencies = [
 
 [package.dev-dependencies]
 dev = [
+    { name = "basedpyright" },
     { name = "pytest" },
 ]
 
@@ -84,7 +97,10 @@ requires-dist = [
 ]
 
 [package.metadata.requires-dev]
-dev = [{ name = "pytest" }]
+dev = [
+    { name = "basedpyright" },
+    { name = "pytest" },
+]
 
 [[package]]
 name = "docutils"
@@ -257,6 +273,22 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/ad/18/3cf5f8b188424313833f085f4ae0452b7db858a687281e2e2ca893a296cc/narwhals-1.13.5-py3-none-any.whl", hash = "sha256:91fe95ffdece9e3837780b6cd32f4309a41f39b285bc9d42d60eaff47d48b39a", size = 208096, upload-time = "2024-11-13T12:42:40.878Z" },
 ]
 
+[[package]]
+name = "nodejs-wheel-binaries"
+version = "24.16.0"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/a3/22/2a5beb4e21417c73233d9f65cf6f3e96e891b80d2f550a8f630ebc6b88c6/nodejs_wheel_binaries-24.16.0.tar.gz", hash = "sha256:c973cb69dc5fd16e6f6dc6e579e2c3d5534e2a1f57619dddf5ba070efa7dde37", size = 8056, upload-time = "2026-05-30T16:52:09.807Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/83/d1/68b43b53cd0fa83ae6fd406705023ca988d9e0ca41c724d82e66fbeb2ef6/nodejs_wheel_binaries-24.16.0-py2.py3-none-macosx_13_0_arm64.whl", hash = "sha256:d9f8f677dcf30e37ac244f07869726abe043f01eb0f45722b1df31cc2af7093c", size = 55666374, upload-time = "2026-05-30T16:51:39.588Z" },
+    { url = "https://files.pythonhosted.org/packages/e9/b2/40a989159599080da485de966c4c2d207e852ac7aa7864702626d96c8bf5/nodejs_wheel_binaries-24.16.0-py2.py3-none-macosx_13_0_x86_64.whl", hash = "sha256:3d0370fe7120ce9697a4f60d40480d2bd8808d9f30131458d5afc0040d4e5a51", size = 55838487, upload-time = "2026-05-30T16:51:43.383Z" },
+    { url = "https://files.pythonhosted.org/packages/d7/a7/cd42174fb5ff6faff7fa8d326a18914d8f232098ab5de055b57c16fa13ca/nodejs_wheel_binaries-24.16.0-py2.py3-none-manylinux_2_28_aarch64.whl", hash = "sha256:85dc92bbb79c851569c5925dcc2a4c915a034efab375f99e4e7e6bbe9cca8342", size = 60179540, upload-time = "2026-05-30T16:51:47.036Z" },
+    { url = "https://files.pythonhosted.org/packages/2b/95/c8a1f9ae140aa28df8744d984d01d4b3af7cdd6555af12127f40ceb45a7d/nodejs_wheel_binaries-24.16.0-py2.py3-none-manylinux_2_28_x86_64.whl", hash = "sha256:2f3036292811514ba847b3708492644764f88a833ac425c5f55007014308ddfd", size = 60716262, upload-time = "2026-05-30T16:51:50.711Z" },
+    { url = "https://files.pythonhosted.org/packages/64/c9/7c35b3737f59e36d0249c265397b7bff570519b95301d6e16ea361e904ad/nodejs_wheel_binaries-24.16.0-py2.py3-none-musllinux_1_2_aarch64.whl", hash = "sha256:db8a8a76ebd2b28ecbfc9ad464baa3707241b9e050a30e2efdf6f60c0f886502", size = 62230592, upload-time = "2026-05-30T16:51:55Z" },
+    { url = "https://files.pythonhosted.org/packages/04/96/d931255cf9d11a84d6b54d882dba7434646467d568ccf070ea3418638df3/nodejs_wheel_binaries-24.16.0-py2.py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:f1a3d8f7b4491cbbd023ba3fc4e901fcca2d9fb80d57f24ba3890de8b1dbac03", size = 62841759, upload-time = "2026-05-30T16:51:59.407Z" },
+    { url = "https://files.pythonhosted.org/packages/a2/7b/8b7a3f41bc255411be30b6d7d288aab8ffd9ea2055db8555ced3548007b9/nodejs_wheel_binaries-24.16.0-py2.py3-none-win_amd64.whl", hash = "sha256:bb136be9944f0662dcf1120f45193a6b75b13fac378971a95cc42c9f879a81aa", size = 42027734, upload-time = "2026-05-30T16:52:03.348Z" },
+    { url = "https://files.pythonhosted.org/packages/17/66/1ed71f1f529b8ca727d42c7ceb9db0bef145ce4a13dfc86fb50aa44f3be6/nodejs_wheel_binaries-24.16.0-py2.py3-none-win_arm64.whl", hash = "sha256:8308940b5edd0a50dc5267ea36ba21c9f668e83fe0d9f293937174d3a7e31c36", size = 39714528, upload-time = "2026-05-30T16:52:06.421Z" },
+]
+
 [[package]]
 name = "numpy"
 version = "2.1.3"