Added start vote and fixed deocrator bug

GameFiles/Game.py

@@ -1,6 +1,7 @@
 import logging
 from typing import Dict, Callable, List, Union
 from functools import wraps
+from collections import defaultdict
 
 import discord
 
@@ -11,11 +12,27 @@ def started_only(func):
 	"""
 	Decorator for *methods* of Game that need the game to be started.
 	"""
+
 	@wraps(func)
-	def decorated(obj, *args, **kwargs):
+	def decorated(obj: 'Game', *args, **kwargs):
 		if not obj.is_started():
 			raise RuntimeError("This function only works on running games!")
-		func(obj, *args, **kwargs)
+		return func(obj, *args, **kwargs)
+
+	return decorated
+
+
+def vote_running(func):
+	"""
+		Decorator for *methods* of Game that need a vote to be running
+	"""
+
+	@wraps(func)
+	def decorated(obj: 'Game', *args, **kwargs):
+		if not obj.is_vote_running():
+			raise RuntimeError("This function only works if a vote is running!")
+		return func(obj, *args, **kwargs)
+
 	return decorated
 
 
@@ -72,6 +89,25 @@ class Game:
 	def get_game_category(self) -> discord.CategoryChannel:
 		return self.guild.get_channel(self.get_game_category_id())
 
+	@started_only
+	def get_votes_channel_id(self) -> int:
+		return self.config["votes_chan"]
+
+	@started_only
+	def get_votes_channel(self) -> discord.TextChannel:
+		return self.guild.get_channel(self.get_votes_channel_id())
+
+	@started_only
+	def is_vote_running(self) -> bool:
+		return self.config["vote"] is not None
+
+	@vote_running
+	def is_vote_passing(self) -> bool:
+		vote_count = defaultdict(int)
+		for player_id in self.get_players_id():
+			vote_count[self.config["vote"][str(player_id)]] += 1
+		return vote_count[True] > vote_count[False]
+
 	async def start(self, gm_role: discord.Role, player_role: discord.Role):
 		if self.is_started():
 			raise ValueError("Game already started")
@@ -80,6 +116,7 @@ class Game:
 		self.config["player_role"] = player_role.id
 		self.config["players"] = [member.id for member in player_role.members]
 		self.config["player_info"] = {str(player): {} for player in self.config["players"]}
+		self.config["vote"] = None
 
 		permissions = {
 			self.guild.default_role: discord.PermissionOverwrite(send_messages = False),
@@ -108,14 +145,14 @@ class Game:
 		self.config["discussion_chan"] = (await game_category.create_text_channel("discussion")).id  # Permissions are inherited from the category
 		logger.debug(f"[{self.guild.name}] Created discussion channel")
 
-		for player in self.get_players():
+		for player in player_role.members:
 			channel_permissions = {
 				self.guild.default_role: discord.PermissionOverwrite(read_messages = False),
 				player: discord.PermissionOverwrite(read_messages = True),
 				gm_role: discord.PermissionOverwrite(read_messages = True),
 			}
 			player_channel = await game_category.create_text_channel(player.name, overwrites = channel_permissions)
-			self.get_player_info(player)["channel"] = player_channel.id
+			self.config["player_info"][str(player.id)]["channel"] = player_channel.id
 		logger.debug(f"[{self.guild.name}] Created player channels")
 
 		self.config["game_started"] = True
@@ -130,3 +167,52 @@ class Game:
 		self.config.clear()
 		self.config.update(self.new_dict())
 		self.save_function()
+
+	async def start_vote(self, president: discord.Member, chancellor: discord.Member):
+		if self.is_vote_running():
+			raise RuntimeError("A vote is already running")
+		logging.debug(f"[{self.guild.name}] Starting vote")
+
+		self.config["vote"] = {
+			"president": president.id,
+			"chancellor": chancellor.id,
+			"message": None,
+			"revealed": False,
+		}
+		self.config["vote"].update({str(player_id): None for player_id in self.get_players_id()})
+		self.save_function()
+		await self.update_vote_message()
+
+	@vote_running
+	async def update_vote_message(self):
+		logging.debug(f"[{self.guild.name}] Updating vote message")
+		president = self.config["vote"]["president"]
+		chancellor = self.config["vote"]["chancellor"]
+		message_content = [
+			"**Citizens are called to vote**",
+			"Do you want to elect the following government?",
+			f":crown: <@{president}> as president",
+			f":person_in_tuxedo: <@{chancellor}> as chancellor",
+			"",
+		]
+		for player_id in self.get_players_id():
+			player_vote = self.config["vote"][str(player_id)]
+			if player_vote is None:
+				message_content.append(f":black_large_square: <@{player_id}> has not voted")
+			else:
+				if self.config["vote"]["revealed"]:
+					if player_vote:
+						message_content.append(f":green_square: <@{player_id}> has voted JA")
+					else:
+						message_content.append(f":red_square: <@{player_id}> has voted NEIN")
+				else:  # Player has voted but the vote should not be revealed
+					message_content.append(f":white_large_square:  <@{player_id}> has voted")
+
+		message_content_str = "\n".join(message_content)
+		if self.config["vote"]["message"] is None:
+			self.config["vote"]["message"] = (await self.get_votes_channel().send(message_content_str, allowed_mentions = discord.AllowedMentions.none())).id
+		else:
+			await (await self.get_votes_channel().fetch_message(self.config["vote"]["message"])).edit(content = message_content_str, allowed_mentions = discord.AllowedMentions.none())
+
+
+

SecretBot.py

@@ -114,6 +114,18 @@ class SecretBot(commands.Cog):
 		else:
 			await ctx.reply(":x: Game is not running")
 
+	@commands.command("StartVote")
+	async def start_vote(self, ctx:commands.Context, president: discord.Member, chancellor: discord.Member):
+		game = self.games_file[ctx.guild]
+		if not game.is_started():
+			await ctx.reply(":x: Game is not running")
+			return
+		if game.is_vote_running():
+			await ctx.reply(":x: A vote is already running")
+			return
+		await game.start_vote(president, chancellor)
+		await ctx.message.delete()
+
 
 if __name__ == '__main__':
 	argparser = argparse.ArgumentParser(description = "Secret Hitler helper bot", formatter_class = argparse.ArgumentDefaultsHelpFormatter)