python-app-rocketbot/rocketbot/storage.py

"""
Handles storage of persisted and non-persisted data for the bot.
"""
import json
from datetime import datetime, timezone, timedelta
from os.path import exists
from typing import Any, Optional
from discord import Guild

from config import CONFIG
from rocketbot.collections import AgeBoundDict


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:
			config = {}
		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 get_bot_messages(cls, guild: Guild) -> AgeBoundDict[int, Any, datetime, timedelta]:
		"""Returns all the bot messages for a guild."""
		bm = cls.get_state_value(guild, 'bot_messages')
		if bm is None:
			far_future = datetime.now(timezone.utc) + timedelta(days=1000)
			bm = AgeBoundDict(timedelta(seconds=600),
				lambda k, v : v.message_sent_at() or far_future)
			Storage.set_state_value(guild, 'bot_messages', bm)
		return bm

	@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}')
		config['_guild_name'] = guild.name  # Just for making JSON files easier to identify
		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