python-app-rocketbot/rocketbot/storage.py

"""
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

class ConfigKey:
	"""
	Common keys in persisted guild storage.
	"""
	WARNING_CHANNEL_ID = 'warning_channel_id'
	WARNING_MENTION = 'warning_mention'

class Storage:
	"""
	Static class for managing persisted bot configuration and transient state
	on a per-guild basis.
	"""

	# -- Transient state management -----------------------------------------

	__guild_id_to_state: dict[int, dict[str, Any]] = {}

	@classmethod
	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[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) -> 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: 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.
		"""
		cls.set_state_values(guild, { key: value })

	@classmethod
	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
		transient state.
		"""
		if values is None or len(values) == 0:
			return
		state: dict[str, Any] = cls.get_state(guild)
		for key, value in values.items():
			if value is None:
				del state[key]
			else:
				state[key] = value
		# XXX: Superstitious. Should update by ref already but saw weirdness once.
		cls.__guild_id_to_state[guild.id] = state

	# -- Persisted configuration management ---------------------------------

	# discord.Guild.id -> dict
	__guild_id_to_config: dict[int, dict[str, Any]] = {}

	@classmethod
	def get_config(cls, guild: Guild) -> dict[str, Any]:
		"""
		Returns all persisted configuration for the given guild.
		"""
		config: dict[str, Any] = cls.__guild_id_to_config.get(guild.id)
		if config is not None:
			# Already in memory
			return config
		# Load from disk if possible
		cls.__trace(f'No loaded config for guild {guild.id}. Attempting to ' +
			'load from disk.')
		config = cls.__read_guild_config(guild)
		if config is None:
			return {}
		cls.__guild_id_to_config[guild.id] = config
		return config

	@classmethod
	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.
		"""
		return cls.get_config(guild).get(key)

	@classmethod
	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
		config.
		"""
		cls.set_config_values(guild, { key: value })

	@classmethod
	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`
		will be raised. Keys with associated values of `None` will be removed
		from the persisted config.
		"""
		if values is None or len(values) == 0:
			return
		config: dict[str, Any] = cls.get_config(guild)
		try:
			json.dumps(values)
		except Exception as e:
			raise ValueError(f'values not JSON encodable - {values}') from e
		for key, value in values.items():
			if value is None:
				del config[key]
			else:
				config[key] = value
		cls.__write_guild_config(guild, config)

	@classmethod
	def __write_guild_config(cls, guild: Guild, config: dict[str, Any]) -> None:
		"""
		Saves config for a guild to a JSON file on disk.
		"""
		path: str = cls.__guild_config_path(guild)
		cls.__trace(f'Saving config for guild {guild.id} to {path}')
		cls.__trace(f'config = {config}')
		with open(path, 'w', encoding='utf8') as file:
			# Pretty printing to make more legible for debugging
			# Sorting keys to help with diffs
			json.dump(config, file, indent='\t', sort_keys=True)
		cls.__trace('State saved')

	@classmethod
	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.
		"""
		path: str = cls.__guild_config_path(guild)
		if not exists(path):
			cls.__trace(f'No config on disk for guild {guild.id}. Returning None.')
			return None
		cls.__trace(f'Loading config from disk for guild {guild.id}')
		with open(path, 'r', encoding='utf8') as file:
			config = json.load(file)
		cls.__trace('State loaded')
		return config

	@classmethod
	def __guild_config_path(cls, guild: Guild) -> str:
		"""
		Returns the JSON file path where guild config should be written.
		"""
		config_value: str = CONFIG['config_path']
		path: str = config_value if config_value.endswith('/') else f'{config_value}/'
		return f'{path}guild_{guild.id}.json'

	@classmethod
	def __trace(cls, message: Any) -> None:
		# print(f'{cls.__name__}: {str(message)}')
		pass