covariant populations and generators

src/icepool/deal.py

@@ -9,17 +9,17 @@
 from functools import cached_property
 import math
 
-T = TypeVar('T', bound=Hashable)
+T_co = TypeVar('T_co', bound=Hashable, covariant=True)
 """Type variable representing the outcome type."""
 
 
-class Deal(OutcomeCountGenerator[T]):
+class Deal(OutcomeCountGenerator[T_co]):
     """EXPERIMENTAL: Represents an sorted/unordered deal of cards from a `Deck`. """
 
     _deck: 'icepool.Deck'
     _hand_sizes: tuple[int, ...]
 
-    def __init__(self, deck: 'icepool.Deck[T]', *hand_sizes: int):
+    def __init__(self, deck: 'icepool.Deck[T_co]', *hand_sizes: int):
         """Constructor.
 
         For algorithmic reasons, you must pre-commit to the number of cards to
@@ -41,7 +41,7 @@ def __init__(self, deck: 'icepool.Deck[T]', *hand_sizes: int):
                 'The total number of cards dealt cannot exceed the size of the deck.'
             )
 
-    def deck(self) -> 'icepool.Deck[T]':
+    def deck(self) -> 'icepool.Deck[T_co]':
         """The `Deck` the cards are dealt from."""
         return self._deck
 
@@ -53,7 +53,7 @@ def total_cards_dealt(self) -> int:
         """The total number of cards dealt."""
         return sum(self.hand_sizes())
 
-    def outcomes(self) -> CountsKeysView[T]:
+    def outcomes(self) -> CountsKeysView[T_co]:
         """The outcomes of the `Deck` in sorted order.
 
         These are also the `keys` of the `Deck` as a `Mapping`.
@@ -79,7 +79,7 @@ def _denomiator(self) -> int:
     def denominator(self) -> int:
         return self._denomiator
 
-    def _generate_common(self, popped_deck: 'icepool.Deck[T]',
+    def _generate_common(self, popped_deck: 'icepool.Deck[T_co]',
                          deck_count: int) -> NextOutcomeCountGenerator:
         """Common implementation for _generate_min and _generate_max."""
         min_count = max(
@@ -96,7 +96,7 @@ def _generate_common(self, popped_deck: 'icepool.Deck[T]',
                 weight = weight_total * weight_split
                 yield popped_deal, counts, weight
 
-    def _generate_min(self, min_outcome: T) -> NextOutcomeCountGenerator:
+    def _generate_min(self, min_outcome) -> NextOutcomeCountGenerator:
         if not self.outcomes() or min_outcome != self.min_outcome():
             yield self, (0,), 1
             return
@@ -105,7 +105,7 @@ def _generate_min(self, min_outcome: T) -> NextOutcomeCountGenerator:
 
         yield from self._generate_common(popped_deck, deck_count)
 
-    def _generate_max(self, max_outcome: T) -> NextOutcomeCountGenerator:
+    def _generate_max(self, max_outcome) -> NextOutcomeCountGenerator:
         if not self.outcomes() or max_outcome != self.max_outcome():
             yield self, (0,), 1
             return

src/icepool/deck.py

@@ -13,20 +13,20 @@
 
 from typing import Any, Callable, Hashable, Iterator, Mapping, MutableMapping, Sequence, TypeVar, overload
 
-T = TypeVar('T', bound=Hashable)
+T_co = TypeVar('T_co', bound=Hashable, covariant=True)
 """Type variable representing the outcome type."""
 
 U = TypeVar('U', bound=Hashable)
 """Type variable representing another outcome type."""
 
 
-class Deck(Population[T]):
+class Deck(Population[T_co]):
     """Sampling without replacement (within a single evaluation).
 
     Quantities represent duplicates.
     """
 
-    _data: Counts[T]
+    _data: Counts[T_co]
     _deal: int
 
     def _new_type(self) -> type:
@@ -87,7 +87,7 @@ def __new__(cls, outcomes, times: Sequence[int] | int = 1) -> 'Deck':
         return Deck._new_deck(data)
 
     @classmethod
-    def _new_deck(cls, data: Counts[T]) -> 'Deck[T]':
+    def _new_deck(cls, data: Counts[T_co]) -> 'Deck[T_co]':
         """Creates a new `Deck` using already-processed arguments.
 
         Args:
@@ -97,19 +97,19 @@ def _new_deck(cls, data: Counts[T]) -> 'Deck[T]':
         self._data = data
         return self
 
-    def keys(self) -> CountsKeysView[T]:
+    def keys(self) -> CountsKeysView[T_co]:
         return self._data.keys()
 
     def values(self) -> CountsValuesView:
         return self._data.values()
 
-    def items(self) -> CountsItemsView[T]:
+    def items(self) -> CountsItemsView[T_co]:
         return self._data.items()
 
     def __getitem__(self, outcome) -> int:
         return self._data[outcome]
 
-    def __iter__(self) -> Iterator[T]:
+    def __iter__(self) -> Iterator[T_co]:
         return iter(self.keys())
 
     def __len__(self) -> int:
@@ -118,30 +118,30 @@ def __len__(self) -> int:
     size = icepool.Population.denominator
 
     @cached_property
-    def _popped_min(self) -> tuple['Deck[T]', int]:
+    def _popped_min(self) -> tuple['Deck[T_co]', int]:
         return self._new_deck(self._data.remove_min()), self.quantities()[0]
 
-    def _pop_min(self) -> tuple['Deck[T]', int]:
+    def _pop_min(self) -> tuple['Deck[T_co]', int]:
         """A `Deck` with the min outcome removed."""
         return self._popped_min
 
     @cached_property
-    def _popped_max(self) -> tuple['Deck[T]', int]:
+    def _popped_max(self) -> tuple['Deck[T_co]', int]:
         return self._new_deck(self._data.remove_max()), self.quantities()[-1]
 
-    def _pop_max(self) -> tuple['Deck[T]', int]:
+    def _pop_max(self) -> tuple['Deck[T_co]', int]:
         """A `Deck` with the max outcome removed."""
         return self._popped_max
 
-    def deal(self, *hand_sizes: int) -> 'icepool.Deal[T]':
+    def deal(self, *hand_sizes: int) -> 'icepool.Deal[T_co]':
         """Creates a `Deal` object from this deck.
 
         See `Deal()` for details.
         """
         return icepool.Deal(self, *hand_sizes)
 
     def map(self,
-            repl: Callable[..., U] | Mapping[T, U],
+            repl: Callable[..., U] | Mapping[T_co, U],
             /,
             star: int = 0) -> 'Deck[U]':
         """Maps outcomes of this `Deck` to other outcomes.