WIP: added CVR.add_pool_contests() method

shangrla/shangrla/Audit.py

@@ -165,18 +165,18 @@ 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)}'
 
     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] :

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'],