78 lines
2.5 KiB
Python
78 lines
2.5 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"
|
|
Peers = "peers"
|
|
|
|
DEFAULTS = {
|
|
Data.Roles: [],
|
|
Data.Peers: [],
|
|
}
|
|
|
|
STATE_FILE = Path("data") / "state.json"
|
|
|
|
def __init__(self, path=None):
|
|
self._path = Path(path) if path is not None else self.STATE_FILE
|
|
self._path.parent.mkdir(parents=True, exist_ok=True)
|
|
self.logger = get_logger(self.__class__.__name__)
|
|
self._lock = threading.RLock()
|
|
self._data = {}
|
|
self._load()
|
|
|
|
def _load(self):
|
|
with self._lock:
|
|
if self._path.exists():
|
|
try:
|
|
with open(self._path, "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._path.with_name(self._path.name + ".corrupt")
|
|
self.logger.warning(f"Could not use {self._path} ({error}). Keeping a copy as {backup} and resetting to defaults")
|
|
try:
|
|
shutil.copy2(self._path, backup)
|
|
except OSError as copy_error:
|
|
self.logger.warning(f"Could not back up {self._path}: {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._path, 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()
|