Code refactor

pyRCV2/method/base_stv.py

@@ -41,8 +41,7 @@ class BaseSTVCounter:
 		
 		if options is not None:
 			self.options.update(options)
-	
+		
-	def reset(self):
 		self.candidates = SafeDict([(c, CountCard()) for c in self.election.candidates])
 		self.exhausted = CountCard()
 		self.loss_fraction = CountCard()
@@ -57,12 +56,55 @@ class BaseSTVCounter:
 			self.candidates[candidate].state = CandidateState.WITHDRAWN
 			__pragma__('noopov')
 	
+	def reset(self):
+		"""
+		Public function:
+		Reset the count and perform the first step
+		Subclasses must override this function
+		"""
+		raise NotImplementedError('Method not implemented')
+	
 	def step(self):
+		"""
+		Public function:
+		Perform one step of the STV count
+		"""
+		
 		# Step count cards
+		self.step_count_cards()
+		
+		# Check if done
+		result = self.check_if_done()
+		if result:
+			return result
+		
+		# Distribute surpluses
+		result = self.distribute_surpluses()
+		if result:
+			return result
+		
+		# Insufficient winners and no surpluses to distribute
+		# Exclude the lowest ranked hopeful
+		result = self.exclude_candidate()
+		if result:
+			return result
+		
+		raise STVException('Unable to complete step')
+	
+	def step_count_cards(self):
+		"""
+		Reset the count cards for the beginning of a new step
+		"""
+		
 		for candidate, count_card in self.candidates.items():
 			count_card.step()
 		self.exhausted.step()
 		self.loss_fraction.step()
+	
+	def check_if_done(self):
+		"""
+		Check if the count can be completed
+		"""
 		
 		# Have sufficient candidates been elected?
 		if self.num_elected >= self.election.seats:
@@ -89,6 +131,11 @@ class BaseSTVCounter:
 				self.quota
 			)
 			__pragma__('noopov')
+	
+	def distribute_surpluses(self):
+		"""
+		Distribute surpluses, if any
+		"""
 		
 		# Do surpluses need to be distributed?
 		__pragma__('opov')
@@ -112,13 +159,13 @@ class BaseSTVCounter:
 			__pragma__('noopov')
 			
 			# Transfer surplus
-			self.transfer_surplus(candidate_surplus, count_card, surplus)
+			self.do_surplus(candidate_surplus, count_card, surplus)
 			
 			__pragma__('opov')
 			count_card.transfers -= surplus
 			__pragma__('noopov')
 			
-			# Declare any candidates meeting the quota as a result of surpluses
+			# Declare elected any candidates meeting the quota as a result of surpluses
 			self.compute_quota()
 			self.elect_meeting_quota()
 			
@@ -132,9 +179,19 @@ class BaseSTVCounter:
 				self.quota
 			)
 			__pragma__('noopov')
+	
+	def do_surplus(self, candidate_surplus, count_card, surplus):
+		"""
+		Transfer the surplus of the given candidate
+		Subclasses must override this function
+		"""
+		raise NotImplementedError('Method not implemented')
+	
+	def exclude_candidate(self):
+		"""
+		Exclude the lowest ranked hopeful
+		"""
 		
-		# Insufficient winners and no surpluses to distribute
-		# Exclude the lowest ranked hopeful
 		hopefuls = [(c, cc) for c, cc in self.candidates.items() if cc.state == CandidateState.HOPEFUL]
 		hopefuls.sort(lambda x: x[1].votes)
 		
@@ -143,7 +200,7 @@ class BaseSTVCounter:
 		count_card.state = CandidateState.EXCLUDED
 		
 		# Exclude this candidate
-		self.exclude_candidate(candidate_excluded, count_card)
+		self.do_exclusion(candidate_excluded, count_card)
 		
 		__pragma__('opov')
 		count_card.transfers -= count_card.votes
@@ -164,8 +221,18 @@ class BaseSTVCounter:
 		)
 		__pragma__('noopov')
 	
+	def do_exclusion(self, candidate_excluded, count_card):
+		"""
+		Exclude the given candidate and transfer the votes
+		Subclasses must override this function
+		"""
+		raise NotImplementedError('Method not implemented')
+	
 	def compute_quota(self):
-		"""Recount total votes and (if applicable) recalculate the quota"""
+		"""
+		Recount total votes and (if applicable) recalculate the quota
+		"""
+		
 		__pragma__('opov')
 		self.total = sum((cc.votes for c, cc in self.candidates.items()), Num('0'))
 		self.loss_fraction.transfers += (self.total_orig - self.total - self.exhausted.votes) - self.loss_fraction.votes
@@ -184,6 +251,10 @@ class BaseSTVCounter:
 		__pragma__('noopov')
 	
 	def meets_quota(self, count_card):
+		"""
+		Determine if the given candidate meets the quota
+		"""
+		
 		if self.options['quota_criterion'] == 'geq':
 			__pragma__('opov')
 			return count_card.votes >= self.quota
@@ -196,6 +267,10 @@ class BaseSTVCounter:
 			raise STVException('Invalid quota criterion')
 	
 	def elect_meeting_quota(self):
+		"""
+		Elect all candidates meeting the quota
+		"""
+		
 		# Does a candidate meet the quota?
 		meets_quota = [(c, cc) for c, cc in self.candidates.items() if cc.state == CandidateState.HOPEFUL and self.meets_quota(cc)]
 		
@@ -212,8 +287,6 @@ class BaseWIGSTVCounter(BaseSTVCounter):
 	"""
 	
 	def reset(self):
-		BaseSTVCounter.reset(self)
-		
 		# Distribute first preferences
 		for ballot in self.election.ballots:
 			__pragma__('opov')
@@ -242,7 +315,7 @@ class BaseWIGSTVCounter(BaseSTVCounter):
 		)
 		__pragma__('noopov')
 	
-	def transfer_surplus(self, candidate_surplus, count_card, surplus):
+	def do_surplus(self, candidate_surplus, count_card, surplus):
 		for ballot, ballot_value in count_card.ballots:
 			__pragma__('opov')
 			new_value = (ballot_value * surplus) / count_card.votes # Multiply first to avoid rounding errors
@@ -257,7 +330,7 @@ class BaseWIGSTVCounter(BaseSTVCounter):
 				self.exhausted.ballots.append((ballot, new_value))
 			__pragma__('noopov')
 	
-	def exclude_candidate(self, candidate_excluded, count_card):
+	def do_exclusion(self, candidate_excluded, count_card):
 		for ballot, ballot_value in count_card.ballots:
 			__pragma__('opov')
 			candidate = next((c for c in ballot.preferences if self.candidates[c].state == CandidateState.HOPEFUL), None)
@@ -276,8 +349,6 @@ class BaseUIGSTVCounter(BaseSTVCounter):
 	"""
 	
 	def reset(self):
-		BaseSTVCounter.reset(self)
-		
 		# Distribute first preferences
 		for ballot in self.election.ballots:
 			__pragma__('opov')
@@ -306,7 +377,7 @@ class BaseUIGSTVCounter(BaseSTVCounter):
 		)
 		__pragma__('noopov')
 	
-	def transfer_surplus(self, candidate_surplus, count_card, surplus):
+	def do_surplus(self, candidate_surplus, count_card, surplus):
 		# FIXME: Is it okay to use native int's here?
 		next_preferences = SafeDict([(c, []) for c, cc in self.candidates.items()])
 		next_exhausted = []
@@ -352,7 +423,7 @@ class BaseUIGSTVCounter(BaseSTVCounter):
 			__pragma__('noopov')
 			self.exhausted.ballots.append((ballot, new_value))
 	
-	def exclude_candidate(self, candidate_excluded, count_card):
+	def do_exclusion(self, candidate_excluded, count_card):
 		next_preferences = SafeDict([(c, Rational('0')) for c, cc in self.candidates.items()])
 		next_exhausted = Rational('0')
 		total_votes = Rational('0')