python-app-rocketbot/storage.py

import json
from os.path import exists

from discord import Guild

from config import CONFIG

class StateKey:
	WARNING_CHANNEL_ID = 'warning_channel_id'
	WARNING_MENTION = 'warning_mention'

class Storage:
	"""
	Static class for managing persisted bot state.
	"""

	# discord.Guild.id -> dict
	__guild_id_to_state = {}

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

	@classmethod
	def get_state_value(cls, guild: Guild, key: str):
		"""
		Returns a persisted state value stored under the given key. Returns
		None if not present.
		"""
		return cls.get_state(guild).get(key)

	@classmethod
	def set_state_value(cls, guild: Guild, key: str, value) -> None:
		"""
		Adds the given key-value pair to the persisted state for the given
		Guild. If `value` is `None` the key will be removed from persisted
		state.
		"""
		cls.set_state_values(guild, { key: value })

	@classmethod
	def set_state_values(cls, guild: Guild, vars: dict) -> None:
		"""
		Merges the given `vars` dict with the saved state for the given guild
		and saves it to disk. `vars` must be JSON-encodable or a ValueError will
		be raised. Keys with associated values of `None` will be removed from the
		state.
		"""
		if vars is None or len(vars) == 0:
			return
		state: dict = cls.get_state(guild)
		try:
			json.dumps(vars)
		except:
			raise ValueError(f'vars not JSON encodable - {vars}')
		for key, value in vars.items():
			if value is None:
				del state[key]
			else:
				state[key] = value
		cls.__save_guild_state(guild, state)

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

	@classmethod
	def __load_guild_state(cls, guild: Guild) -> dict:
		"""
		Loads state for a guild from a JSON file on disk, or None if not found.
		"""
		path: str = cls.__guild_path(guild)
		if not exists(path):
			cls.__trace(f'No state on disk for guild {guild.id}. Returning None.')
			return None
		cls.__trace(f'Loading state from disk for guild {guild.id}')
		with open(path, 'r') as file:
			state = json.load(file)
		cls.__trace('State loaded')
		return state

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

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