74 lines
2.3 KiB
Python
74 lines
2.3 KiB
Python
import copy
|
|
import json
|
|
import shutil
|
|
import threading
|
|
from enum import Enum
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
from colony.core.logger.logging import get_logger
|
|
from colony.utils.files import atomic_write_json
|
|
|
|
|
|
class StateManager:
|
|
class Data(Enum):
|
|
Roles = "roles"
|
|
|
|
DEFAULTS = {
|
|
Data.Roles: [],
|
|
}
|
|
|
|
STATE_FILE = Path("state.json")
|
|
|
|
def __init__(self):
|
|
self.logger = get_logger(self.__class__.__name__)
|
|
self._lock = threading.RLock()
|
|
self._data = {}
|
|
self._load()
|
|
|
|
def _load(self):
|
|
with self._lock:
|
|
if self.STATE_FILE.exists():
|
|
try:
|
|
with open(self.STATE_FILE, "r", encoding="utf-8") as f:
|
|
data = json.load(f)
|
|
if not isinstance(data, dict):
|
|
raise ValueError("top level is not an object")
|
|
for key, default in self.DEFAULTS.items():
|
|
if key.value in data and not isinstance(data[key.value], type(default)):
|
|
raise ValueError(f"'{key.value}' has the wrong type")
|
|
self._data = data
|
|
except (ValueError, OSError) as e:
|
|
self._recover_from_corrupt_file(e)
|
|
else:
|
|
self._data = {}
|
|
self._apply_defaults()
|
|
self._save()
|
|
|
|
self._apply_defaults()
|
|
|
|
def _recover_from_corrupt_file(self, error):
|
|
backup = self.STATE_FILE.with_name(self.STATE_FILE.name + ".corrupt")
|
|
self.logger.warning(f"Could not use {self.STATE_FILE} ({error}). Keeping a copy as {backup} and resetting to defaults")
|
|
try:
|
|
shutil.copy2(self.STATE_FILE, backup)
|
|
except OSError as copy_error:
|
|
self.logger.warning(f"Could not back up {self.STATE_FILE}: {copy_error}")
|
|
self._data = {}
|
|
|
|
def _apply_defaults(self):
|
|
for key, default in self.DEFAULTS.items():
|
|
self._data.setdefault(key.value, copy.deepcopy(default))
|
|
|
|
def _save(self):
|
|
atomic_write_json(self.STATE_FILE, self._data, indent=4)
|
|
|
|
def get(self, key: Data) -> Any:
|
|
with self._lock:
|
|
return copy.deepcopy(self._data[key.value])
|
|
|
|
def set(self, key: Data, value: Any):
|
|
with self._lock:
|
|
self._data[key.value] = value
|
|
self._save()
|