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'

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

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

	@classmethod
	def state_for_guild(cls, guild: Guild) -> dict:
		"""
		Returns the state for the given guild, loading from disk if necessary.
		Always returns a dict.
		"""
		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)
		cls.__guild_id_to_state[guild.id] = state
		return state

	@classmethod
	def save_guild_state(cls, guild: Guild) -> None:
		"""
		Saves state for the given guild to disk, if any exists.
		"""
		state: dict = cls.__guild_id_to_state.get(guild.id)
		if state is None:
			# Nothing to save
			return
		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('Saving state for guild {guild.id} to {path}')
		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.
		"""
		path: str = cls.__guild_path(guild)
		if not exists(path):
			cls.__trace(f'No state on disk for guild {guild.id}. Returning {{}}.')
			return {}
		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}')