161 lines
5.1 KiB
Python
161 lines
5.1 KiB
Python
from core.models.manage.validate_dataclass import ValidatedDataclass
|
|
from core.utils.basic_operations.file_tools import delete_file
|
|
from core.errors.logger import logger
|
|
|
|
from dataclasses import fields, MISSING
|
|
import json
|
|
import os
|
|
from functools import wraps
|
|
from enum import Enum
|
|
from typing import Any
|
|
|
|
|
|
def singleton(cls):
|
|
instances = {}
|
|
|
|
@wraps(cls)
|
|
def get_instance(*args, **kwargs):
|
|
instance_id = kwargs.get('instance_id')
|
|
|
|
# If instance_id provided, use localized mode (one instance per ID)
|
|
# Otherwise, use global mode (single instance)
|
|
key = (cls, instance_id) if instance_id else cls
|
|
|
|
if key not in instances:
|
|
print(f"Creating new instance with key: {key}")
|
|
instances[key] = cls(*args, **kwargs)
|
|
else:
|
|
print(f"Reusing existing instance with key: {key}")
|
|
|
|
# either way:
|
|
return instances[key]
|
|
|
|
return get_instance
|
|
|
|
|
|
# for saving.
|
|
EXCLUDED_FIELDS = {"parent_folder", "folder_path", "_data_file"}
|
|
|
|
|
|
class StatePersistenceMixin(ValidatedDataclass):
|
|
def _serialize_value(self, value):
|
|
"""Helper method to ensure values are JSON-serializable"""
|
|
if isinstance(value, Enum):
|
|
return value.value
|
|
return value
|
|
|
|
|
|
def prep_payload(self) -> dict|None:
|
|
"""Orchestrate Save, Validation, and Payload preparation of relevant fields"""
|
|
|
|
if not self.save():
|
|
logger.error("Failed to Save! Not returning a payload")
|
|
return None
|
|
|
|
payload = {}
|
|
|
|
for f in fields(self):
|
|
# Only add the field if it's in our allowed list
|
|
if f.name in self.PAYLOAD_FIELDS:
|
|
# Get the raw value
|
|
raw_value = getattr(self, f.name)
|
|
# Clean it (handle Enums)
|
|
clean_value = self._serialize_value(raw_value)
|
|
# Add it to the dictionary
|
|
payload[f.name] = clean_value
|
|
|
|
return payload
|
|
|
|
|
|
def save(self) -> bool:
|
|
"""Save state to file"""
|
|
|
|
if not self.validate_payload():
|
|
return False
|
|
|
|
os.makedirs(self.folder_path, exist_ok=True)
|
|
|
|
# 1. Start with an empty dictionary
|
|
data = {}
|
|
|
|
# 2. Loop through the fields
|
|
for f in fields(self):
|
|
# Only add the field if it's NOT in the excluded list
|
|
if f.name not in EXCLUDED_FIELDS:
|
|
# Get the raw value
|
|
raw_value = getattr(self, f.name)
|
|
# Clean it (handle Enums)
|
|
clean_value = self._serialize_value(raw_value)
|
|
# Add it to the dictionary
|
|
data[f.name] = clean_value
|
|
|
|
print(f"saving to {self._data_file}")
|
|
with open(self._data_file, "w") as f:
|
|
json.dump(data, f, indent=4, sort_keys=True)
|
|
|
|
return True
|
|
|
|
|
|
def load(self):
|
|
"""Load state from file, if each item passes a type check. With Enum translation"""
|
|
|
|
logger.info("Starting load of pre-existing data")
|
|
|
|
if not os.path.exists(self._data_file):
|
|
logger.info("does NOT exist already")
|
|
return
|
|
|
|
with open(self._data_file, "r") as f:
|
|
data = json.load(f)
|
|
|
|
for field in fields(self):
|
|
value = data.get(field.name)
|
|
logger.info(f"testing {field.name} of {value}")
|
|
if value is None:
|
|
continue
|
|
|
|
# Translation: If the expected type is an Enum, convert the raw value back into an Enum member
|
|
if isinstance(field.type, type) and issubclass(field.type, Enum):
|
|
logger.info("testing enum")
|
|
try:
|
|
value = field.type(value)
|
|
except ValueError:
|
|
logger.error(f"Invalid enum value {value} for {field.name}")
|
|
continue
|
|
|
|
# type check works because 'value' is now an actual Enum member
|
|
if isinstance(value, field.type):
|
|
setattr(self, field.name, value)
|
|
|
|
|
|
def wipe_data_on_disk(self):
|
|
logger.info("Starting wipe of pre-existing data")
|
|
|
|
if not os.path.exists(self._data_file):
|
|
logger.info("File does NOT exist already")
|
|
return
|
|
|
|
return delete_file(self._data_file)
|
|
|
|
|
|
def exists_already(self) -> bool:
|
|
return os.path.exists(self._data_file)
|
|
|
|
|
|
|
|
def wipe_model(model: Any, KEEP_FIELDS: list) -> Any:
|
|
for field in fields(model):
|
|
if field.name not in KEEP_FIELDS:
|
|
# 1. Check for a static default value (e.g., attempts: int = 0)
|
|
if field.default is not MISSING:
|
|
setattr(model, field.name, field.default)
|
|
|
|
# 2. Check for a default factory (e.g., items: list = field(default_factory=list))
|
|
elif field.default_factory is not MISSING:
|
|
setattr(model, field.name, field.default_factory())
|
|
|
|
# 3. Fallback: if no default is defined at all, set it to None.
|
|
else:
|
|
setattr(model, field.name, None)
|
|
|
|
return model
|