""" Core Hook System for Atomizer Defines hook points in the optimization lifecycle and hook registration mechanism. """ from enum import Enum from typing import Callable, Dict, Any, Optional from dataclasses import dataclass import logging logger = logging.getLogger(__name__) class HookPoint(Enum): """Enumeration of available hook points in the optimization lifecycle.""" PRE_MESH = "pre_mesh" # Before meshing POST_MESH = "post_mesh" # After meshing, before solve PRE_SOLVE = "pre_solve" # Before solver execution POST_SOLVE = "post_solve" # After solve, before extraction POST_EXTRACTION = "post_extraction" # After result extraction CUSTOM_OBJECTIVE = "custom_objective" # Custom objective functions @dataclass class Hook: """ Represents a single hook function to be executed at a specific point. Attributes: name: Unique identifier for this hook hook_point: When this hook should execute (HookPoint enum) function: The callable to execute description: Human-readable description of what this hook does priority: Execution order (lower = earlier, default=100) enabled: Whether this hook is currently active """ name: str hook_point: HookPoint function: Callable[[Dict[str, Any]], Optional[Dict[str, Any]]] description: str priority: int = 100 enabled: bool = True def execute(self, context: Dict[str, Any]) -> Optional[Dict[str, Any]]: """ Execute this hook with the given context. Args: context: Dictionary containing relevant data for this hook point Common keys: - trial_number: Current trial number - design_variables: Current design variable values - sim_file: Path to simulation file - working_dir: Current working directory Returns: Optional dictionary with results or modifications to context Raises: Exception: Any exception from the hook function is logged and re-raised """ if not self.enabled: logger.debug(f"Hook '{self.name}' is disabled, skipping") return None try: logger.info(f"Executing hook '{self.name}' at {self.hook_point.value}") result = self.function(context) logger.debug(f"Hook '{self.name}' completed successfully") return result except Exception as e: logger.error(f"Hook '{self.name}' failed: {e}", exc_info=True) raise def __repr__(self) -> str: status = "enabled" if self.enabled else "disabled" return f"Hook(name='{self.name}', point={self.hook_point.value}, priority={self.priority}, {status})" class HookContext: """ Context object passed to hooks containing all relevant data. This is a convenience wrapper around a dictionary that provides both dict-like access and attribute access. """ def __init__(self, **kwargs): self._data = kwargs def __getitem__(self, key: str) -> Any: return self._data[key] def __setitem__(self, key: str, value: Any): self._data[key] = value def __contains__(self, key: str) -> bool: return key in self._data def get(self, key: str, default: Any = None) -> Any: return self._data.get(key, default) def update(self, other: Dict[str, Any]): self._data.update(other) def to_dict(self) -> Dict[str, Any]: return self._data.copy() def __repr__(self) -> str: keys = list(self._data.keys()) return f"HookContext({', '.join(keys)})"