Improvements/fixes to type annotations

pyrankvote/helpers.py

@@ -4,9 +4,10 @@
 """
 from pyrankvote.models import Candidate, Ballot
 
+from enum import Enum
 import random
 import functools
-from typing import List, NamedTuple
+from typing import List, NamedTuple, Optional
 from tabulate import tabulate
 
 
@@ -17,7 +18,7 @@ def almost_equal(value1: float, value2: float) -> bool:
     return abs(value1 - value2) < CONSIDERED_EQUAL_MARGIN
 
 
-class CandidateStatus:
+class CandidateStatus(Enum):
     Elected = "Elected"
     Hopeful = "Hopeful"
     Rejected = "Rejected"
@@ -134,7 +135,7 @@ def __init__(
     ):
 
         self._ballots = ballots
-        self._candidate_vote_counts: dict[Candidate:CandidateVoteCount] = {
+        self._candidate_vote_counts: dict[Candidate, CandidateVoteCount] = {
             candidate: CandidateVoteCount(candidate) for candidate in candidates
         }
 
@@ -331,7 +332,7 @@ def get_results(self) -> RoundResult:
     # INTERNAL METHODS
     def _get_ballot_candidate_nr_x_in_race_or_none(
         self, ballot: Ballot, x: int
-    ) -> Candidate:
+    ) -> Optional[Candidate]:
         ranked_candidates_in_race = [
             candidate
             for candidate in ballot.ranked_candidates

pyrankvote/models.py

@@ -3,7 +3,7 @@
 
 You can create and use your own Candidate and Ballot models as long as they implement the same properties and methods.
 """
-from typing import List
+from typing import List, Tuple
 
 
 class Candidate:
@@ -40,7 +40,7 @@ class Ballot:
     """
 
     def __init__(self, ranked_candidates: List[Candidate]):
-        self.ranked_candidates: List[Candidate] = tuple(ranked_candidates)
+        self.ranked_candidates: Tuple[Candidate, ...] = tuple(ranked_candidates)
 
         if Ballot._is_duplicates(ranked_candidates):
             raise DuplicateCandidatesError