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}')