Annotating storage.py

rocketbot/storage.py

@@ -3,6 +3,7 @@ Handles storage of persisted and non-persisted data for the bot.
 """
 import json
 from os.path import exists
+from typing import Any, Optional
 from discord import Guild
 
 from config import CONFIG
@@ -22,29 +23,29 @@ class Storage:
 
 	# -- Transient state management -----------------------------------------
 
-	__guild_id_to_state = {}
+	__guild_id_to_state: dict[int, dict[str, Any]] = {}
 
 	@classmethod
-	def get_state(cls, guild: Guild) -> dict:
+	def get_state(cls, guild: Guild) -> dict[str, Any]:
 		"""
 		Returns transient state for the given guild. This state is not preserved
 		if the bot is restarted.
 		"""
-		state: dict = cls.__guild_id_to_state.get(guild.id)
+		state: dict[str, Any] = cls.__guild_id_to_state.get(guild.id)
 		if state is None:
 			state = {}
 			cls.__guild_id_to_state[guild.id] = state
 		return state
 
 	@classmethod
-	def get_state_value(cls, guild: Guild, key: str):
+	def get_state_value(cls, guild: Guild, key: str) -> Optional[Any]:
 		"""
 		Returns a state value for the given guild and key, or `None` if not set.
 		"""
 		return cls.get_state(guild).get(key)
 
 	@classmethod
-	def set_state_value(cls, guild: Guild, key: str, value) -> None:
+	def set_state_value(cls, guild: Guild, key: str, value: Optional[Any]) -> None:
 		"""
 		Updates a transient value associated with the given guild and key name.
 		A value of `None` removes any previous value for that key.
@@ -52,7 +53,7 @@ class Storage:
 		cls.set_state_values(guild, { key: value })
 
 	@classmethod
-	def set_state_values(cls, guild: Guild, values: dict) -> None:
+	def set_state_values(cls, guild: Guild, values: Optional[dict[str, Optional[Any]]]) -> None:
 		"""
 		Merges in a set of key-value pairs into the transient state for the
 		given guild. Any pairs with a value of `None` will be removed from the
@@ -60,7 +61,7 @@ class Storage:
 		"""
 		if values is None or len(values) == 0:
 			return
-		state: dict = cls.get_state(guild)
+		state: dict[str, Any] = cls.get_state(guild)
 		for key, value in values.items():
 			if value is None:
 				del state[key]
@@ -72,14 +73,14 @@ class Storage:
 	# -- Persisted configuration management ---------------------------------
 
 	# discord.Guild.id -> dict
-	__guild_id_to_config = {}
+	__guild_id_to_config: dict[int, dict[str, Any]] = {}
 
 	@classmethod
-	def get_config(cls, guild: Guild) -> dict:
+	def get_config(cls, guild: Guild) -> dict[str, Any]:
 		"""
 		Returns all persisted configuration for the given guild.
 		"""
-		config: dict = cls.__guild_id_to_config.get(guild.id)
+		config: dict[str, Any] = cls.__guild_id_to_config.get(guild.id)
 		if config is not None:
 			# Already in memory
 			return config
@@ -93,7 +94,7 @@ class Storage:
 		return config
 
 	@classmethod
-	def get_config_value(cls, guild: Guild, key: str):
+	def get_config_value(cls, guild: Guild, key: str) -> Optional[Any]:
 		"""
 		Returns a persisted guild config value stored under the given key.
 		Returns `None` if not present.
@@ -101,7 +102,7 @@ class Storage:
 		return cls.get_config(guild).get(key)
 
 	@classmethod
-	def set_config_value(cls, guild: Guild, key: str, value) -> None:
+	def set_config_value(cls, guild: Guild, key: str, value: Optional[Any]) -> None:
 		"""
 		Adds/updates the given key-value pair to the persisted config for the
 		given Guild. If `value` is `None` the key will be removed from persisted
@@ -110,7 +111,7 @@ class Storage:
 		cls.set_config_values(guild, { key: value })
 
 	@classmethod
-	def set_config_values(cls, guild: Guild, values: dict) -> None:
+	def set_config_values(cls, guild: Guild, values: Optional[dict[str, Optional[Any]]]) -> None:
 		"""
 		Merges the given `values` dict with the saved config for the given guild
 		and writes it to disk. `values` must be JSON-encodable or a `ValueError`
@@ -119,7 +120,7 @@ class Storage:
 		"""
 		if values is None or len(values) == 0:
 			return
-		config: dict = cls.get_config(guild)
+		config: dict[str, Any] = cls.get_config(guild)
 		try:
 			json.dumps(values)
 		except Exception as e:
@@ -132,7 +133,7 @@ class Storage:
 		cls.__write_guild_config(guild, config)
 
 	@classmethod
-	def __write_guild_config(cls, guild: Guild, config: dict) -> None:
+	def __write_guild_config(cls, guild: Guild, config: dict[str, Any]) -> None:
 		"""
 		Saves config for a guild to a JSON file on disk.
 		"""
@@ -146,7 +147,7 @@ class Storage:
 		cls.__trace('State saved')
 
 	@classmethod
-	def __read_guild_config(cls, guild: Guild) -> dict:
+	def __read_guild_config(cls, guild: Guild) -> Optional[dict[str, Any]]:
 		"""
 		Loads config for a guild from a JSON file on disk, or `None` if not
 		found.
@@ -171,6 +172,6 @@ class Storage:
 		return f'{path}guild_{guild.id}.json'
 
 	@classmethod
-	def __trace(cls, message: str) -> None:
-		# print(f'{cls.__name__}: {message}')
+	def __trace(cls, message: Any) -> None:
+		# print(f'{cls.__name__}: {str(message)}')
 		pass