Feat/allow int candidates

src/votekit/ballot.py

@@ -1,35 +1,45 @@
 from __future__ import annotations
 
+import warnings
 from numbers import Real
-from typing import Iterable, Optional, Sequence, TypeAlias, Union, overload
+from typing import Iterable, Mapping, Optional, Sequence, TypeAlias, Union, overload
 
-Ranking: TypeAlias = Optional[tuple[frozenset[str], ...]]
-RankingLike: TypeAlias = Optional[Sequence[str | Iterable[str]]]
+from votekit.types import Candidate
+
+Ranking: TypeAlias = Optional[tuple[frozenset[Candidate], ...]]
+RankingLike: TypeAlias = Optional[Sequence[Candidate | Iterable[Candidate]]]
+ScoresLike: TypeAlias = Optional[
+    Mapping[Candidate, float | int] | Mapping[str, float | int] | Mapping[int, float | int]
+]
 
 
 class Ballot:
     """
     Ballot parent class, contains voter set and assigned weight.
 
     Args:
-        ranking (Optional[Sequence[str | Iterable[str]]]): Candidate ranking.
+        ranking (Optional[Sequence[Candidate | Iterable[Candidate]]]): Candidate ranking.
             Entry i of the sequence is a candidate or iterable of candidates ranked in position i.
-            Defaults to None. Will be coerced to tuple[frozenset[str], ...].
+            Candidates can be strings, integers, or mix of both.
+            Defaults to None. Will be coerced to tuple[frozenset[str | int], ...].
         weight (Union[float, int]): Weight assigned to a given ballot. Defaults to 1.0
             Can be input as int or float, and will be coerced to float.
         voter_set (Union[set[str], frozenset[str]]): Set of voters who cast the ballot.
             Defaults to frozenset(). Will be coerced to frozenset.
-        scores (Optional[dict[str, Union[int, float]]): Scores for individual candidates.
-            Defaults to None. Values can be input as int or float but will be coerced to float.
+        scores (Optional[Mapping[Candidate, float | int] | Mapping[str, float | int]
+            | Mapping[int, float | int]]): Scores for individual candidates. Defaults to None.
+            Values can be input as int or float but will be coerced to float.
+            Candidates can be strings, integers, or mix of both.
+            Stored internally as a dict[str | int, float].
             Only retains non-zero scores.
 
     Attributes:
-        ranking (Optional[tuple[frozenset[str], ...]]): Tuple of candidate ranking.
-            Entry i of the tuple is a
-            frozenset of candidates ranked in position i.
+        ranking (Optional[tuple[frozenset[Candidate], ...]]): Tuple of candidate ranking.
+            Entry i of the tuple is a frozenset of candidates ranked in position i.
+            Candidates can be strings, integers, or mix of both.
         weight (float): Weight assigned to a given ballot.
         voter_set (frozenset[str]): Set of voters who cast the ballot.
-        scores (Optional[dict[str, float]]): Scores for individual candidates.
+        scores (Optional[Mapping[Candidate, float | int]): Scores for individual candidates.
 
     Raises:
         TypeError: Only one of ranking or scores can be provided.
@@ -49,7 +59,7 @@ class Ballot:
     def __new__(
         cls,
         *,
-        ranking: Sequence[str | Iterable[str]],
+        ranking: RankingLike,
         scores: None = None,
         weight: Union[float, int] = 1.0,
         voter_set: Union[set[str], frozenset[str]] = frozenset(),
@@ -60,7 +70,7 @@ def __new__(
         cls,
         *,
         ranking: None = None,
-        scores: dict[str, Union[int, float]],
+        scores: ScoresLike,
         weight: Union[float, int] = 1.0,
         voter_set: Union[set[str], frozenset[str]] = frozenset(),
     ) -> ScoreBallot: ...
@@ -69,17 +79,17 @@ def __new__(
     def __new__(
         cls,
         *,
-        ranking: Optional[Sequence[str | Iterable[str]]] = None,
-        scores: Optional[dict[str, Union[int, float]]] = None,
+        ranking: RankingLike = None,
+        scores: ScoresLike = None,
         weight: Union[float, int] = 1.0,
         voter_set: Union[set[str], frozenset[str]] = frozenset(),
     ) -> Ballot: ...
 
     def __new__(
         cls,
         *,
-        ranking: Optional[Sequence[str | Iterable[str]]] = None,
-        scores: Optional[dict[str, Union[int, float]]] = None,
+        ranking: RankingLike = None,
+        scores: ScoresLike = None,
         weight: Union[float, int] = 1.0,
         voter_set: Union[set[str], frozenset[str]] = frozenset(),
     ):
@@ -95,8 +105,8 @@ def __new__(
     def __init__(
         self,
         *,
-        ranking: Optional[Sequence[str | Iterable[str]]] = None,
-        scores: Optional[dict[str, Union[int, float]]] = None,
+        ranking: RankingLike = None,
+        scores: ScoresLike = None,
         weight: Union[float, int] = 1.0,
         voter_set: Union[set[str], frozenset[str]] = frozenset(),
     ):
@@ -152,25 +162,29 @@ class RankBallot(Ballot):
 
     Args:
         ranking (RankingLike): Ranking of candidates, defaults to None.
+            RankingLike = Sequence[Candidate | Iterable[Candidate]] | None
+            Canidates can be strings, integers, or mix of both.
         weight (Union[int, float]): Weight of the ballot, defaults to 1.0.
         voter_set (Union[set[str], frozenset[str]]): Voter set of the ballot,
             defaults to frozenset().
 
     Attributes:
-        ranking (RankingLike): Ranking of candidates.
+        ranking (Ranking): Ranking of candidates.
+            Ranking = tuple[frozenset[Candidate], ...] | None
         weight (float): Weight of the ballot.
         voter_set (frozenset[str]): Voter set of the ballot.
 
     Raises:
         ValueError: Candidate '~' found in ballot ranking.
         ValueError: Ballot weight cannot be negative.
+        UserWarning: '1' and 1 candidates are treated as separate candidates.
     """
 
     def __init__(
         self,
         *,
         ranking: RankingLike = None,
-        scores: Optional[dict[str, Union[int, float]]] = None,
+        scores: ScoresLike = None,
         weight: Union[int, float] = 1.0,
         voter_set: Union[set[str], frozenset[str]] = frozenset(),
     ):
@@ -194,22 +208,40 @@ def _convert_ranking_candidates_to_frozenset_strip_whitespace(
 
         normalized_ranking = []
         for cand_set in ranking:
-            if isinstance(cand_set, str):
-                normalized_ranking.append(frozenset({cand_set.strip()}))
+            if isinstance(cand_set, Candidate):
+                normalized_ranking.append(
+                    frozenset({cand_set.strip() if isinstance(cand_set, str) else cand_set})
+                )
             else:
-                normalized_ranking.append(frozenset(c.strip() for c in cand_set))
+                normalized_ranking.append(
+                    frozenset(c.strip() if isinstance(c, str) else c for c in cand_set)
+                )
         return tuple(normalized_ranking)
 
     def _validate_ranking_candidates(self, ranking: Ranking):
         if ranking is None:
             return
-        if any(c == "~" for cand_set in ranking for c in cand_set):
+        if any(cand == "~" for cand_set in ranking for cand in cand_set):
             raise ValueError(
                 f"Candidate '~' found in ballot ranking {ranking}."
                 " '~' is a reserved character and cannot be used for"
                 " candidate names."
             )
 
+        str_cands = {cand for cand_set in ranking for cand in cand_set if isinstance(cand, str)}
+        int_cands = {cand for cand_set in ranking for cand in cand_set if isinstance(cand, int)}
+        collisions = {
+            str_cand
+            for str_cand in str_cands
+            if str_cand.lstrip("-").isdigit() and int(str_cand) in int_cands
+        }
+        if collisions:
+            warnings.warn(
+                f"Candidates {collisions} appear as both str and int (e.g. '1' and 1)."
+                " These will be treated as separate candidates.",
+                UserWarning,
+            )
+
     def __eq__(self, other):
         if not isinstance(other, RankBallot):
             return False
@@ -246,56 +278,76 @@ class ScoreBallot(Ballot):
     Class to handle ballots with scores. Strips whitespace from candidate names.
 
     Args:
-        scores (Optional[dict[str, Union[int, float]]]): Scores of candidates, defaults to None.
+        scores (ScoresLike): Scores of candidates, defaults to None.
+            ScoresLike = Mapping[Candidate, int | float] | Mapping[str, int | float]
+            | Mapping[int, int | float] | None
+            Candidates can be strings, integers, or mix of both.
         weight (Union[int, float]): Weight of the ballot, defaults to 1.0.
         voter_set (Union[set[str], frozenset[str]]): Voter set of the ballot,
             defaults to frozenset().
 
     Attributes:
-        scores (Optional[dict[str, float]]): Scores of candidates.
+        scores (Optional[dict[Candidate, float]]): Scores of candidates.
         weight (float): Weight of the ballot.
         voter_set (frozenset[str]): Voter set of the ballot.
 
     Raises:
         ValueError: Candidate '~' found in ballot scores.
         ValueError: Ballot weight cannot be negative.
         TypeError: Score values must be numeric.
+        UserWarning: '1' and 1 candidates are treated as separate candidates.
     """
 
     def __init__(
         self,
         *,
         ranking: RankingLike = None,
-        scores: Optional[dict[str, Union[int, float]]] = None,
+        scores: ScoresLike = None,
         weight: Union[int, float] = 1.0,
         voter_set: Union[set[str], frozenset[str]] = frozenset(),
     ):
         if ranking is not None:
             raise TypeError("Only one of ranking or scores can be provided.")
+        scores = self._convert_scores_to_float_strip_whitespace(scores)
         self._validate_scores_candidates(scores)
-        self.scores = self._convert_scores_to_float_strip_whitespace(scores)
+        self.scores = scores
 
         super().__init__(weight=weight, voter_set=voter_set)
 
-    def _validate_scores_candidates(self, scores: Optional[dict[str, Union[int, float]]]):
-        if scores is not None:
-            if "~" in scores:
-                raise ValueError(
-                    f"Candidate '~' found in ballot scores {list(scores.keys())}."
-                    " '~' is a reserved character and cannot be used for"
-                    " candidate names."
-                )
-
     def _convert_scores_to_float_strip_whitespace(
-        self, scores: Optional[dict[str, float]]
-    ) -> Optional[dict[str, float]]:
+        self, scores: ScoresLike
+    ) -> Optional[dict[Candidate, float]]:
         if scores is None:
             return None
 
         if any(not isinstance(s, Real) for s in scores.values()):
             raise TypeError("Score values must be numeric.")
 
-        return {c.strip(): float(s) for c, s in scores.items() if s != 0}
+        return {
+            c.strip() if isinstance(c, str) else c: float(s) for c, s in scores.items() if s != 0
+        }
+
+    def _validate_scores_candidates(self, scores: ScoresLike):
+        if scores is not None:
+            if "~" in scores:
+                raise ValueError(
+                    f"Candidate '~' found in ballot scores {list(scores.keys())}."
+                    " '~' is a reserved character and cannot be used for"
+                    " candidate names."
+                )
+            str_cands = {cand for cand in scores.keys() if isinstance(cand, str)}
+            int_cands = {cand for cand in scores.keys() if isinstance(cand, int)}
+            collisions = {
+                str_cand
+                for str_cand in str_cands
+                if str_cand.lstrip("-").isdigit() and int(str_cand) in int_cands
+            }
+            if collisions:
+                warnings.warn(
+                    f"Candidates {collisions} appear as both str and int (e.g. '1' and 1)."
+                    " These will be treated as separate candidates.",
+                    UserWarning,
+                )
 
     def __eq__(self, other):
         if not isinstance(other, ScoreBallot):
@@ -306,11 +358,7 @@ def __eq__(self, other):
 
     def __hash__(self):
         return (
-            hash(
-                tuple(sorted((c, s) for c, s in self.scores.items()))
-                if self.scores is not None
-                else self.scores
-            )
+            hash(frozenset(self.scores.items()) if self.scores is not None else self.scores)
             + super().__hash__()
         )
 

src/votekit/ballot_generator/bloc_slate_generator/name_bradley_terry.py

@@ -27,6 +27,7 @@
 from votekit.ballot_generator.bloc_slate_generator.config import BlocSlateConfig
 from votekit.ballot_generator.utils import system_memory
 from votekit.pref_profile import RankProfile
+from votekit.types import Candidate
 
 # ====================================================
 # ================= Helper Functions =================
@@ -227,7 +228,7 @@ def _inner_name_bradley_terry(config: BlocSlateConfig) -> dict[str, RankProfile]
 # - Other speed improvements
 def _bradley_terry_mcmc(
     n_ballots: int,
-    pref_interval: Mapping[str, float],
+    pref_interval: Mapping[Candidate, float],
     seed_ballot: RankBallot,
     verbose: bool = False,
     burn_in_time: int = 0,
@@ -239,7 +240,8 @@ def _bradley_terry_mcmc(
 
     Args:
         n_ballots (int): the number of ballots to sample
-        pref_interval (Mapping[str, float]): the preference interval to determine BT distribution
+        pref_interval (Mapping[Candidate, float]): the preference interval
+            to determine BT distribution. Candidate can be a str or int.
         seed_ballot (RankBallot):  the seed ballot for the Markov chain
         verbose (bool): If True, print the acceptance ratio of the chain. Defaults to False.
         burn_in_time (int): the number of ballots discarded in the beginning of the chain

src/votekit/ballot_generator/std_generator/impartial_anon_culture.py

@@ -13,6 +13,7 @@
 from typing import Optional, Sequence
 
 from votekit.pref_profile import RankProfile
+from votekit.types import Candidate
 from votekit.utils import build_df_from_ballot_samples, index_to_lexicographic_ballot
 
 # ====================================================
@@ -154,7 +155,7 @@ def _sample_anonymous_profile_ballot_counts(
 
 
 def iac_profile_generator(
-    candidates: Sequence[str],
+    candidates: Sequence[Candidate],
     number_of_ballots: int,
     max_ballot_length: Optional[int] = None,
 ) -> RankProfile:
@@ -163,7 +164,8 @@ def iac_profile_generator(
     is equally likely.
 
     Args:
-        candidates (Sequence[str]): List of candidate strings.
+        candidates (Sequence[Candidate]): List of candidate strings.
+            Candidates can be strings, integers, or mix of both.
         number_of_ballots (int): Number of ballots to generate.
         max_ballot_length (Optional[int]): Maximum length of each ballot. If None, defaults to
             the number of candidates.