diff --git a/shangrla/shangrla/Audit.py b/shangrla/shangrla/Audit.py index 2794b6e..f349ff8 100644 --- a/shangrla/shangrla/Audit.py +++ b/shangrla/shangrla/Audit.py @@ -165,7 +165,7 @@ def __init__(self, self.p = p # sampling probability self.sampled = sampled # is this CVR in the sample? - def __str__(self): + def __str__(self) -> str: return f'id: {str(self.id)} votes: {str(self.votes)} phantom: {str(self.phantom)} ' + \ f'tally_batch: {str(self.tally_batch)} pool: {str(self.pool)}' @@ -173,10 +173,10 @@ def get_vote_for(self, contest_id: str, candidate: str): return (False if (contest_id not in self.votes or candidate not in self.votes[contest_id]) else self.votes[contest_id][candidate]) - def has_contest(self, contest_id: str): + def has_contest(self, contest_id: str) -> bool: return contest_id in self.votes - def update_votes(self, votes: dict): + def update_votes(self, votes: dict) -> bool: ''' Update the votes for any contests the CVR already contains; add any contests and votes not already contained @@ -309,7 +309,13 @@ def from_dict(cls, cvr_dict: list[dict]) -> list: cvr_list = [] for c in cvr_dict: phantom = False if 'phantom' not in c.keys() else c['phantom'] - cvr_list.append(CVR(id = c['id'], votes = c['votes'], phantom=phantom)) + pool = None if 'pool' not in c.keys() else c['pool'] + tally_batch = None if 'tally_batch' not in c.keys() else c['tally_batch'] + sample_num = None if 'sample_num' not in c.keys() else c['sample_num'] + p = None if 'p' not in c.keys() else c['p'] + sampled = None if 'sampled' not in c.keys() else c['sampled'] + cvr_list.append(CVR(id=c['id'], votes=c['votes'], phantom=phantom, pool=pool, tally_batch=tally_batch, + sample_num=sample_num, p=p, sampled=sampled)) return cvr_list @classmethod @@ -434,7 +440,47 @@ def as_vote(cls, v) -> int: def as_rank(cls, v) -> int: return int(v) + @classmethod + def pool_contests(cls, cvrs: list["CVR"]) -> set: + ''' + create a set containing all contest ids in the list of CVRs + + Parameters + ---------- + cvrs : list of CVR objects + the set to collect contests from + Returns + ------- + a set containing the ID of every contest mentioned in the CVR list + ''' + contests = set() + for c in cvrs: + contests = contests.union(c.votes.keys()) + return contests + + @classmethod + def add_pool_contests(cls, cvrs: list["CVR"], pools: dict) -> bool: + ''' + for each pool, ensure every CVR in that pool has every contest in that pool + + Parameters + ---------- + cvrs : list of CVR objects + the set to update with additional contests as needed + + pools : dict + keys are pool ids, values are sets of contests every CVR in that pool should have + + Returns + ------- + bool : True if any contests are added + ''' + added = False + for c in cvrs: + added = c.update_votes({con: {} for con in pools[c.pool]}) or added # note: order of terms matters! + return added + @classmethod def make_phantoms(cls, audit: dict=None, contests: dict=None, cvr_list: list=None, prefix: str='phantom-') -> Tuple[list, int] : diff --git a/shangrla/shangrla/tests/test_CVR.py b/shangrla/shangrla/tests/test_CVR.py index 41c7017..1bef5c8 100644 --- a/shangrla/shangrla/tests/test_CVR.py +++ b/shangrla/shangrla/tests/test_CVR.py @@ -46,9 +46,10 @@ def test_rcv_votefor_cand(self): assert votes.rcv_votefor_cand("AvB", "Aaron", remaining) == 0 def test_cvr_from_dict(self): - cvr_dict = [{'id': 1, 'votes': {'AvB': {'Alice':True}, 'CvD': {'Candy':True}}}, - {'id': 2, 'votes': {'AvB': {'Bob':True}, 'CvD': {'Elvis':True, 'Candy':False}}}, - {'id': 3, 'votes': {'EvF': {'Bob':1, 'Edie':2}, 'CvD': {'Elvis':False, 'Candy':True}}}] + cvr_dict = [{'id': 1, 'pool': '1', 'votes': {'AvB': {'Alice':True}, 'CvD': {'Candy':True}}}, + {'id': 2, 'sample_num': 0.2, 'p': 0.5, 'sampled': True, + 'votes': {'AvB': {'Bob':True}, 'CvD': {'Elvis':True, 'Candy':False}}}, + {'id': 3, 'tally_batch': 'abc', 'votes': {'EvF': {'Bob':1, 'Edie':2}, 'CvD': {'Elvis':False, 'Candy':True}}}] cvr_list = CVR.from_dict(cvr_dict) assert len(cvr_list) == 3 assert cvr_list[0].id == 1 @@ -75,6 +76,12 @@ def test_cvr_from_dict(self): assert cvr_list[2].get_vote_for('EvF', 'Bob') == 1 assert cvr_list[2].get_vote_for('EvF', 'Edie') == 2 assert cvr_list[2].get_vote_for('EvF', 'Alice') == False + + assert cvr_list[0].pool == '1' + assert cvr_list[1].sample_num == 0.2 + assert cvr_list[1].p == 0.5 + assert cvr_list[1].sampled + assert cvr_list[2].tally_batch == 'abc' def test_cvr_has_contest(self): cvr_dict = [{'id': 1, 'votes': {'AvB': {}, 'CvD': {'Candy':True}}}, @@ -89,9 +96,9 @@ def test_cvr_has_contest(self): assert not cvr_list[1].has_contest('EvF') def test_cvr_add_votes(self): - cvr_dict = [{'id': 1, 'votes': {'AvB': {}, 'CvD': {'Candy':True}}}, - {'id': 2, 'votes': {'CvD': {'Elvis':True, 'Candy':False}}}] - cvr_list = CVR.from_dict(cvr_dict) + cvr_dicts = [{'id': 1, 'votes': {'AvB': {}, 'CvD': {'Candy':True}}}, + {'id': 2, 'votes': {'CvD': {'Elvis':True, 'Candy':False}}}] + cvr_list = CVR.from_dict(cvr_dicts) assert not cvr_list[0].has_contest('QvR') assert not cvr_list[1].has_contest('AvB') assert cvr_list[0].update_votes({'QvR': {}}) @@ -107,7 +114,35 @@ def test_cvr_add_votes(self): assert cvr_list[1].get_vote_for('CvD', 'Dan') == 7 assert cvr_list[1].get_vote_for('CvD', 'Candy') assert not cvr_list[1].get_vote_for('CvD', 'Elvis') - + + + def test_cvr_pool_contests(self): + cvr_dicts = [{'id': 1, 'sample_num': 1, 'votes': {'AvB': {}, 'CvD': {'Candy':True}}}, + {'id': 2, 'p': 0.5, 'votes': {'CvD': {'Elvis':True, 'Candy':False}, 'EvF': {}}}, + {'id': 3, 'tally_batch': 'abc', 'sampled': True, 'votes': {'GvH': {}}} + ] + cvr_list = CVR.from_dict(cvr_dicts) + assert CVR.pool_contests(cvr_list) == {'AvB', 'CvD', 'EvF', 'GvH'} + + def test_add_pool_contests(self): + cvr_dicts = [{'id': 1, 'pool': 1, 'votes': {'AvB': {}, 'CvD': {'Candy':True}}}, + {'id': 2, 'pool': 1, 'votes': {'CvD': {'Elvis':True, 'Candy':False}, 'EvF': {}}}, + {'id': 3, 'pool': 1, 'votes': {'GvH': {}}}, + {'id': 4, 'pool': 2, 'votes': {'AvB': {}, 'CvD': {'Candy':True}}}, + {'id': 5, 'pool': 2, 'votes': {'CvD': {'Elvis':True, 'Candy':False}, 'EvF': {}}} + ] + cvr_list = CVR.from_dict(cvr_dicts) + pool_set = set(c.pool for c in cvr_list) + print(f'{pool_set=}') + pools = {} + for p in pool_set: + pools[p] = CVR.pool_contests(list([c for c in cvr_list if c.pool == p])) + assert CVR.add_pool_contests(cvr_list, pools) + for i in range(3): + assert set(cvr_list[i].votes.keys()) == {'AvB', 'CvD', 'EvF', 'GvH'} + for i in range(3,5): + assert set(cvr_list[i].votes.keys()) == {'AvB', 'CvD', 'EvF'} + assert not CVR.add_pool_contests(cvr_list, pools) def test_cvr_from_raire(self): raire_cvrs = [['1'],