Edit
IABSD.fr/xenocara/lib/mesa/bin/ci/structured_logger.py
jsg
- lib/mesa/bin/ci/structured_logger.py
-
""" A structured logging utility supporting multiple data formats such as CSV, JSON, and YAML. The main purpose of this script, besides having relevant information available in a condensed and deserialized. This script defines a protocol for different file handling strategies and provides implementations for CSV, JSON, and YAML formats. The main class, StructuredLogger, allows for easy interaction with log data, enabling users to load, save, increment, set, and append fields in the log. The script also includes context managers for file locking and editing log data to ensure data integrity and avoid race conditions. """ import json import os from collections.abc import MutableMapping, MutableSequence from contextlib import contextmanager from datetime import datetime from pathlib import Path from typing import Any, Protocol import fire from filelock import FileLock try: import polars as pl CSV_LIB_EXCEPTION = None except ImportError as e: CSV_LIB_EXCEPTION: ImportError = e try: from ruamel.yaml import YAML YAML_LIB_EXCEPTION = None except ImportError as e: YAML_LIB_EXCEPTION: ImportError = e class ContainerProxy: """ A proxy class that wraps a mutable container object (such as a dictionary or a list) and calls a provided save_callback function whenever the container or its contents are changed. """ def __init__(self, container, save_callback): self.container = container self.save_callback = save_callback def __getitem__(self, key): value = self.container[key] if isinstance(value, (MutableMapping, MutableSequence)): return ContainerProxy(value, self.save_callback) return value def __setitem__(self, key, value): self.container[key] = value self.save_callback() def __delitem__(self, key): del self.container[key] self.save_callback() def __getattr__(self, name): attr = getattr(self.container, name) if callable(attr): def wrapper(*args, **kwargs): result = attr(*args, **kwargs) self.save_callback() return result return wrapper return attr def __iter__(self): return iter(self.container) def __len__(self): return len(self.container) def __repr__(self): return repr(self.container) class AutoSaveDict(dict): """ A subclass of the built-in dict class with additional functionality to automatically save changes to the dictionary. It maintains a timestamp of the last modification and automatically wraps nested mutable containers using ContainerProxy. """ timestamp_key = "_timestamp" def __init__(self, *args, save_callback, register_timestamp=True, **kwargs): self.save_callback = save_callback self.__register_timestamp = register_timestamp self.__heartbeat() super().__init__(*args, **kwargs) self.__wrap_dictionaries() def __heartbeat(self): if self.__register_timestamp: self[AutoSaveDict.timestamp_key] = datetime.now().isoformat() def __save(self): self.__heartbeat() self.save_callback() def __wrap_dictionaries(self): for key, value in self.items(): if isinstance(value, MutableMapping) and not isinstance( value, AutoSaveDict ): self[key] = AutoSaveDict( value, save_callback=self.save_callback, register_timestamp=False ) def __setitem__(self, key, value): if isinstance(value, MutableMapping) and not isinstance(value, AutoSaveDict): value = AutoSaveDict( value, save_callback=self.save_callback, register_timestamp=False ) super().__setitem__(key, value) if self.__register_timestamp and key == AutoSaveDict.timestamp_key: return self.__save() def __getitem__(self, key): value = super().__getitem__(key) if isinstance(value, (MutableMapping, MutableSequence)): return ContainerProxy(value, self.__save) return value def __delitem__(self, key): super().__delitem__(key) self.__save() def pop(self, *args, **kwargs): result = super().pop(*args, **kwargs) self.__save() return result def update(self, *args, **kwargs): super().update(*args, **kwargs) self.__wrap_dictionaries() self.__save() class StructuredLoggerStrategy(Protocol): def load_data(self, file_path: Path) -> dict: pass def save_data(self, file_path: Path, data: dict) -> None: pass class CSVStrategy: def __init__(self) -> None: if CSV_LIB_EXCEPTION: raise RuntimeError( "Can't parse CSV files. Missing library" ) from CSV_LIB_EXCEPTION def load_data(self, file_path: Path) -> dict: dicts: list[dict[str, Any]] = pl.read_csv( file_path, try_parse_dates=True ).to_dicts() data = {} for d in dicts: for k, v in d.items(): if k != AutoSaveDict.timestamp_key and k in data: if isinstance(data[k], list): data[k].append(v) continue data[k] = [data[k], v] else: data[k] = v return data def save_data(self, file_path: Path, data: dict) -> None: pl.DataFrame(data).write_csv(file_path) class JSONStrategy: def load_data(self, file_path: Path) -> dict: return json.loads(file_path.read_text()) def save_data(self, file_path: Path, data: dict) -> None: with open(file_path, "w") as f: json.dump(data, f, indent=2) class YAMLStrategy: def __init__(self): if YAML_LIB_EXCEPTION: raise RuntimeError( "Can't parse YAML files. Missing library" ) from YAML_LIB_EXCEPTION self.yaml = YAML() self.yaml.indent(sequence=4, offset=2) self.yaml.default_flow_style = False self.yaml.representer.add_representer(AutoSaveDict, self.represent_dict) @classmethod def represent_dict(cls, dumper, data): return dumper.represent_mapping("tag:yaml.org,2002:map", data) def load_data(self, file_path: Path) -> dict: return self.yaml.load(file_path.read_text()) def save_data(self, file_path: Path, data: dict) -> None: with open(file_path, "w") as f: self.yaml.dump(data, f) class StructuredLogger: def __init__( self, file_name: str, strategy: StructuredLoggerStrategy = None, truncate=False ): self.file_name: str = file_name self.file_path = Path(self.file_name) self._data: AutoSaveDict = AutoSaveDict(save_callback=self.save_data) if strategy is None: self.strategy: StructuredLoggerStrategy = self.guess_strategy_from_file( self.file_path ) else: self.strategy = strategy if not self.file_path.exists(): Path.mkdir(self.file_path.parent, exist_ok=True) self.save_data() return if truncate: with self.get_lock(): os.truncate(self.file_path, 0) self.save_data() def load_data(self): self._data = self.strategy.load_data(self.file_path) def save_data(self): self.strategy.save_data(self.file_path, self._data) @property def data(self) -> AutoSaveDict: return self._data @contextmanager def get_lock(self): with FileLock(f"{self.file_path}.lock", timeout=10): yield @contextmanager def edit_context(self): """ Context manager that ensures proper loading and saving of log data when performing multiple modifications. """ with self.get_lock(): try: self.load_data() yield finally: self.save_data() @staticmethod def guess_strategy_from_file(file_path: Path) -> StructuredLoggerStrategy: file_extension = file_path.suffix.lower().lstrip(".") return StructuredLogger.get_strategy(file_extension) @staticmethod def get_strategy(strategy_name: str) -> StructuredLoggerStrategy: strategies = { "csv": CSVStrategy, "json": JSONStrategy, "yaml": YAMLStrategy, "yml": YAMLStrategy, } try: return strategies[strategy_name]() except KeyError as e: raise ValueError(f"Unknown strategy for: {strategy_name}") from e if __name__ == "__main__": fire.Fire(StructuredLogger)