diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 9e97e00c..0f3d1a06 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -31,3 +31,13 @@ repos: - id: ruff-check args: [--fix] - id: ruff-format + + - repo: https://github.com/pre-commit/mirrors-mypy + rev: v2.3.1 + hooks: + - id: mypy + args: [--strict, --ignore-missing-imports] + additional_dependencies: + - pydantic + - pandas-stubs + - types-requests diff --git a/pyproject.toml b/pyproject.toml index 92236b73..ca9c9e79 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -37,7 +37,16 @@ dynamic = ["version"] version_file = "src/badger/_version.py" [project.optional-dependencies] -dev = ["pytest", "pytest-cov", "pytest-qt", "pytest-mock", "pyqt5-stubs"] +dev = [ + "pytest", + "pytest-cov", + "pytest-qt", + "pytest-mock", + "pandas-stubs", + "types-requests", + "types-PyYAML", + "mypy", +] [project.urls] Homepage = "https://github.com/xopt-org/Badger" @@ -57,7 +66,7 @@ include = ["badger"] namespaces = false [tool.pytest.ini_options] -addopts = "--cov=badger" +addopts = "--cov=badger/" log_cli_level = "INFO" log_level = "DEBUG" testpaths = ["src/badger/tests"] diff --git a/src/badger/__main__.py b/src/badger/__main__.py index 70e47f6e..1966fa77 100644 --- a/src/badger/__main__.py +++ b/src/badger/__main__.py @@ -18,7 +18,7 @@ logger = logging.getLogger("badger") -def main(): +def main() -> None: # Create the top-level parser parser = argparse.ArgumentParser(description="Badger the optimizer") parser.add_argument("-g", "--gui", action="store_true", help="launch the GUI") diff --git a/src/badger/actions/__init__.py b/src/badger/actions/__init__.py index c38fc53e..c8e4efe9 100644 --- a/src/badger/actions/__init__.py +++ b/src/badger/actions/__init__.py @@ -2,8 +2,8 @@ GUI or prints config details; sub-command handlers are re-exported from their own modules (doctor, routine, env, generator, etc.).""" -import argparse import os +from argparse import Namespace from importlib import metadata from badger.actions.doctor import check_n_config_paths @@ -11,7 +11,7 @@ from badger.utils import yprint -def show_info(args: argparse.Namespace) -> None: +def show_info(args: Namespace) -> None: config_path = None if args.config_filepath: @@ -25,9 +25,9 @@ def show_info(args: argparse.Namespace) -> None: return if args.mini and check_n_config_paths(args.config_filepath): - from badger.gui.mini import launch_gui + from badger.gui.mini import launch_mini_gui - launch_gui(config_path, template_filename=args.template) + launch_mini_gui(config_path, args.template) return diff --git a/src/badger/actions/config.py b/src/badger/actions/config.py index b448b87a..0a1c2b1e 100644 --- a/src/badger/actions/config.py +++ b/src/badger/actions/config.py @@ -4,6 +4,7 @@ import logging import os +from argparse import Namespace from badger.settings import init_settings from badger.utils import convert_str_to_value, yprint @@ -11,13 +12,13 @@ logger = logging.getLogger(__name__) -def config_settings(args): +def config_settings(args: Namespace) -> bool | None: config_singleton = init_settings() key = args.key if key is None: yprint(config_singleton.list_path_settings()) - return + return None try: try: @@ -29,12 +30,13 @@ def config_settings(args): except IndexError: pass except KeyboardInterrupt: - return + return None logger.error(f"{key} is not a valid Badger config key!") + return None -def _config_path_var(var_name): +def _config_path_var(var_name: str) -> bool: config_singleton = init_settings() is_path = config_singleton.read_is_path(var_name) @@ -100,7 +102,7 @@ def _config_path_var(var_name): return False -def _config_core_var(var_name): +def _config_core_var(var_name: str) -> bool: config_singleton = init_settings() display_name = config_singleton.read_display_name(var_name) @@ -132,6 +134,11 @@ def _config_core_var(var_name): if res == "R": config_singleton.write_value(var_name, default) print(f"You reset the {display_name} setting") + return False # fault value elif res != "S": config_singleton.write_value(var_name, convert_str_to_value(res)) print(f"You set {display_name} to {res}") + return True # success value + else: + print(f"You skipped the {display_name} setting") + return False diff --git a/src/badger/actions/doctor.py b/src/badger/actions/doctor.py index b2b77d06..c3ebd941 100644 --- a/src/badger/actions/doctor.py +++ b/src/badger/actions/doctor.py @@ -2,11 +2,13 @@ offers to fix missing ones interactively, and can factory-reset Badger back to its default state.""" +from argparse import Namespace + from badger.actions.config import _config_path_var from badger.settings import init_settings, mock_settings -def self_check(args): +def self_check(args: Namespace) -> None: config = init_settings() # Reset Badger @@ -35,7 +37,7 @@ def self_check(args): print("Badger is healthy!") -def check_n_config_paths(config_filepath=None): +def check_n_config_paths(config_filepath: str | None = None) -> bool: if config_filepath is not None: config = init_settings(config_filepath) else: @@ -43,9 +45,14 @@ def check_n_config_paths(config_filepath=None): good = True all_bad = True # if all config paths are empty, we'll suggest initialization - issue_list = [] - - for pname, pvalue in config._config.model_dump(by_alias=True).items(): + issue_list: list[str] = [] + + for ( + pname, + pvalue, + ) in config._config.model_dump( # pyright: ignore[reportPrivateUsage] + by_alias=True + ).items(): if config.read_value(pname) is None: good = False issue_list.append(pname) diff --git a/src/badger/actions/env.py b/src/badger/actions/env.py index 4cfcfe6a..994f063f 100644 --- a/src/badger/actions/env.py +++ b/src/badger/actions/env.py @@ -1,8 +1,8 @@ """The `badger env` command. Lists available environment plugins or shows the details (variables, observations, parameters) of a specific one.""" -import argparse import logging +from argparse import Namespace from badger.errors import BadgerInvalidPluginError, BadgerPluginNotFoundError from badger.utils import range_to_str, yprint @@ -10,7 +10,7 @@ logger = logging.getLogger(__name__) -def show_env(args: argparse.Namespace) -> None: +def show_env(args: Namespace) -> None: try: from badger.factory import get_env, list_env except Exception: diff --git a/src/badger/actions/generator.py b/src/badger/actions/generator.py index 7f7a0059..4c0cfc63 100644 --- a/src/badger/actions/generator.py +++ b/src/badger/actions/generator.py @@ -2,6 +2,7 @@ or shows the default configuration of a specific one.""" import logging +from argparse import Namespace from xopt.errors import XoptError @@ -10,7 +11,7 @@ logger = logging.getLogger(__name__) -def show_generator(args): +def show_generator(args: Namespace) -> None: try: from badger.factory import get_generator, list_generators except Exception as e: # noqa: BLE001 - import triggers config/plugin loading; report and exit diff --git a/src/badger/actions/install.py b/src/badger/actions/install.py index 3bd25545..2aa66d1e 100644 --- a/src/badger/actions/install.py +++ b/src/badger/actions/install.py @@ -3,6 +3,7 @@ the current plugin installation workflow.""" import logging +from argparse import Namespace from typing_extensions import deprecated @@ -10,7 +11,7 @@ @deprecated("The `badger install` command is currently disabled.") -def plugin_install(args): +def plugin_install(args: Namespace) -> None: print( "This command is currently disabled.\n" "Please refer to the Badger documentation for plugin management.\n\n" diff --git a/src/badger/actions/intf.py b/src/badger/actions/intf.py index 31d3b5f0..9687abd6 100644 --- a/src/badger/actions/intf.py +++ b/src/badger/actions/intf.py @@ -2,6 +2,7 @@ shows the details of a specific one.""" import logging +from typing import Any from badger.errors import BadgerInvalidPluginError, BadgerPluginNotFoundError from badger.utils import yprint @@ -9,7 +10,7 @@ logger = logging.getLogger(__name__) -def show_intf(args): +def show_intf(args: Any) -> None: try: from badger.factory import get_intf, list_intf except Exception as e: # noqa: BLE001 - import triggers config/plugin loading; report and exit diff --git a/src/badger/actions/routine.py b/src/badger/actions/routine.py index 7d8ad9e5..32ecf56e 100644 --- a/src/badger/actions/routine.py +++ b/src/badger/actions/routine.py @@ -3,6 +3,7 @@ backwards compatibility.""" import logging +from argparse import Namespace from typing_extensions import deprecated @@ -10,7 +11,7 @@ @deprecated("The `badger routine` command is deprecated. Please use the GUI.") -def show_routine(args): +def show_routine(args: Namespace) -> None: print( "This command is deprecated.\n" "Please use 'badger -g' to launch the Badger GUI " diff --git a/src/badger/actions/run.py b/src/badger/actions/run.py index b205aff5..0c66fc97 100644 --- a/src/badger/actions/run.py +++ b/src/badger/actions/run.py @@ -13,6 +13,8 @@ import signal import sys import time +from argparse import Namespace +from typing import Any, TypedDict from pandas import DataFrame from typing_extensions import deprecated @@ -26,9 +28,20 @@ logger = logging.getLogger(__name__) +class Storage(TypedDict): + states: Any + ts_last_dump: float | None + paused: bool + + def run_n_archive( - routine: Routine, yes=False, save=False, verbose=2, sleep=0, flush_prompt=False -): + routine: Routine, + yes: bool = False, + save: bool = False, + verbose: int = 2, + sleep: float = 0, + flush_prompt: bool = False, +) -> None: try: from badger.archive import archive_run except Exception as e: # noqa: BLE001 - import triggers config/plugin loading; report and exit @@ -36,13 +49,13 @@ def run_n_archive( return # Store system states and other stuff - storage = { + storage: Storage = { "states": None, "ts_last_dump": None, "paused": False, } - def handler(*args): + def handler(*args: Any) -> None: if storage["paused"]: print() # start a new line if flush_prompt: # erase the last prompt @@ -52,10 +65,10 @@ def handler(*args): signal.signal(signal.SIGINT, handler) - def check_run_status(): + def check_run_status() -> int: return 0 - def before_evaluate(candidates: DataFrame): + def before_evaluate(candidates: DataFrame) -> None: if storage["paused"]: res = input( "Optimization paused. Press Enter to resume or Ctrl/Cmd + C to terminate: " @@ -70,7 +83,7 @@ def before_evaluate(candidates: DataFrame): sys.stdout.write("\033[F") storage["paused"] = False - def after_evaluate(data: DataFrame): + def after_evaluate(data: DataFrame) -> None: # vars: ndarray # obses: ndarray # cons: ndarray @@ -88,16 +101,16 @@ def after_evaluate(data: DataFrame): try: path = _run["path"] filename = _run["filename"][:-4] + "pickle" - routine.environment.interface.dump_recording( - os.path.join(path, filename) - ) + interface = getattr(routine.environment, "interface", None) + if interface is not None: + interface.dump_recording(os.path.join(path, filename)) except Exception: # noqa: BLE001 - interface dump is best-effort logger.warning("Failed to dump interface logs") # take a break to let the outside signal to change the status time.sleep(sleep) - def states_ready(states): + def states_ready(states: Any) -> None: storage["states"] = states try: @@ -114,7 +127,7 @@ def states_ready(states): logger.error(e) # Save the run when at least one solution has been evaluated - if len(routine.data): + if routine.data and len(routine.data): _run = archive_run(routine, storage["states"]) # Try dump the interface logs try: @@ -126,7 +139,7 @@ def states_ready(states): @deprecated("The `badger run` command is deprecated. Please use the GUI.") -def run_routine(args): +def run_routine(args: Namespace) -> None: print( "This command is deprecated.\n" "Please use 'badger -g' to launch the Badger GUI " diff --git a/src/badger/actions/uninstall.py b/src/badger/actions/uninstall.py index a1a0b470..00ce8b6b 100644 --- a/src/badger/actions/uninstall.py +++ b/src/badger/actions/uninstall.py @@ -2,6 +2,7 @@ removing plugins — see docs for the current workflow.""" import logging +from argparse import Namespace from typing_extensions import deprecated @@ -9,7 +10,7 @@ @deprecated("The `badger uninstall` command is currently disabled.") -def plugin_remove(args): +def plugin_remove(args: Namespace) -> None: print( "This command is currently disabled.\n" "Please refer to the Badger documentation for plugin management.\n\n" diff --git a/src/badger/archive.py b/src/badger/archive.py index 7f85e842..afed6741 100644 --- a/src/badger/archive.py +++ b/src/badger/archive.py @@ -11,6 +11,8 @@ import os import time import warnings +from collections.abc import Hashable +from typing import Any, TypedDict from badger.errors import BadgerConfigError from badger.routine import Routine @@ -30,12 +32,23 @@ logger.info(f"Badger run root {BADGER_ARCHIVE_ROOT} created") -def archive_run(routine, states=None): - # routine: Routine +class RunArchive(TypedDict, total=False): + filename: str + routine: Routine + data: dict[Hashable, Any] # TODO: make this more specific + system_states: Any | None # TODO: make this more specific + path: str + + +def archive_run( + routine: Routine, states: Any | None = None +) -> RunArchive: # TODO: make states more specific data = routine.sorted_data + if not data: + raise ValueError("No data to archive") data_dict = data.to_dict("list") - if hasattr(routine, "creation_ts"): + if hasattr(routine, "creation_ts") and routine.creation_ts is not None: suffix = routine.creation_ts else: # compatibility with old routines ts_float = data_dict["timestamp"][0] # time of the first evaluated point @@ -49,7 +62,7 @@ def archive_run(routine, states=None): # algo_name = routine.generator.name fname = f"{env_name}-{suffix}.yaml" - run = { + run: RunArchive = { "filename": fname, "routine": routine, "data": data_dict, @@ -67,14 +80,14 @@ def archive_run(routine, states=None): return run -def clear_tmp_runs(): +def clear_tmp_runs() -> None: path = os.path.join(BADGER_ARCHIVE_ROOT, ".tmp") if os.path.exists(path): for f in os.listdir(path): os.remove(os.path.join(path, f)) -def save_tmp_run(routine): +def save_tmp_run(routine: Routine) -> str: # routine: Routine path = os.path.join(BADGER_ARCHIVE_ROOT, ".tmp") suffix = time.strftime("%Y-%m-%d-%H%M%S") @@ -86,8 +99,8 @@ def save_tmp_run(routine): return fname -def list_run(): - runs = {} +def list_run() -> dict[str, dict[str, dict[str, list[str]]]]: + runs: dict[str, dict[str, dict[str, list[str]]]] = {} # Get years, latest first years = sorted( [ @@ -137,7 +150,7 @@ def list_run(): return runs -def get_runs(): +def get_runs() -> list[str]: runs = list_run() run_list = [] for months in runs.values(): @@ -176,11 +189,11 @@ def load_run(run_fname: str) -> Routine: return routine -def update_run(routine: Routine): +def update_run(routine: Routine) -> None: pass -def delete_run(run_fname): +def delete_run(run_fname: str) -> None: tokens = run_fname.split("-") first_level = tokens[1] second_level = f"{tokens[1]}-{tokens[2]}" @@ -199,7 +212,7 @@ def delete_run(run_fname): os.remove(os.path.join(prefix, run_fname)) -def get_base_run_filename(run_filename): +def get_base_run_filename(run_filename: str) -> str: if run_filename.endswith(" (failed to load)"): base_name = run_filename[:-17] else: diff --git a/src/badger/built_in_plugins/environments/sphere_2d/__init__.py b/src/badger/built_in_plugins/environments/sphere_2d/__init__.py index 0bcf25cb..6e457d0f 100644 --- a/src/badger/built_in_plugins/environments/sphere_2d/__init__.py +++ b/src/badger/built_in_plugins/environments/sphere_2d/__init__.py @@ -24,12 +24,12 @@ class Environment(environment.Environment): "g": 0.0, } - def get_variables(self, variable_names): + def get_variables(self, variable_names: list[str]) -> dict[str, float]: variable_outputs = {v: self._variables[v] for v in variable_names} return variable_outputs - def set_variables(self, variable_inputs: dict[str, float]): + def set_variables(self, variable_inputs: dict[str, float]) -> None: for var, x in variable_inputs.items(): self._variables[var] = x @@ -40,5 +40,7 @@ def set_variables(self, variable_inputs: dict[str, float]): self._observations["f"] = f self._observations["g"] = g - def get_observables(self, observable_names): + def get_observables( + self, observable_names: list[str] + ) -> dict[str, float | list[float]]: return {k: self._observations[k] for k in observable_names} diff --git a/src/badger/built_in_plugins/interfaces/default/__init__.py b/src/badger/built_in_plugins/interfaces/default/__init__.py index ae6447d3..34f39c74 100644 --- a/src/badger/built_in_plugins/interfaces/default/__init__.py +++ b/src/badger/built_in_plugins/interfaces/default/__init__.py @@ -2,6 +2,9 @@ Used by test environments and as the fallback when no hardware interface is configured.""" +from collections.abc import Mapping +from typing import Any + from badger import interface @@ -10,14 +13,14 @@ class Interface(interface.Interface): # If params not specified, it would be an empty dict # Private variables - _states: dict + _states: dict[str, float | list[float]] - def __init__(self, **data): + def __init__(self, **data: Any) -> None: super().__init__(**data) - self._states = {} + self._states: dict[str, float | list[float]] = {} - def get_values(self, channel_names): + def get_values(self, channel_names: list[str]) -> dict[str, float | list[float]]: channel_outputs = {} for channel in channel_names: @@ -30,6 +33,6 @@ def get_values(self, channel_names): return channel_outputs - def set_values(self, channel_inputs): + def set_values(self, channel_inputs: Mapping[str, float | list[float]]) -> None: for channel, value in channel_inputs.items(): self._states[channel] = value diff --git a/src/badger/core.py b/src/badger/core.py index 76c47d4b..fe8b711f 100644 --- a/src/badger/core.py +++ b/src/badger/core.py @@ -11,18 +11,19 @@ import time from collections.abc import Callable +from typing import Any from pandas import DataFrame, concat from xopt.vocs import select_best from badger.errors import BadgerRunTerminated from badger.logger import _get_default_logger -from badger.logger.event import Events +from badger.logger.event import Events, Solution from badger.routine import Routine from badger.utils import curr_ts_to_str, dump_state -def check_run_status(active_callback): +def check_run_status(active_callback: Callable[[], int]) -> None: while True: status = active_callback() if status == 2: @@ -34,13 +35,14 @@ def check_run_status(active_callback): break -def convert_to_solution(result: DataFrame, routine: Routine): +def convert_to_solution(result: DataFrame, routine: Routine) -> Solution: vocs = routine.vocs + try: best_idx, _, _ = select_best(vocs, routine.sorted_data, n=1) if best_idx.size > 0: best_idx = int(best_idx[0]) # convert numpy array to int - if best_idx != len(routine.data) - 1: + if routine.data is None or best_idx != len(routine.data) - 1: is_optimal = False else: is_optimal = True @@ -56,7 +58,7 @@ def convert_to_solution(result: DataFrame, routine: Routine): cons = list(result[vocs.constraint_names].to_numpy()[0]) stas = list(result[vocs.observable_names].to_numpy()[0]) - solution = ( + solution = Solution( vars, objs, cons, @@ -73,11 +75,11 @@ def convert_to_solution(result: DataFrame, routine: Routine): def run_routine( routine: Routine, - active_callback: Callable, - generate_callback: Callable, - evaluate_callback: Callable, - states_callback: Callable, - dump_file_callback: Callable | None = None, + active_callback: Callable[..., Any], + generate_callback: Callable[..., Any], + evaluate_callback: Callable[..., Any], + states_callback: Callable[..., Any], + dump_file_callback: Callable[..., Any] | None = None, verbose: int = 2, ) -> None: """ @@ -115,32 +117,32 @@ def run_routine( # Save system states if applicable states = environment.get_system_states() - if states_callback and (states is not None): + if states is not None: states_callback(states) # Optimization starts - solution_meta = ( - None, + solution_meta = Solution( None, None, None, None, + False, # TODO: was previously None, but type of is_optimal is bool, need to check if this causes any issue routine.vocs.variable_names, routine.vocs.objective_names, routine.vocs.constraint_names, routine.vocs.observable_names, ) + solution_meta = Solution(*solution_meta) opt_logger.update(Events.OPTIMIZATION_START, solution_meta) # evaluate initial points: # Nikita: more care about the setting var logic, # wait or consider timeout/retry - for _, ele in initial_points.iterrows(): - result = routine.evaluate_data(ele.to_dict()) + for _, ele in initial_points.iterrows() if initial_points is not None else []: + result = routine.evaluate_data(DataFrame([ele.to_dict()])) solution = convert_to_solution(result, routine) opt_logger.update(Events.OPTIMIZATION_STEP, solution) - if evaluate_callback: - evaluate_callback(result) + evaluate_callback(result) # Prepare for dumping file if dump_file_callback: @@ -172,8 +174,7 @@ def run_routine( result = routine.evaluate_data(candidates) solution = convert_to_solution(result, routine) opt_logger.update(Events.OPTIMIZATION_STEP, solution) - if evaluate_callback: - evaluate_callback(result) + evaluate_callback(result) # Dump Xopt state after each step if dump_file_callback: diff --git a/src/badger/core_subprocess.py b/src/badger/core_subprocess.py index c67d7d7a..7d2c9000 100644 --- a/src/badger/core_subprocess.py +++ b/src/badger/core_subprocess.py @@ -18,6 +18,9 @@ import time import traceback from copy import deepcopy +from dataclasses import dataclass +from multiprocessing.connection import Connection +from multiprocessing.synchronize import Event from queue import Empty from typing import Any @@ -28,23 +31,29 @@ from badger.errors import ( MEASUREMENT_ACTION_ABORT, MEASUREMENT_ACTION_RETRY, - MEASUREMENT_ACTION_TYPE, - MEASUREMENT_ERROR_TYPE, TERMINATION_ACTION_CONTINUE, TERMINATION_ACTION_END, - TERMINATION_ACTION_TYPE, - TERMINATION_REACHED_TYPE, BadgerEnvObsError, BadgerRunTerminated, ) from badger.log import configure_process_logging from badger.logger import _get_default_logger -from badger.logger.event import Events +from badger.logger.event import Events, Solution from badger.routine import Routine from badger.settings import ( apply_pytorch_multiprocess_tensor_sharing_setting, init_settings, ) +from badger.types import ( + ArgumentQueueType, + DataQueueMessage, + DialogActionMessage, + MeasurementActionMessage, + MeasurementErrorMessage, + RoutineErrorMessage, + TerminationActionMessage, + TerminationReachedMessage, +) logger = logging.getLogger(__name__) @@ -52,9 +61,9 @@ def evaluate_measurement_with_retry( routine: Routine, point: Any, - queue: mp.Queue, - stop_process: mp.Event, - dialog_action_queue: mp.Queue, + data_and_error_queue: mp.Queue[DataQueueMessage], + stop_process: Event, + dialog_action_queue: mp.Queue[DialogActionMessage], ) -> DataFrame: while True: try: @@ -63,12 +72,11 @@ def evaluate_measurement_with_retry( error_title = f"{type(e).__name__}: {e}" error_traceback = traceback.format_exc() logger.error(f"Measurement failed: {error_title}\n{error_traceback}") - queue.put( - { - "type": MEASUREMENT_ERROR_TYPE, - "title": error_title, - "traceback": error_traceback, - } + data_and_error_queue.put( + MeasurementErrorMessage( + title=error_title, + traceback=error_traceback, + ) ) # Wait for response back from retry dialog @@ -82,32 +90,35 @@ def evaluate_measurement_with_retry( except Empty: continue - if ( - isinstance(msg, dict) - and msg.get("type") == MEASUREMENT_ACTION_TYPE - and msg.get("action") - in [MEASUREMENT_ACTION_RETRY, MEASUREMENT_ACTION_ABORT] - ): - if msg["action"] == MEASUREMENT_ACTION_RETRY: + if isinstance(msg, MeasurementActionMessage) and msg.action in [ + MEASUREMENT_ACTION_RETRY, + MEASUREMENT_ACTION_ABORT, + ]: + if msg.action == MEASUREMENT_ACTION_RETRY: break raise BadgerRunTerminated( f"Run terminated after measurement error: {error_title}" ) +@dataclass +class TerminationCondition: + type: str + condition: dict[str, Any] + + def pause_for_termination_dialog_action( - queue: mp.Queue, - stop_process: mp.Event, - pause_process: mp.Event, - dialog_action_queue: mp.Queue, - tc_condition: dict, + data_and_error_queue: mp.Queue[DataQueueMessage], + stop_process: Event, + pause_process: Event, + dialog_action_queue: mp.Queue[DialogActionMessage], + tc_condition: TerminationCondition, ) -> None: """Pause the run and wait for user action when run-until condition is reached.""" - queue.put( - { - "type": TERMINATION_REACHED_TYPE, - "tc_condition": tc_condition, - } + data_and_error_queue.put( + TerminationReachedMessage( + tc_condition={"type": tc_condition.type, **tc_condition.condition}, + ) ) while True: @@ -121,13 +132,11 @@ def pause_for_termination_dialog_action( except Empty: continue - if ( - isinstance(msg, dict) - and msg.get("type") == TERMINATION_ACTION_TYPE - and msg.get("action") - in [TERMINATION_ACTION_CONTINUE, TERMINATION_ACTION_END] - ): - if msg["action"] == TERMINATION_ACTION_CONTINUE: + if isinstance(msg, TerminationActionMessage) and msg.action in [ + TERMINATION_ACTION_CONTINUE, + TERMINATION_ACTION_END, + ]: + if msg.action == TERMINATION_ACTION_CONTINUE: pause_process.set() return @@ -136,7 +145,7 @@ def pause_for_termination_dialog_action( ) -def convert_to_solution(result: DataFrame, routine: Routine): +def convert_to_solution(result: DataFrame, routine: Routine) -> Solution: """ This method is passed the latest evaluated solution and converts that to a printable format for the terminal. This method is for the GUI version of Badger. @@ -146,11 +155,13 @@ def convert_to_solution(result: DataFrame, routine: Routine): result : DataFrame routine : Routine """ + if routine.sorted_data is None or routine.sorted_data.empty: + raise ValueError("Routine has no sorted data") vocs = routine.vocs try: best_idx, _, _ = select_best(vocs, routine.sorted_data, n=1) logger.debug(f"Selected best index: {best_idx}") - if best_idx.size > 0: + if best_idx.size > 0 and routine.data is not None: if best_idx[0] != len(routine.data) - 1: is_optimal = False else: @@ -172,7 +183,7 @@ def convert_to_solution(result: DataFrame, routine: Routine): stas = list(result[vocs.observable_names].to_numpy()[0]) # TODO: This structure needs improvement - solution = ( + solution = Solution( vars, objs, cons, @@ -188,28 +199,30 @@ def convert_to_solution(result: DataFrame, routine: Routine): def run_routine_subprocess( - args_queue: mp.Queue, - queue: mp.Queue, - evaluate_queue: mp.Pipe, - stop_process: mp.Event, - pause_process: mp.Event, - wait_event: mp.Event, + args_queue: mp.Queue[ArgumentQueueType], + data_and_error_queue: mp.Queue[DataQueueMessage], + evaluate_queue: tuple[Connection, Connection], + stop_process: Event, + pause_process: Event, + wait_event: Event, + dialog_action_queue: mp.Queue[DialogActionMessage], config_path: str | None = None, - log_queue: mp.Queue | None = None, - dialog_action_queue: mp.Queue | None = None, + log_queue: mp.Queue[logging.LogRecord] | None = None, ) -> None: """ Run the provided routine object using Xopt. This method is run as a subproccess Parameters ---------- - queue: mp.Queue - evaluate_queue: mp.Pipe - stop_process: mp.Event - pause_process: mp.Event - wait_event: mp.Event - config_path: str - log_queue: mp.Queue + args_queue: mp.Queue[ArgumentQueueType], + data_and_error_queue: mp.Queue[DataQueueMessage], + evaluate_queue: tuple[Connection, Connection], + stop_process: Event, + pause_process: Event, + wait_event: Event, + dialog_action_queue: mp.Queue[DialogActionMessage], + config_path: str | None = None, + log_queue: mp.Queue[logging.LogRecord] | None = None, """ # Setup logging for this subprocess if log_queue is not None: @@ -218,7 +231,7 @@ def run_routine_subprocess( logger_name=__name__, # Always make this level DEBUG so no logs are filtered out until get sent to main, # where logs from all sub-processes get filtered before written. - log_level=logging.DEBUG, + log_level="DEBUG", process_name=f"{os.path.basename(__name__)}-{mp.current_process().pid}", ) @@ -238,17 +251,17 @@ def run_routine_subprocess( logger.info("Waiting for wait_event to be set...") wait_event.wait() - args: dict[str, Any] = {} try: args = args_queue.get(timeout=1) logger.debug(f"Received args from queue: {args}") - except Exception as e: # noqa: BLE001 - subprocess queue read boundary + except Exception as e: logger.error(f"Error in subprocess queue.get: {type(e).__name__}, {e!s}") + raise RuntimeError("Failed to get arguments from queue") from e # set required arguments try: - logger.info(f"Loading routine from file: {args['routine_filename']}") - routine = load_run(args["routine_filename"]) + logger.info(f"Loading routine from file: {args.routine_filename}") + routine = load_run(args.routine_filename) logger.info("Resetting environment global state") routine.environment.reset_environment() if routine.vrange_hard_limit: @@ -258,7 +271,7 @@ def run_routine_subprocess( routine.environment.variables.update(routine.vrange_hard_limit) # Reset data if run_data option is False - if not args["run_data"] and routine.data is not None: + if not args.run_data and routine.data is not None: logger.info("Resetting routine data") routine.data = routine.data.iloc[0:0] # reset the data @@ -266,7 +279,9 @@ def run_routine_subprocess( error_title = f"{type(e).__name__}: {e}" error_traceback = traceback.format_exc() logger.error(f"Error initializing routine: {error_title}\n{error_traceback}") - queue.put((error_title, error_traceback)) + data_and_error_queue.put( + RoutineErrorMessage(title=error_title, traceback=error_traceback) + ) raise # TODO look into this bug with serializing of turbo. Fix might be needed in Xopt @@ -283,18 +298,18 @@ def run_routine_subprocess( logger.warning("TypeError when converting turbo_controller dtype") # Assign the initial points and bounds - logger.info(f"Setting routine variable ranges: {args['variable_ranges']}") - routine.vocs.variables = args["variable_ranges"] + logger.info(f"Setting routine variable ranges: {args.variable_ranges}") + routine.vocs.variables = args.variable_ranges logger.info("Setting routine initial points") - routine.initial_points = args["initial_points"] + routine.initial_points = args.initial_points # set optional arguments - evaluate = args.pop("evaluate", None) - archive = args.pop("archive", False) - termination_condition = args.pop("termination_condition", None) - start_time = args.pop("start_time", None) - verbose = args.pop("verbose", 2) - testing = args.pop("testing", False) + evaluate = args.evaluate + archive = args.archive + termination_condition = args.termination_condition + start_time = args.start_time + verbose = 2 # default value as before + testing = args.testing # setup variables of routine properties for code readablilty initial_points = routine.initial_points @@ -305,12 +320,12 @@ def run_routine_subprocess( # Optimization starts # This is used by the logger to print to the terminal. - solution_meta = ( - None, + solution_meta = Solution( None, None, None, None, + False, # TODO: Was set to None before, but that doesn't make sense given the context. Need to verify this is correct. routine.vocs.variable_names, routine.vocs.objective_names, routine.vocs.constraint_names, @@ -323,24 +338,32 @@ def run_routine_subprocess( # timeout logic will be handled in the specific environment try: # initial sampling - if args["init_points"]: + if args.init_points: logger.info("Evaluating initial points...") - for _, ele in initial_points.iterrows(): - logger.debug(f"Evaluating initial point: {ele.to_dict()}") - result = evaluate_measurement_with_retry( - routine, ele.to_dict(), queue, stop_process, dialog_action_queue - ) - solution = convert_to_solution(result, routine) - opt_logger.update(Events.OPTIMIZATION_STEP, solution) - if evaluate: - time.sleep(0.1) # give it some break tp catch up - evaluate_queue[0].send((routine.data, routine.generator)) + if initial_points is None or initial_points.empty: + logger.info("No initial points provided.") + else: + for _, ele in initial_points.iterrows(): + logger.debug(f"Evaluating initial point: {ele.to_dict()}") + result = evaluate_measurement_with_retry( + routine, + ele.to_dict(), + data_and_error_queue, + stop_process, + dialog_action_queue, + ) + solution = convert_to_solution(result, routine) + opt_logger.update(Events.OPTIMIZATION_STEP, solution) + if evaluate: + time.sleep(0.1) # give it some break tp catch up + evaluate_queue[0].send((routine.data, routine.generator)) logger.info("Starting optimization loop...") while True: if stop_process.is_set(): logger.info("Stop process set. Terminating optimization.") evaluate_queue[0].close() + pause_process.clear() raise BadgerRunTerminated elif not pause_process.is_set(): logger.info("Pause process not set. Waiting...") @@ -348,9 +371,9 @@ def run_routine_subprocess( if termination_condition and start_time: tc_config = termination_condition - idx = tc_config["tc_idx"] + idx = tc_config.tc_idx if idx == 0: - max_eval = tc_config["max_eval"] + max_eval = tc_config.max_eval if routine.data is not None: if "live" in routine.data.columns: # Only count number of live data points @@ -370,22 +393,26 @@ def run_routine_subprocess( "Max evaluations reached. Pausing optimization and waiting for user action." ) pause_process.clear() + tc_condition = TerminationCondition( + type="max_eval", + condition={ + "config": max_eval, + "state": count, + }, + ) + pause_for_termination_dialog_action( - queue=queue, + data_and_error_queue=data_and_error_queue, stop_process=stop_process, pause_process=pause_process, dialog_action_queue=dialog_action_queue, - tc_condition={ - "type": "max_eval", - "config": max_eval, - "state": count, - }, + tc_condition=tc_condition, ) # reset termination condition termination_condition = None continue elif idx == 1: - max_time = tc_config["max_time"] + max_time = tc_config.max_time dt = time.time() - start_time logger.debug(f"Checking max_time termination: {dt} >= {max_time}") if dt >= max_time: @@ -394,15 +421,17 @@ def run_routine_subprocess( ) pause_process.clear() pause_for_termination_dialog_action( - queue=queue, + data_and_error_queue=data_and_error_queue, stop_process=stop_process, pause_process=pause_process, dialog_action_queue=dialog_action_queue, - tc_condition={ - "type": "max_time", - "config": max_time, - "state": dt, - }, + tc_condition=TerminationCondition( + type="max_time", + condition={ + "config": max_time, + "state": dt, + }, + ), ) # reset termination condition termination_condition = None @@ -423,7 +452,11 @@ def run_routine_subprocess( pause_process.wait() result = evaluate_measurement_with_retry( - routine, candidates, queue, stop_process, dialog_action_queue + routine, + candidates, + data_and_error_queue, + stop_process, + dialog_action_queue, ) solution = convert_to_solution(result, routine) opt_logger.update(Events.OPTIMIZATION_STEP, solution) @@ -446,7 +479,9 @@ def run_routine_subprocess( logger.error(f"XoptError during optimization: {e}") opt_logger.update(Events.OPTIMIZATION_END, solution_meta) error_title = "BadgerEnvObsError: There was an error getting observables from the environment. See the traceback for more details." - queue.put((error_title, traceback.format_exc())) + data_and_error_queue.put( + RoutineErrorMessage(title=error_title, traceback=traceback.format_exc()) + ) evaluate_queue[0].close() raise BadgerEnvObsError(e) except Exception as e: @@ -456,6 +491,8 @@ def run_routine_subprocess( opt_logger.update(Events.OPTIMIZATION_END, solution_meta) error_title = f"{type(e).__name__}: {e}" error_traceback = traceback.format_exc() - queue.put((error_title, error_traceback)) + data_and_error_queue.put( + RoutineErrorMessage(title=error_title, traceback=error_traceback) + ) evaluate_queue[0].close() raise diff --git a/src/badger/db.py b/src/badger/db.py index 28d8b6dd..ac20fb45 100644 --- a/src/badger/db.py +++ b/src/badger/db.py @@ -11,7 +11,9 @@ import sqlite3 import uuid import warnings +from collections.abc import Callable from datetime import UTC, datetime +from typing import Any, ParamSpec, TypeVar, cast import yaml @@ -24,26 +26,35 @@ # Check badger database root flag_use_db = True + config_singleton = init_settings() try: - BADGER_DB_ROOT = config_singleton.read_value("BADGER_DB_ROOT") + badger_db_root = cast(str, config_singleton.read_value("BADGER_DB_ROOT")) except KeyError: + badger_db_root = "" flag_use_db = False +BADGER_DB_ROOT = badger_db_root + if flag_use_db: - if BADGER_DB_ROOT is None: + if BADGER_DB_ROOT == "": raise BadgerConfigError("Please set the BADGER_DB_ROOT env var!") elif not os.path.exists(BADGER_DB_ROOT): os.makedirs(BADGER_DB_ROOT) logger.info(f"Badger database root {BADGER_DB_ROOT} created") +T = TypeVar("T") +P = ParamSpec("P") -def ensure_routines_db_exists(func): + +def ensure_routines_db_exists( + func: Callable[P, T], +) -> Callable[P, T]: """ Create the routines database (a SQL table) if it does not already exist. """ - def func_safe(*args, **kwargs): + def func_safe(*args: P.args, **kwargs: P.kwargs) -> T: db_routine = os.path.join(BADGER_DB_ROOT, "routines.db") con = sqlite3.connect(db_routine) @@ -61,12 +72,12 @@ def func_safe(*args, **kwargs): return func_safe -def ensure_runs_db_exists(func): +def ensure_runs_db_exists(func: Callable[P, T]) -> Callable[P, T]: """ Create the runs database (a SQL table) if it does not already exist. """ - def func_safe(*args, **kwargs): + def func_safe(*args: P.args, **kwargs: P.kwargs) -> T: db_run = os.path.join(BADGER_DB_ROOT, "runs.db") con = sqlite3.connect(db_run) @@ -84,7 +95,7 @@ def func_safe(*args, **kwargs): return func_safe -def filter_routines(records, tags): +def filter_routines(records: list[tuple], tags: dict) -> list[tuple]: records_filtered = [] for record in records: try: @@ -97,7 +108,7 @@ def filter_routines(records, tags): return records_filtered -def extract_metadata(records): +def extract_metadata(records: list[tuple]) -> tuple[list[str], list[str]]: env_list = [] descr_list = [] for record in records: @@ -115,7 +126,7 @@ def extract_metadata(records): @ensure_routines_db_exists -def save_routine(routine: Routine): +def save_routine(routine: Routine) -> None: db_routine = os.path.join(BADGER_DB_ROOT, "routines.db") con = sqlite3.connect(db_routine) @@ -134,7 +145,7 @@ def save_routine(routine: Routine): # This function is not safe and might break database! Use with caution! @ensure_routines_db_exists -def update_routine(routine: Routine): +def update_routine(routine: Routine) -> None: db_routine = os.path.join(BADGER_DB_ROOT, "routines.db") con = sqlite3.connect(db_routine) @@ -155,7 +166,7 @@ def update_routine(routine: Routine): @ensure_routines_db_exists @ensure_runs_db_exists -def remove_routine(id: str, remove_runs=True): +def remove_routine(id: str, remove_runs: bool = True) -> None: db_routine = os.path.join(BADGER_DB_ROOT, "routines.db") con = sqlite3.connect(db_routine) @@ -180,7 +191,7 @@ def remove_routine(id: str, remove_runs=True): @ensure_routines_db_exists -def load_routine(id: str): +def load_routine(id: str) -> tuple[Routine, Any]: db_routine = os.path.join(BADGER_DB_ROOT, "routines.db") con = sqlite3.connect(db_routine) cur = con.cursor() @@ -215,7 +226,7 @@ def load_routine(id: str): @ensure_routines_db_exists -def list_routine(keyword="", tags: dict[str, str] | None = None): +def list_routine(keyword: str = "", tags: dict[str, str] | None = None) -> list[tuple]: if tags is None: tags = {} db_routine = os.path.join(BADGER_DB_ROOT, "routines.db") @@ -307,7 +318,7 @@ def list_routine(keyword="", tags: dict[str, str] | None = None): @ensure_runs_db_exists -def save_run(run): +def save_run(run: dict[str, Any]) -> int | None: db_run = os.path.join(BADGER_DB_ROOT, "runs.db") con = sqlite3.connect(db_run, timeout=30.0) @@ -324,6 +335,8 @@ def save_run(run): cur.execute("select id from run where filename = ?", (run_filename,)) existing_row = cur.fetchone() + rid: int | None = None + if existing_row: cur.execute( "update run set finishedAt = ? where filename = ?", @@ -344,7 +357,7 @@ def save_run(run): @ensure_runs_db_exists -def get_runs_by_routine(routine_id: str): +def get_runs_by_routine(routine_id: str) -> list[str]: db_run = os.path.join(BADGER_DB_ROOT, "runs.db") con = sqlite3.connect(db_run) @@ -362,7 +375,7 @@ def get_runs_by_routine(routine_id: str): @ensure_runs_db_exists -def get_runs(): +def get_runs() -> list[str]: db_run = os.path.join(BADGER_DB_ROOT, "runs.db") con = sqlite3.connect(db_run) @@ -379,7 +392,7 @@ def get_runs(): @ensure_runs_db_exists -def remove_run_by_filename(filename): +def remove_run_by_filename(filename: str) -> None: db_run = os.path.join(BADGER_DB_ROOT, "runs.db") con = sqlite3.connect(db_run) @@ -392,7 +405,7 @@ def remove_run_by_filename(filename): @ensure_runs_db_exists -def remove_run_by_id(rid): +def remove_run_by_id(rid: int) -> None: db_run = os.path.join(BADGER_DB_ROOT, "runs.db") con = sqlite3.connect(db_run) @@ -404,7 +417,7 @@ def remove_run_by_id(rid): con.close() -def import_routines(filename): +def import_routines(filename: str) -> None: con = sqlite3.connect(filename) cur = con.cursor() @@ -436,7 +449,7 @@ def import_routines(filename): raise BadgerDBError(get_yaml_string(failed_list)) -def export_routines(filename, routine_id_list): +def export_routines(filename: str, routine_id_list: list[str]) -> None: con = sqlite3.connect(filename) cur = con.cursor() diff --git a/src/badger/environment.py b/src/badger/environment.py index bfd7d628..23746fe7 100644 --- a/src/badger/environment.py +++ b/src/badger/environment.py @@ -12,7 +12,8 @@ import logging from abc import abstractmethod -from typing import TYPE_CHECKING, Any, ClassVar +from collections.abc import Callable +from typing import TYPE_CHECKING, Any, ClassVar, cast from pydantic import BaseModel, ConfigDict, Field, SerializeAsAny from pydantic._internal._model_construction import ModelMetaclass @@ -26,13 +27,15 @@ from badger.factory import BadgerPluginConfig from badger.formula import extract_variable_keys, interpret_expression -from badger.interface import Interface +from badger.interface import Interface, InterfaceInfo logger = logging.getLogger(__name__) -def validate_setpoints(func): - def validate(cls, variable_inputs: dict[str, float]): +def validate_setpoints( + func: Callable[[Any, dict[str, float]], Any], +) -> Callable[[Any, dict[str, float]], Any]: + def validate(cls: Any, variable_inputs: dict[str, float]) -> Any: _bounds = cls.get_bounds(list(variable_inputs.keys())) for name, value in variable_inputs.items(): lower = _bounds[name][0] @@ -49,13 +52,17 @@ def validate(cls, variable_inputs: dict[str, float]): return validate -def process_formulas(func): +def process_formulas( + func: Callable[[Any, list[str]], dict[str, float | list[float]]], +) -> Callable[[Any, list[str]], dict[str, float | list[float]]]: """ Decorator function that wraps get_observables method to process formulas if they exist in the observable names. """ - def process(cls, observable_names: list[str]) -> dict[str, float]: + def process( + cls: Any, observable_names: list[str] + ) -> dict[str, float | list[float]]: # get the list of observable names needed by themselves and any formulas formula_observables = [] basic_observables = [] @@ -95,8 +102,10 @@ def process(cls, observable_names: list[str]) -> dict[str, float]: return process -def validate_bounds(func): - def validate(cls, variable_names: list[str]): +def validate_bounds( + func: Callable[[Any, list[str]], dict[str, tuple[float, float]]], +) -> Callable[[Any, list[str]], dict[str, tuple[float, float]]]: + def validate(cls: Any, variable_names: list[str]) -> dict[str, tuple[float, float]]: bounds = func(cls, variable_names) for name, bound in bounds.items(): @@ -120,7 +129,10 @@ def validate(cls, variable_names: list[str]): class EnvMeta(ModelMetaclass): - def __new__(mcs, name: str, bases: tuple[type, ...], namespace: dict[str, Any]): + # mypy forbids Self as a metaclass __new__ return type, so return EnvMeta. + def __new__( # noqa: PYI034 + mcs, name: str, bases: tuple[type, ...], namespace: dict[str, Any] + ) -> "EnvMeta": # Wrap get_bounds with validate_bounds if defined if "get_bounds" in namespace: namespace["get_bounds"] = validate_bounds(namespace["get_bounds"]) @@ -132,7 +144,7 @@ def __new__(mcs, name: str, bases: tuple[type, ...], namespace: dict[str, Any]): # Wrap set_variables with validate_setpoints if defined if "set_variables" in namespace: namespace["set_variables"] = validate_setpoints(namespace["set_variables"]) - return super().__new__(mcs, name, bases, namespace) + return super().__new__(mcs, name, bases, namespace) # type: ignore[return-value] class BaseEnvironment(BaseModel, metaclass=EnvMeta): @@ -166,7 +178,7 @@ def get_variables(self, variable_names: list[str]) -> dict[str, float]: """ @abstractmethod - def set_variables(self, variable_inputs: dict[str, float]): + def set_variables(self, variable_inputs: dict[str, float]) -> None: """ Set the values of the specified variables in the environment. @@ -188,17 +200,17 @@ def get_observables( Parameters ---------- - observable_names : List[str] + observable_names : list[str] A list of observable names to retrieve from the environment. Returns ------- - Dict[str, float | List[float]] + dict[str, float | list[float]] A dictionary mapping observable names to their values. """ - def reset_environment(self): + def reset_environment(self) -> None: """ Reset the environment to its initial state. This method is called at the start of each run. @@ -211,14 +223,12 @@ def get_system_states(self) -> dict[str, Any]: Returns ------- - Dict[str, Any] + dict[str, Any] A dictionary containing the current system states. """ return {} - def get_bounds( - self, variable_names: dict[str, list[float]] - ) -> dict[str, list[float]]: + def get_bounds(self, variable_names: list[str]) -> dict[str, list[float]]: """ Get the bounds for the specified variables in the environment. The bounds are returned as a dictionary with variable names as keys @@ -232,7 +242,7 @@ def get_bounds( """ return {name: self.variables[name] for name in variable_names} - def search(self, keyword: str, callback: callable): + def search(self, keyword: str, callback: Callable[[Any], None]) -> None: """ Search for a keyword in the environment and call the callback function with the results. @@ -267,7 +277,7 @@ def get_variable(self, variable_name: str) -> float: """ return self.get_variables([variable_name])[variable_name] - def set_variable(self, variable_name: str, value: float): + def set_variable(self, variable_name: str, value: float) -> None: """ Set the value of a single variable in the environment. @@ -280,7 +290,7 @@ def set_variable(self, variable_name: str, value: float): """ self.set_variables({variable_name: value}) - def get_observable(self, observable_name: str) -> float: + def get_observable(self, observable_name: str) -> float | list[float]: """ Get the value of a single observable from the environment. @@ -307,41 +317,45 @@ class Environment(BaseEnvironment): # Optional methods to inherit ############################################################ - def get_variables(self, variable_names: list[str]) -> dict: + def get_variables(self, variable_names: list[str]) -> dict[str, float]: if not self.interface: raise BadgerNoInterfaceError - return self.interface.get_values(variable_names) + # Variables are scalar setpoints, but the interface exposes the broader + # channel type shared with observables. + return cast(dict[str, float], self.interface.get_values(variable_names)) - def set_variables(self, variable_inputs: dict[str, float]): + def set_variables(self, variable_inputs: dict[str, float]) -> None: if not self.interface: raise BadgerNoInterfaceError return self.interface.set_values(variable_inputs) - def get_observables(self, observable_names: list[str]) -> dict: + def get_observables( + self, observable_names: list[str] + ) -> dict[str, float | list[float]]: if not self.interface: raise BadgerNoInterfaceError return self.interface.get_values(observable_names) - def reset_environment(self): + def reset_environment(self) -> None: if self.interface: return self.interface.reset_interface() - def get_info(self, variable_names: list[str]) -> dict | None: + def get_info(self, variable_names: list[str]) -> InterfaceInfo | None: if not self.interface: return None return self.interface.get_info(variable_names) @property - def variable_names(self): + def variable_names(self) -> list[str]: return [k for k in self.variables] def instantiate_env( - env_class: type[Environment], configs: "BadgerPluginConfig", manager=None + env_class: type[Environment], configs: "BadgerPluginConfig", manager: Any = None ) -> Environment: # Configure interface # TODO: figure out the correct logic diff --git a/src/badger/errors.py b/src/badger/errors.py index c7a8efed..d05408cb 100644 --- a/src/badger/errors.py +++ b/src/badger/errors.py @@ -10,7 +10,7 @@ class BadgerError(Exception): - def __init__(self, message="", detailed_text=None): + def __init__(self, message: str = "", detailed_text: str | None = None) -> None: if detailed_text is None: detailed_text = self.capture_traceback_or_stack() @@ -18,7 +18,7 @@ def __init__(self, message="", detailed_text=None): self.detailed_text = detailed_text self.show_message_box() - def show_message_box(self): + def show_message_box(self) -> None: """ Method to create and display a popup window with the error message. """ @@ -33,7 +33,7 @@ def show_message_box(self): dialog.setIcon(QMessageBox.Critical) dialog.exec_() - def capture_traceback_or_stack(self): + def capture_traceback_or_stack(self) -> str: """ Captures the current traceback if an exception is active, otherwise captures the call stack. """ @@ -79,7 +79,7 @@ class BadgerEnvObsError(Exception): class BadgerNoInterfaceError(Exception): - def __init__(self, message="Must provide an interface!"): + def __init__(self, message: str = "Must provide an interface!") -> None: super().__init__(message) @@ -106,7 +106,7 @@ class BadgerLogbookError(Exception): class BadgerLoadConfigError(Exception): - def __init__(self, message="Config Error!"): + def __init__(self, message: str = "Config Error!") -> None: super().__init__(message) @@ -115,7 +115,7 @@ class BadgerRoutineError(Exception): class BadgerRunTerminated(Exception): - def __init__(self, message="Optimization run has been terminated!"): + def __init__(self, message: str = "Optimization run has been terminated!") -> None: super().__init__(message) @@ -132,3 +132,6 @@ def __init__(self, message="Optimization run has been terminated!"): TERMINATION_ACTION_TYPE = "termination_action" TERMINATION_ACTION_CONTINUE = "continue" TERMINATION_ACTION_END = "end" + +# Type tag for fatal errors that end a run, sent from subprocess to routine runner. +ROUTINE_ERROR_TYPE = "routine_error" diff --git a/src/badger/extension.py b/src/badger/extension.py index 1d60e312..51bf1812 100644 --- a/src/badger/extension.py +++ b/src/badger/extension.py @@ -3,30 +3,36 @@ Third-party packages can subclass Extension to plug in custom algorithms.""" from abc import ABC, abstractmethod +from collections.abc import Callable +from typing import Any class Extension(ABC): @property @abstractmethod def name(self) -> str: - pass + raise NotImplementedError("Extension must implement name property") @abstractmethod - def __init__(self): - pass + def __init__(self) -> None: + raise NotImplementedError("Extension must implement __init__ method") # List all available generators @abstractmethod def list_generator(self) -> list[str]: - pass + raise NotImplementedError("Extension must implement list_generator method") # Get config of an generator @abstractmethod - def get_generator_config(self, name: str): - pass + def get_generator_config(self, name: str) -> dict[str, Any]: + raise NotImplementedError( + "Extension must implement get_generator_config method" + ) # Run an optimization with an array-style evaluate function # and a configs dict @abstractmethod - def optimize(self, evaluate, configs: dict): - pass + def optimize( + self, evaluate: Callable[[Any], Any], configs: dict[str, Any] + ) -> dict[str, Any]: + raise NotImplementedError("Extension must implement optimize method") diff --git a/src/badger/factory.py b/src/badger/factory.py index cfa2e5fe..a9550fa2 100644 --- a/src/badger/factory.py +++ b/src/badger/factory.py @@ -15,9 +15,10 @@ import re import sys from pathlib import Path -from typing import TYPE_CHECKING, Any, TypedDict, cast +from typing import TYPE_CHECKING, Any, Literal, TypedDict, cast import yaml +from xopt.generator import Generator from xopt.generators import generators, get_generator_defaults from badger.errors import ( @@ -47,7 +48,7 @@ ] -class BadgerPluginConfig(TypedDict): +class BadgerPluginConfig(TypedDict, total=False): name: str description: str version: str @@ -75,10 +76,21 @@ class BadgerPluginConfig(TypedDict): sys.path.append(BADGER_PLUGIN_ROOT) -def scan_plugins(root: str): - factory: dict[str, Any] = {} +class BadgerFactoryType(TypedDict): + generator: dict[str, dict[str, tuple[type[Generator], BadgerPluginConfig | None]]] + interface: dict[ + str, dict[str, tuple[type["BadgerInterface"], BadgerPluginConfig | None]] + ] + environment: dict[ + str, dict[str, tuple[type["BadgerEnvironment"], BadgerPluginConfig | None]] + ] + + +def scan_plugins(root: str) -> BadgerFactoryType: + factory: BadgerFactoryType = {"environment": {}, "interface": {}, "generator": {}} # Do not scan local generators if option disabled + ptype_list: list[Literal["generator", "interface", "environment"]] if LOAD_LOCAL_ALGO: ptype_list = ["generator", "interface", "environment"] else: @@ -86,7 +98,7 @@ def scan_plugins(root: str): factory["generator"] = {} for ptype in ptype_list: - factory[ptype] = {} + factory[ptype] = {} # type: ignore proot = os.path.join(root, f"{ptype}s") @@ -102,14 +114,14 @@ def scan_plugins(root: str): for pname in plugins: # TODO: Also load the configs here # So that list plugins can access the metadata of the plugins - factory[ptype][pname] = None + factory[ptype][pname] = None # type: ignore return factory def load_plugin( - root: str, pname: str, ptype: str -) -> tuple[Any | None, BadgerPluginConfig | None]: + root: str, pname: str, ptype: Literal["environment", "interface", "generator"] +) -> "tuple[type[BadgerEnvironment | BadgerInterface | Generator] | None, BadgerPluginConfig | None]": assert ptype in [ "generator", "interface", @@ -212,6 +224,7 @@ def load_badger_docs(name: str, ptype: str | None = None) -> str: Formatted markdown string containing both the README content and the plugin class docstring if applicable in a code block. """ + # .../Badger/src/badger/factory.py PROJECT_ROOT = Path(__file__).resolve().parent.parent.parent BADGER_GUIDES_DIR = PROJECT_ROOT / "documentation" / "docs" / "guides" @@ -223,6 +236,12 @@ def load_badger_docs(name: str, ptype: str | None = None) -> str: if subdir.is_dir(): docs_dir = subdir + assert ptype in [ + "generator", + "interface", + "environment", + ], f"Invalid plugin type {ptype}" + # Create header with links to other guides files = [x.stem for x in BADGER_GUIDES_DIR.iterdir() if str(x).endswith(".md")] headers = [f"{filename.title()}" for filename in files] @@ -241,7 +260,7 @@ def load_badger_docs(name: str, ptype: str | None = None) -> str: if ptype == "generator": docstring = generators[name].__doc__ - help_md = _format_docs_str(readme, docstring, ptype) + help_md = _format_docs_str(readme, docstring or "", ptype or "") return f"{header}
{help_md}" except FileNotFoundError: @@ -322,7 +341,7 @@ def _format_docs_str(readme: str, docstring: str, ptype: str) -> str: return help_md -def _format_md_docs(text: str): +def _format_md_docs(text: str) -> str: """ Helper function to format markdown docs for display in QTextBrowser. Removes the first '---' section and replaces double newlines with
for better rendering. @@ -355,7 +374,7 @@ def _format_md_docs(text: str): def _md_images_to_html( text: str, - base_prefix: str | None = None, + base_prefix: str | Path | None = None, width: int = 575, ) -> str: """ @@ -367,15 +386,17 @@ def _md_images_to_html( # Get absolute path to image folder relative to this module base_prefix = Path(__file__).parent.parent.parent / "documentation" / "static" - def repl(m: re.Match) -> str: - url = m.group(1).strip().strip("'\"").lstrip("./") - url = Path(base_prefix / url) + def repl(m: re.Match[str]) -> str: + url_str = m.group(1).strip().strip("'\"").lstrip("./") + url = Path(base_prefix / url_str) # type: ignore[operator] return f'' return _MD_IMG.sub(repl, text) -def get_plug(root: str, name: str, ptype: str): +def get_plug( + root: str, name: str, ptype: Literal["environment", "interface", "generator"] +) -> "tuple[type[BadgerEnvironment | BadgerInterface | Generator] | None, BadgerPluginConfig | None]": try: plug = BADGER_FACTORY[ptype][name] if plug is None: # lazy loading @@ -393,21 +414,21 @@ def get_plug(root: str, name: str, ptype: str): return plug -def scan_extensions(root): - extensions = {} +def scan_extensions(root: str) -> dict[str, Any]: + extensions: dict[str, Any] = {} return extensions -def get_env_docs(name: str): +def get_env_docs(name: str) -> str: return load_plugin_docs(name, "environment") -def get_intf(name: str): +def get_intf(name: str) -> tuple[Any, Any]: return get_plug(BADGER_PLUGIN_ROOT, name, "interface") -def get_env(name: str): +def get_env(name: str) -> tuple[Any, Any]: return get_plug(BADGER_PLUGIN_ROOT, name, "environment") @@ -415,7 +436,7 @@ def list_generators() -> list[str]: try: from xopt.generators import try_load_all_generators - try_load_all_generators() + try_load_all_generators() # type: ignore[no-untyped-call] except ImportError: # this API changed somehow pass # there is nothing we can do... generator_names = list(generators.keys()) @@ -427,11 +448,11 @@ def list_generators() -> list[str]: get_generator = get_generator_defaults -def list_intf(): +def list_intf() -> list[str]: return sorted(BADGER_FACTORY["interface"]) -def list_env(): +def list_env() -> list[str]: return sorted(BADGER_FACTORY["environment"]) diff --git a/src/badger/formula.py b/src/badger/formula.py index 51208bbd..d8ee63de 100644 --- a/src/badger/formula.py +++ b/src/badger/formula.py @@ -10,19 +10,21 @@ import ast import difflib import re +from collections.abc import Iterable +from typing import Any import numpy as np -def safe_var_name(var_name): +def safe_var_name(var_name: str) -> str: return re.sub(r"[^0-9a-zA-Z_]", "_", var_name) -def extract_variable_keys(expr): +def extract_variable_keys(expr: str) -> list[str]: return re.findall(r"`([^`]+)`", expr) -def find_used_names(expr): +def find_used_names(expr: str) -> set[str]: try: tree = ast.parse(expr, mode="eval") return {node.id for node in ast.walk(tree) if isinstance(node, ast.Name)} @@ -30,16 +32,16 @@ def find_used_names(expr): raise SyntaxError(f"Invalid syntax in expression: {e}") -def suggest_name(unknown, known_names): - suggestions = {} +def suggest_name(unknown: Iterable[str], known_names: Iterable[str]) -> dict[str, str]: + suggestions: dict[str, str] = {} for name in unknown: - matches = difflib.get_close_matches(name, known_names, n=1, cutoff=0.7) + matches = difflib.get_close_matches(name, list(known_names), n=1, cutoff=0.7) if matches: suggestions[name] = matches[0] return suggestions -def interpret_expression(expr, variables): +def interpret_expression(expr: str, variables: dict[str, Any]) -> Any: """ Interpret a mathematical expression with variables and functions. The expression can contain variables in single or double quotes, @@ -59,7 +61,7 @@ def interpret_expression(expr, variables): """ quoted_vars = extract_variable_keys(expr) - missing_vars = quoted_vars - variables.keys() + missing_vars = set(quoted_vars) - variables.keys() if missing_vars: raise KeyError(f"Missing variables for expression: {sorted(missing_vars)}") @@ -88,7 +90,7 @@ def interpret_expression(expr, variables): msg += f" - {bad} → {good}\n" raise NameError(msg.strip()) - safe_namespace = {name: getattr(np, name) for name in np_funcs} + safe_namespace: dict[str, Any] = {name: getattr(np, name) for name in np_funcs} safe_namespace["percentile"] = np.percentile # Add common built-in functions for func_name in builtin_funcs: diff --git a/src/badger/gui/__init__.py b/src/badger/gui/__init__.py index 7803fc8e..c5f25a91 100644 --- a/src/badger/gui/__init__.py +++ b/src/badger/gui/__init__.py @@ -8,7 +8,7 @@ import traceback from importlib import resources from types import TracebackType -from typing import NoReturn +from typing import Any, NoReturn from PyQt5 import QtCore from PyQt5.QtGui import QFont, QIcon @@ -35,10 +35,10 @@ # QApplication.setAttribute( # QtCore.Qt.HighDpiScaleFactorRoundingPolicy.PassThrough) -TIMER = {"time": None} +TIMER: dict[str, float | None] = {"time": None} -def on_exit(*args): +def on_exit(*args: Any) -> None: if TIMER["time"] is None: print("Press Ctrl/Cmd + C again within 1s to quit Badger") TIMER["time"] = time.time() @@ -49,7 +49,7 @@ def on_exit(*args): QApplication.quit() -def on_timeout(): +def on_timeout() -> None: if TIMER["time"] is None: return @@ -60,7 +60,7 @@ def on_timeout(): def error_handler( - etype: type[BaseException], value: BaseException, tb: TracebackType + etype: type[BaseException], value: BaseException, tb: TracebackType | None ) -> NoReturn: """ Custom exception handler that formats uncaught exceptions and raises a BadgerError. @@ -71,7 +71,7 @@ def error_handler( The class of the exception that was raised. value : BaseException The exception instance. - tb : TracebackType + tb : TracebackType | None The traceback object associated with the exception. Raises @@ -84,7 +84,7 @@ def error_handler( raise BadgerError(error_title, error_msg) -def launch_gui(config_path=None): +def launch_gui(config_path: str | None = None) -> None: sys.excepthook = error_handler app = QApplication(sys.argv) diff --git a/src/badger/gui/components/action_bar.py b/src/badger/gui/components/action_bar.py index 7bb58b9f..2fd50353 100644 --- a/src/badger/gui/components/action_bar.py +++ b/src/badger/gui/components/action_bar.py @@ -2,9 +2,10 @@ docs access, and the extensions palette launcher.""" from importlib import resources +from typing import cast -from PyQt5.QtCore import QEvent, QSize, pyqtSignal -from PyQt5.QtGui import QFont, QIcon +from PyQt5.QtCore import QEvent, QPoint, QSize, pyqtSignal +from PyQt5.QtGui import QFont, QHelpEvent, QIcon from PyQt5.QtWidgets import ( QAction, QHBoxLayout, @@ -17,6 +18,7 @@ from badger.gui.utils import create_button from badger.gui.windows.docs_window import BadgerDocsWindow +from badger.types import TerminationConditionConfig class SplitTooltipToolButton(QToolButton): @@ -25,7 +27,7 @@ class SplitTooltipToolButton(QToolButton): Use arg menu_tooltip="desired tooltip" to set the menu tooltip """ - def __init__(self, menu_tooltip="", parent=None): + def __init__(self, menu_tooltip: str = "", parent: QWidget | None = None) -> None: """ Parameters ---------- @@ -35,20 +37,28 @@ def __init__(self, menu_tooltip="", parent=None): super().__init__(parent) self.menu_tooltip = menu_tooltip - def _over_menu_arrow(self, pos): + def _over_menu_arrow(self, pos: QPoint) -> bool: opt = QStyleOptionToolButton() self.initStyleOption(opt) - rect = self.style().subControlRect( - QStyle.CC_ToolButton, opt, QStyle.SC_ToolButtonMenu, self + style = self.style() + if style is None: + raise RuntimeError("Style not found") + rect = style.subControlRect( + QStyle.ComplexControl.CC_ToolButton, + opt, + QStyle.SubControl.SC_ToolButtonMenu, + self, ) return rect.contains(pos) - def event(self, event): - if event.type() == QEvent.ToolTip and self._over_menu_arrow(event.pos()): - from PyQt5.QtWidgets import QToolTip + def event(self, event: QEvent | None) -> bool: + if event is not None and event.type() == QEvent.Type.ToolTip: + help_event = cast(QHelpEvent, event) + if self._over_menu_arrow(help_event.pos()): + from PyQt5.QtWidgets import QToolTip - QToolTip.showText(event.globalPos(), self.menu_tooltip, self) - return True + QToolTip.showText(help_event.globalPos(), self.menu_tooltip, self) + return True return super().event(event) @@ -150,13 +160,13 @@ class BadgerActionBar(QWidget): sig_edit_checkpoint = pyqtSignal() sig_load_checkpoint = pyqtSignal() - def __init__(self, parent=None): + def __init__(self, parent: QWidget | None = None) -> None: super().__init__(parent) self.docs_name = "gui-usage" self.init_ui() self.config_logic() - def init_ui(self): + def init_ui(self) -> None: def load_internal_icon(name: str) -> QIcon: icon_ref = resources.files(__package__) / f"../images/{name}" with resources.as_file(icon_ref) as icon_path: @@ -189,13 +199,14 @@ def load_internal_icon(name: str) -> QIcon: self.btn_help.setIconSize(QSize(20, 20)) self.btn_reset = create_button("undo.png", "Reset environment") - self.btn_checkpoint = create_button( - "flag.png", "Checkpoint", size=(48, 32), tool_button=True + self.btn_checkpoint = cast( + QToolButton, + create_button("flag.png", "Checkpoint", size=(48, 32), tool_button=True), ) self.btn_opt = create_button("star.png", "Jump to optimum") self.btn_set = create_button("set.png", "Dial in solution") self.btn_ctrl = create_button("pause.png", "Pause") - self.btn_ctrl._status = "pause" + self.ctrl_status = "pause" self.btn_del.setDisabled(True) self.btn_log.setDisabled(True) @@ -291,7 +302,7 @@ def load_internal_icon(name: str) -> QIcon: } """) - def config_logic(self): + def config_logic(self) -> None: self.btn_del.clicked.connect(self.delete_run) self.btn_log.clicked.connect(self.logbook) self.btn_help.clicked.connect(self.open_docs) @@ -318,7 +329,7 @@ def config_logic(self): ) self.btn_open_extensions_palette.clicked.connect(self.open_extensions_palette) - def lock(self): + def lock(self) -> None: self.btn_del.setDisabled(True) self.btn_log.setDisabled(True) self.btn_reset.setDisabled(True) @@ -328,7 +339,7 @@ def lock(self): self.btn_opt.setDisabled(True) self.btn_set.setDisabled(True) - def unlock(self): + def unlock(self) -> None: self.btn_del.setDisabled(False) self.btn_log.setDisabled(False) self.btn_reset.setDisabled(False) @@ -338,13 +349,13 @@ def unlock(self): self.btn_opt.setDisabled(False) self.btn_set.setDisabled(False) - def routine_invalid(self): + def routine_invalid(self) -> None: self.btn_stop.setDisabled(False) - def routine_finished(self): + def routine_finished(self) -> None: self.btn_ctrl.setIcon(self.icon_pause) self.btn_ctrl.setToolTip("Pause") - self.btn_ctrl._status = "pause" + self.ctrl_status = "pause" self.btn_ctrl.setDisabled(True) # Note the order of the following two lines cannot be changed! @@ -363,19 +374,19 @@ def routine_finished(self): self.btn_set.setDisabled(False) self.btn_del.setDisabled(False) - def toggle_reset(self, locked): + def toggle_reset(self, locked: bool) -> None: self.btn_reset.setDisabled(locked) - def toggle_run(self, locked): + def toggle_run(self, locked: bool) -> None: self.btn_stop.setDisabled(locked) - def toggle_other(self, locked): + def toggle_other(self, locked: bool) -> None: self.btn_del.setDisabled(locked) self.btn_log.setDisabled(locked) self.btn_opt.setDisabled(locked) self.btn_set.setDisabled(locked) - def run_start(self): + def run_start(self) -> None: self.btn_stop.setStyleSheet(stylesheet_stop) self.btn_stop.setPopupMode(QToolButton.DelayedPopup) self.btn_stop.setDisabled(False) @@ -389,7 +400,7 @@ def run_start(self): self.btn_ctrl.setDisabled(False) self.btn_set.setDisabled(True) - def set_run_action(self): + def set_run_action(self) -> None: if self.btn_stop.defaultAction() is not self.run_action: self.btn_stop.setDefaultAction(self.run_action) @@ -400,7 +411,7 @@ def set_run_action(self): self.btn_stop.setDisabled(True) self.sig_stop.emit() - def set_run_until_action(self, from_menu=False): + def set_run_until_action(self, from_menu: bool = False) -> None: if self.btn_stop.defaultAction() is not self.run_until_action: self.btn_stop.setDefaultAction(self.run_until_action) @@ -410,62 +421,61 @@ def set_run_until_action(self, from_menu=False): self.btn_stop.setDisabled(True) self.sig_stop.emit() - def _on_run_action_triggered(self): + def _on_run_action_triggered(self) -> None: self.set_run_action() - def _on_run_until_action_triggered(self): + def _on_run_until_action_triggered(self) -> None: self.set_run_until_action(from_menu=False) - def _on_run_until_menu_action_triggered(self): + def _on_run_until_menu_action_triggered(self) -> None: self.set_run_until_action(from_menu=True) - def delete_run(self): + def delete_run(self) -> None: self.sig_delete_run.emit() - def logbook(self): + def logbook(self) -> None: self.sig_logbook.emit() - def open_docs(self): + def open_docs(self) -> None: self.window_docs.update_docs(self.docs_name) self.window_docs.show() - def reset_env(self): + def reset_env(self) -> None: self.sig_reset_env.emit() - def jump_to_optimal(self): + def jump_to_optimal(self) -> None: self.sig_jump_to_optimal.emit() - def dial_in(self): + def dial_in(self) -> None: self.sig_dial_in.emit() - def ctrl_routine(self): - if self.btn_ctrl._status == "pause": + def ctrl_routine(self) -> None: + if self.ctrl_status == "pause": self.sig_ctrl.emit(True) self.btn_ctrl.setIcon(self.icon_play) self.btn_ctrl.setToolTip("Resume") - self.btn_ctrl._status = "play" + self.ctrl_status = "play" else: self.sig_ctrl.emit(False) self.btn_ctrl.setIcon(self.icon_pause) self.btn_ctrl.setToolTip("Pause") - self.btn_ctrl._status = "pause" + self.ctrl_status = "pause" - def open_extensions_palette(self): + def open_extensions_palette(self) -> None: self.sig_open_extensions_palette.emit() - def env_ready(self): + def env_ready(self) -> None: self.btn_log.setDisabled(False) self.btn_opt.setDisabled(False) - def update_run_tooltip(self, tc=None): - """Update btn_stop tooltip: tc dict for run-until mode, or None.""" + def update_run_tooltip(self, tc: TerminationConditionConfig | None = None) -> None: + """Update btn_stop tooltip: tc config for run-until mode, or None.""" if tc is None: self.run_action.setToolTip("Run") else: - tc_idx = tc.get("tc_idx", 0) - if tc_idx == 0: - tip = f"Run until: n iterations = {tc.get('max_eval')}" - elif tc_idx == 1: - tip = f"Run until: timeout = {tc.get('max_time')}s" + if tc.tc_idx == 0: + tip = f"Run until: n iterations = {tc.max_eval}" + elif tc.tc_idx == 1: + tip = f"Run until: timeout = {tc.max_time}s" self.run_until_action.setToolTip(tip) self.run_until_menu_action.setToolTip(tip) diff --git a/src/badger/gui/components/analysis_extensions.py b/src/badger/gui/components/analysis_extensions.py index 8246a19e..c9b034c4 100644 --- a/src/badger/gui/components/analysis_extensions.py +++ b/src/badger/gui/components/analysis_extensions.py @@ -2,7 +2,7 @@ viewer). Each dialog receives live data updates from the run monitor.""" import logging -from typing import cast +from typing import Any, cast from PyQt5.QtCore import Qt, pyqtSignal from PyQt5.QtGui import QCloseEvent @@ -24,7 +24,7 @@ class AnalysisExtension(QWidget): window_closed = pyqtSignal(object) generator_type: type[Generator] - widget: AnalysisWidget + widget: AnalysisWidget[Any] def __init__(self, parent: QWidget | None = None): super().__init__(parent=parent) @@ -43,7 +43,7 @@ def update_window(self, routine: Routine) -> None: def initialize_extension( self, - extension_widget: AnalysisWidget, + extension_widget: AnalysisWidget[Any], extension_name: str, generator_type: type[Generator], ) -> None: @@ -77,7 +77,7 @@ def update_extension( self.widget.update_plots(requires_rebuild, interval=self.widget.update_interval) - def closeEvent(self, a0: QCloseEvent) -> None: + def closeEvent(self, a0: QCloseEvent | None) -> None: self.window_closed.emit(self) super().closeEvent(a0) diff --git a/src/badger/gui/components/analysis_widget.py b/src/badger/gui/components/analysis_widget.py index d300e7d2..f7d5e947 100644 --- a/src/badger/gui/components/analysis_widget.py +++ b/src/badger/gui/components/analysis_widget.py @@ -4,7 +4,7 @@ import logging from abc import abstractmethod from collections.abc import Callable -from typing import Any +from typing import Generic, TypeVar from PyQt5.QtWidgets import QWidget from xopt import Generator @@ -15,10 +15,14 @@ logger = logging.getLogger(__name__) +# Each extension supplies its own options container (a TypedDict or dataclass). +ParametersT = TypeVar("ParametersT") -class AnalysisWidget(QWidget): + +class AnalysisWidget(QWidget, Generic[ParametersT]): routine: Routine generator: Generator + parameters: ParametersT df_length: float = float("inf") initialized: bool = False @@ -32,7 +36,6 @@ def __init__( routine: Routine, parent: QWidget | None = None, ): - self.parameters: dict[str, Any] = {} super().__init__(parent=parent) self.routine = routine self.generator = routine.generator diff --git a/src/badger/gui/components/archive_search.py b/src/badger/gui/components/archive_search.py index 1449e0a3..4c4a8bd3 100644 --- a/src/badger/gui/components/archive_search.py +++ b/src/badger/gui/components/archive_search.py @@ -3,6 +3,8 @@ when building a routine.""" import logging +from collections.abc import Callable +from typing import Any from PyQt5.QtCore import ( QAbstractTableModel, @@ -10,7 +12,6 @@ QModelIndex, QObject, Qt, - QVariant, pyqtSignal, ) from PyQt5.QtGui import QDrag, QKeyEvent @@ -26,6 +27,7 @@ QWidget, ) +from badger.environment import Environment from badger.errors import BadgerRoutineError logger = logging.getLogger(__name__) @@ -41,37 +43,42 @@ class ArchiveResultsTableModel(QAbstractTableModel): The parent item of this table """ - def __init__(self, parent: QObject = None) -> None: + def __init__(self, parent: QObject | None = None) -> None: super().__init__(parent=parent) - self.results_list = [] + self.results_list: list[str] = [] self.column_names = ("Variable",) - def rowCount(self, parent: QObject) -> int: + def rowCount(self, parent: QModelIndex | None = None) -> int: """Return the row count of the table""" if parent is not None and parent.isValid(): return 0 return len(self.results_list) - def columnCount(self, parent: QObject) -> int: + def columnCount(self, parent: QModelIndex | None = None) -> int: """Return the column count of the table""" if parent is not None and parent.isValid(): return 0 return len(self.column_names) - def data(self, index: QModelIndex, role: int) -> QVariant: + def data(self, index: QModelIndex, role: int = Qt.ItemDataRole.DisplayRole) -> Any: """Return the data for the associated role. Currently only supporting DisplayRole.""" if not index.isValid(): - return QVariant() + return None - if role != Qt.DisplayRole: - return QVariant() + if role != Qt.ItemDataRole.DisplayRole: + return None return self.results_list[index.row()] - def headerData(self, section, orientation, role=Qt.DisplayRole) -> QVariant: + def headerData( + self, + section: int, + orientation: Qt.Orientation, + role: int = Qt.ItemDataRole.DisplayRole, + ) -> Any: """Return data associated with the header""" - if role != Qt.DisplayRole: + if role != Qt.ItemDataRole.DisplayRole: return super().headerData(section, orientation, role) return str(self.column_names[section]) @@ -79,7 +86,12 @@ def headerData(self, section, orientation, role=Qt.DisplayRole) -> QVariant: def flags(self, index: QModelIndex) -> Qt.ItemFlags: """Return flags that determine how users can interact with the items in the table""" if index.isValid(): - return Qt.ItemIsEnabled | Qt.ItemIsSelectable | Qt.ItemIsDragEnabled + return ( + Qt.ItemFlags(Qt.ItemFlag.ItemIsEnabled) + | Qt.ItemFlag.ItemIsSelectable + | Qt.ItemFlag.ItemIsDragEnabled + ) + return Qt.ItemFlags() def append(self, pv: str) -> None: """Appends a row to this table given the variable name as input""" @@ -104,12 +116,33 @@ def clear(self) -> None: self.endRemoveRows() self.layoutChanged.emit() - def sort(self, col: int, order=Qt.AscendingOrder) -> None: + def sort(self, col: int, order: Qt.SortOrder = Qt.SortOrder.AscendingOrder) -> None: """Sort the table by variable name""" - self.results_list.sort(reverse=order == Qt.DescendingOrder) + self.results_list.sort(reverse=order == Qt.SortOrder.DescendingOrder) self.layoutChanged.emit() +class ResultsTableView(QTableView): + """Table view whose drag payload is the caller-supplied selected-variable text.""" + + def __init__( + self, + get_selected: Callable[[], str], + parent: QWidget | None = None, + ) -> None: + super().__init__(parent) + self._get_selected = get_selected + + def startDrag(self, supportedActions: Qt.DropActions | Qt.DropAction) -> None: + """Start a drag carrying the selected variable names so they can be dropped + onto a plot to begin drawing that variable's data.""" + drag = QDrag(self) + mime_data = QMimeData() + mime_data.setText(self._get_selected()) + drag.setMimeData(mime_data) + drag.exec_() + + class ArchiveSearchWidget(QWidget): """ The ArchiveSearchWidget is a display widget for showing the results of a variable search using, for example, an instance of the @@ -118,19 +151,19 @@ class ArchiveSearchWidget(QWidget): Parameters ---------- - parent : QObject, optional - The parent item of this widget + parent : QWidget, optional + The parent widget of this widget """ append_variables_requested = pyqtSignal(str) - def __init__(self, environment, parent: QObject = None) -> None: + def __init__(self, environment: Environment, parent: QWidget | None = None) -> None: super().__init__(parent=parent) self.env = environment self.resize(400, 800) self.setWindowTitle("Variable Search") - self.layout = QVBoxLayout() + self._layout = QVBoxLayout() self.search_label = QLabel("Pattern") self.search_box = QLineEdit() @@ -143,7 +176,7 @@ def __init__(self, environment, parent: QObject = None) -> None: self.loading_label.hide() self.results_table_model = ArchiveResultsTableModel() - self.results_view = QTableView(self) + self.results_view = ResultsTableView(self.selectedVariables, self) self.results_view.setModel(self.results_table_model) self.results_view.setProperty("showDropIndicator", False) self.results_view.setDragDropOverwriteMode(False) @@ -153,23 +186,26 @@ def __init__(self, environment, parent: QObject = None) -> None: self.results_view.setDropIndicatorShown(True) self.results_view.setCornerButtonEnabled(False) self.results_view.setSortingEnabled(True) - self.results_view.verticalHeader().setVisible(False) - self.results_view.horizontalHeader().setSectionResizeMode(QHeaderView.Stretch) - self.results_view.startDrag = self.startDragAction + vheader = self.results_view.verticalHeader() + if vheader is not None: + vheader.setVisible(False) + hheader = self.results_view.horizontalHeader() + if hheader is not None: + hheader.setSectionResizeMode(QHeaderView.Stretch) self.archive_url_layout = QHBoxLayout() - self.layout.addLayout(self.archive_url_layout) + self._layout.addLayout(self.archive_url_layout) self.search_layout = QHBoxLayout() self.search_layout.addWidget(self.search_label) self.search_layout.addWidget(self.search_box) self.search_layout.addWidget(self.search_button) - self.layout.addLayout(self.search_layout) - self.layout.addWidget(self.loading_label) - self.layout.addWidget(self.results_view) + self._layout.addLayout(self.search_layout) + self._layout.addWidget(self.loading_label) + self._layout.addWidget(self.results_view) self.results_view.doubleClicked.connect( lambda: self.append_variables_requested.emit(self.selectedVariables()) ) - self.setLayout(self.layout) + self.setLayout(self._layout) def selectedVariables(self) -> str: """Figure out based on which indexes were selected, the list of variables (by string name) @@ -182,21 +218,11 @@ def selectedVariables(self) -> str: pv_list += pv_name + ", " return pv_list[:-2] - def startDragAction(self, supported_actions) -> None: - """ - The method to be called when a user initiates a drag action for one of the results in the table. The current - reason for this functionality is the ability to drag a variable name onto a plot to automatically start drawing - data for that variable - """ - drag = QDrag(self) - mime_data = QMimeData() - mime_data.setText(self.selectedVariables()) - drag.setMimeData(mime_data) - drag.exec_() - - def keyPressEvent(self, e: QKeyEvent) -> None: + def keyPressEvent(self, e: QKeyEvent | None) -> None: """Special key press tracker, just so that if enter or return is pressed the formula dialog attempts to submit the formula""" - if e.key() == Qt.Key_Return or e.key() == Qt.Key_Enter: + if e is not None and ( + e.key() == Qt.Key.Key_Return or e.key() == Qt.Key.Key_Enter + ): self.request_variable_search() return super().keyPressEvent(e) diff --git a/src/badger/gui/components/bax_visualizer/bax_widget.py b/src/badger/gui/components/bax_visualizer/bax_widget.py index a1325eaf..6406600f 100644 --- a/src/badger/gui/components/bax_visualizer/bax_widget.py +++ b/src/badger/gui/components/bax_visualizer/bax_widget.py @@ -69,9 +69,9 @@ class Parameters: include_y: bool = True -class BaxWidget(AnalysisWidget): +class BaxWidget(AnalysisWidget[Parameters]): generator: BaxGenerator - parameters: Parameters # type: ignore[assignment] + parameters: Parameters def __init__(self, routine: Routine, parent: QWidget | None = None): logger.debug("Initializing BaxWidget") diff --git a/src/badger/gui/components/bax_visualizer/controls.py b/src/badger/gui/components/bax_visualizer/controls.py index fee4bee5..9511a889 100644 --- a/src/badger/gui/components/bax_visualizer/controls.py +++ b/src/badger/gui/components/bax_visualizer/controls.py @@ -169,7 +169,8 @@ def _create_reference_point_group(self) -> QGroupBox: self.reference_table.setColumnCount(2) self.reference_table.setHorizontalHeaderLabels(["Variable", "Value"]) horizontal_header = self.reference_table.horizontalHeader() - horizontal_header.setSectionResizeMode(QHeaderView.ResizeMode.Stretch) + if horizontal_header is not None: + horizontal_header.setSectionResizeMode(QHeaderView.ResizeMode.Stretch) self.select_latest_reference_point_button = QPushButton("Set Latest") self.reference_point_display = QLabel("") @@ -190,7 +191,7 @@ def populate_reference_table( with BlockSignalsContext(self.reference_table): self.reference_table.setRowCount(len(self.parameters.variables)) - self.ref_inputs: list[QTableWidgetItem] = [] + self.ref_inputs = [] for i, var_name in enumerate(self.parameters.variables): variable_item = QTableWidgetItem(var_name) @@ -248,7 +249,8 @@ def update_reference_point_table_editability(self) -> None: # Force the table to refresh and update its view viewport = self.reference_table.viewport() - viewport.update() + if viewport is not None: + viewport.update() def get_selected_variables(self) -> list[str]: """Get the currently selected variables from the combo boxes.""" diff --git a/src/badger/gui/components/bo_visualizer/bo_widget.py b/src/badger/gui/components/bo_visualizer/bo_widget.py index bfba76d6..a5aacb4f 100644 --- a/src/badger/gui/components/bo_visualizer/bo_widget.py +++ b/src/badger/gui/components/bo_visualizer/bo_widget.py @@ -57,7 +57,7 @@ } -class BOPlotWidget(AnalysisWidget): +class BOPlotWidget(AnalysisWidget[ConfigurableOptions]): generator: BayesianGenerator # pyright: ignore[reportIncompatibleVariableOverride] df_length: float = float("inf") @@ -374,7 +374,8 @@ def update_reference_point_table(self, selected_variables: list[str]) -> None: # Force the table to refresh and update its view viewport = self.ui_components.reference_table.viewport() - viewport.update() + if viewport is not None: + viewport.update() def get_reference_points( self, ref_inputs: list[QTableWidgetItem], variable_names: list[str] diff --git a/src/badger/gui/components/bo_visualizer/plotting_area.py b/src/badger/gui/components/bo_visualizer/plotting_area.py index b593669d..8b0ec17c 100644 --- a/src/badger/gui/components/bo_visualizer/plotting_area.py +++ b/src/badger/gui/components/bo_visualizer/plotting_area.py @@ -87,6 +87,9 @@ def update_plot( ) layout = self.layout() + if layout is None: + logger.error("Layout not found") + return with BlockSignalsContext(layout): # Clear the existing layout (remove previous plot if any) @@ -94,14 +97,14 @@ def update_plot( with MatplotlibFigureContext(fig, ax) as (fig, ax): # Create a new figure and canvas - canvas = FigureCanvas(fig) - toolbar = NavigationToolbar(canvas, self) + canvas = FigureCanvas(fig) # type: ignore[no-untyped-call] + toolbar = NavigationToolbar(canvas, self) # type: ignore[no-untyped-call] variables = parameters["variables"] handler = MatplotlibInteractionHandler( canvas, - parameters, # pyright: ignore[reportArgumentType] + parameters, routine, variables, update_extension, diff --git a/src/badger/gui/components/bo_visualizer/types.py b/src/badger/gui/components/bo_visualizer/types.py index eaf18732..41b171dd 100644 --- a/src/badger/gui/components/bo_visualizer/types.py +++ b/src/badger/gui/components/bo_visualizer/types.py @@ -3,6 +3,8 @@ from typing import TypedDict +from badger.gui.components.types import InteractionParameters + class PlotOptions(TypedDict): n_grid: int @@ -13,10 +15,6 @@ class PlotOptions(TypedDict): show_acq_func: bool -class ConfigurableOptions(TypedDict): +class ConfigurableOptions(InteractionParameters): plot_options: PlotOptions - variable_1: int - variable_2: int - variables: list[str] - reference_points: dict[str, float] include_variable_2: bool diff --git a/src/badger/gui/components/bounds_preview.py b/src/badger/gui/components/bounds_preview.py index a8610f3a..83107afb 100644 --- a/src/badger/gui/components/bounds_preview.py +++ b/src/badger/gui/components/bounds_preview.py @@ -48,7 +48,7 @@ def _to_x(self, value: float, left: float, width: float) -> float: ratio = max(0.0, min(1.0, ratio)) return left + ratio * width - def paintEvent(self, _event: QPaintEvent) -> None: + def paintEvent(self, a0: QPaintEvent | None) -> None: """Paint the track, preview range, current marker, and value labels.""" painter = QPainter(self) painter.setRenderHint(QPainter.Antialiasing) @@ -64,7 +64,7 @@ def paintEvent(self, _event: QPaintEvent) -> None: track_top = track_y - track_h / 2.0 track_rect = QRectF(left, track_top, track_w, float(track_h)) - painter.setPen(Qt.NoPen) + painter.setPen(Qt.PenStyle.NoPen) painter.setBrush(QBrush(QColor(55, 66, 81))) painter.drawRoundedRect(track_rect, 2, 2) diff --git a/src/badger/gui/components/collapsible_box.py b/src/badger/gui/components/collapsible_box.py index 0f67e88e..e45c8d0f 100644 --- a/src/badger/gui/components/collapsible_box.py +++ b/src/badger/gui/components/collapsible_box.py @@ -2,8 +2,21 @@ group environment config, generator parameters, etc. without taking up permanent screen space.""" -from PyQt5 import QtCore, QtGui, QtWidgets -from PyQt5.QtWidgets import QLayout +import logging +from typing import cast + +from PyQt5 import QtCore, QtWidgets +from PyQt5.QtCore import pyqtSignal +from PyQt5.QtGui import QResizeEvent +from PyQt5.QtWidgets import ( + QLayout, + QScrollArea, + QSizePolicy, + QToolButton, + QWidget, +) + +logger = logging.getLogger(__name__) stylesheet_toolbutton = """ QToolButton @@ -14,48 +27,57 @@ # https://stackoverflow.com/a/56293688 -class ScrollArea(QtWidgets.QScrollArea): - resized = QtCore.pyqtSignal() +class ScrollArea(QScrollArea): + resized = pyqtSignal() - def resizeEvent(self, e): + def resizeEvent(self, a0: QResizeEvent | None) -> None: self.resized.emit() - return super().resizeEvent(e) + return super().resizeEvent(a0) # https://stackoverflow.com/a/52617714/4263605 -class CollapsibleBox(QtWidgets.QWidget): - def __init__(self, parent=None, title="", duration=100, tooltip=None): +class CollapsibleBox(QWidget): + def __init__( + self, + parent: QWidget | None = None, + title: str = "", + duration: int = 100, + tooltip: str = "", + ): super().__init__(parent) self.title = title self.duration = duration - cool_font = QtGui.QFont() - cool_font.setWeight(QtGui.QFont.DemiBold) + from PyQt5.QtGui import QFont + + cool_font = QFont() + cool_font.setWeight(QFont.DemiBold) cool_font.setPixelSize(13) - self.toggle_button = QtWidgets.QToolButton( - text=title, checkable=True, checked=False - ) + self.toggle_button = QToolButton() + self.toggle_button.setText(title) + self.toggle_button.setCheckable(True) + self.toggle_button.setChecked(False) self.toggle_button.setFont(cool_font) self.toggle_button.setFixedHeight(28) - self.toggle_button.setSizePolicy( - QtWidgets.QSizePolicy.Expanding, QtWidgets.QSizePolicy.Fixed - ) + self.toggle_button.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Fixed) self.toggle_button.setStyleSheet(stylesheet_toolbutton) - self.toggle_button.setToolButtonStyle(QtCore.Qt.ToolButtonTextBesideIcon) + self.toggle_button.setToolButtonStyle( + QtCore.Qt.ToolButtonStyle.ToolButtonTextBesideIcon + ) self.toggle_button.setIconSize(QtCore.QSize(11, 11)) - self.toggle_button.setArrowType(QtCore.Qt.RightArrow) + self.toggle_button.setArrowType(QtCore.Qt.ArrowType.RightArrow) self.toggle_button.setToolTip(tooltip) self.toggle_button.clicked.connect(self.start_animation) # self.toggle_button.setText(f'+ {title}') self.toggle_animation = QtCore.QParallelAnimationGroup(self) - self.content_area = ScrollArea(maximumHeight=0, minimumHeight=0) - self.content_area.setSizePolicy( - QtWidgets.QSizePolicy.Expanding, QtWidgets.QSizePolicy.Fixed - ) + self.content_area = ScrollArea() + self.content_area.setMaximumHeight(0) + self.content_area.setMinimumHeight(0) + self.content_area.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Fixed) self.content_area.setFrameShape(QtWidgets.QFrame.NoFrame) self.content_area.resized.connect(self.updateContentLayout) @@ -76,29 +98,35 @@ def __init__(self, parent=None, title="", duration=100, tooltip=None): ) @QtCore.pyqtSlot() - def start_animation(self): + def start_animation(self) -> None: checked = self.toggle_button.isChecked() - arrow_type = QtCore.Qt.DownArrow if checked else QtCore.Qt.RightArrow + arrow_type = ( + QtCore.Qt.ArrowType.DownArrow if checked else QtCore.Qt.ArrowType.RightArrow + ) self.toggle_button.setArrowType(arrow_type) # self.toggle_button.setText(f'- {self.title}' if not checked else f'+ {self.title}') direction = ( - QtCore.QAbstractAnimation.Forward + QtCore.QAbstractAnimation.Direction.Forward if checked - else QtCore.QAbstractAnimation.Backward + else QtCore.QAbstractAnimation.Direction.Backward ) self.toggle_animation.setDirection(direction) self.toggle_animation.start() - def expand(self): + def expand(self) -> None: if not self.toggle_button.isChecked(): self.toggle_button.click() - def updateContentLayout(self): + def updateContentLayout(self) -> None: if ( self.toggle_button.isChecked() - and self.toggle_animation.state() != self.toggle_animation.Running + and self.toggle_animation.state() != self.toggle_animation.State.Running ): - content_height = self.content_area.layout().sizeHint().height() + _layout = self.content_area.layout() + if _layout is None: + logger.warning("Content area has no layout.") + return + content_height = _layout.sizeHint().height() self.setMinimumHeight(self.collapsed_height + content_height) self.setMaximumHeight(self.collapsed_height + content_height) self.content_area.setMaximumHeight(content_height) @@ -107,7 +135,7 @@ def updateContentLayout(self): if isinstance(p, ScrollArea): p.resized.emit() - def setContentLayout(self, layout: QLayout): + def setContentLayout(self, layout: QLayout) -> None: lay = self.content_area.layout() del lay self.content_area.setLayout(layout) @@ -116,13 +144,18 @@ def setContentLayout(self, layout: QLayout): ) content_height = layout.sizeHint().height() for i in range(self.toggle_animation.animationCount() - 1): - animation = self.toggle_animation.animationAt(i) + animation = cast( + QtCore.QPropertyAnimation, self.toggle_animation.animationAt(i) + ) animation.setDuration(self.duration) animation.setStartValue(collapsed_height) animation.setEndValue(collapsed_height + content_height) - content_animation = self.toggle_animation.animationAt( - self.toggle_animation.animationCount() - 1 + content_animation = cast( + QtCore.QPropertyAnimation, + self.toggle_animation.animationAt( + self.toggle_animation.animationCount() - 1 + ), ) content_animation.setDuration(self.duration) content_animation.setStartValue(0) diff --git a/src/badger/gui/components/con_table.py b/src/badger/gui/components/con_table.py index bd44b928..bae0019c 100644 --- a/src/badger/gui/components/con_table.py +++ b/src/badger/gui/components/con_table.py @@ -48,16 +48,16 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: ["", "Name", "Relation", "Threshold", "Critical"] ) - def default_info(self) -> list[Any]: + def default_info(self) -> tuple[str, float, bool]: """ Get the default information list for a new item. Returns ------- - list - A list containing default values for a new item. + tuple + A tuple containing default values for a new item. """ - return ["<", 0.0, False] + return ("<", 0.0, False) def new_item_prompt(self) -> str: """ @@ -80,7 +80,7 @@ def heads_up(self, name: str) -> None: f"Constraint {name} already exists!", ) - def create_cell_widgets(self, info: list[Any]) -> tuple[QWidget, ...]: + def create_cell_widgets(self, info: tuple[str, float, bool]) -> tuple[QWidget, ...]: # Relation relation_combo = QComboBox() relation_combo.setItemDelegate(QStyledItemDelegate()) diff --git a/src/badger/gui/components/constraint_item.py b/src/badger/gui/components/constraint_item.py index 8574212c..a086e4fa 100644 --- a/src/badger/gui/components/constraint_item.py +++ b/src/badger/gui/components/constraint_item.py @@ -1,6 +1,9 @@ """Creates a single constraint row widget: observable selector, relation combo (>, <, =), threshold spinbox, criticality checkbox, and remove button.""" +from collections.abc import Callable + +from PyQt5 import QtCore from PyQt5.QtCore import Qt from PyQt5.QtWidgets import ( QAbstractSpinBox, @@ -19,11 +22,33 @@ ) +class ConstraintItem(QWidget): + """A single constraint row that exposes its input widgets as attributes.""" + + cb_obs: NoHoverFocusComboBox + cb_rel: QComboBox + sb: QDoubleSpinBox + check_crit: QCheckBox + btn_del: QPushButton + + def enterEvent(self, a0: QtCore.QEvent | None) -> None: + self.btn_del.show() + + def leaveEvent(self, a0: QtCore.QEvent | None) -> None: + self.btn_del.hide() + + def constraint_item( - options, remove_item, name=None, relation=0, threshold=0, critical=False, decimals=4 -): + options: list[str], + remove_item: Callable[[], None], + name: str = "", + relation: int = 0, + threshold: float = 0, + critical: bool = False, + decimals: int = 4, +) -> ConstraintItem: # relation: 0 for >, 1 for <, 2 for = - widget = QWidget() + widget = ConstraintItem() hbox = QHBoxLayout(widget) hbox.setContentsMargins(2, 2, 2, 2) # hbox.setSpacing(0) @@ -45,7 +70,7 @@ def constraint_item( widget.sb = sb = QDoubleSpinBox() sb.setDecimals(decimals) - sb.setFocusPolicy(Qt.StrongFocus) + sb.setFocusPolicy(Qt.FocusPolicy.StrongFocus) sb.installEventFilter(MouseWheelWidgetAdjustmentGuard(sb)) default_value = threshold lb = default_value - 1e3 @@ -69,13 +94,4 @@ def constraint_item( btn_del.clicked.connect(remove_item) - def show_button(event): - btn_del.show() - - def hide_button(event): - btn_del.hide() - - widget.enterEvent = show_button - widget.leaveEvent = hide_button - return widget diff --git a/src/badger/gui/components/create_process.py b/src/badger/gui/components/create_process.py index a263ce87..325a2319 100644 --- a/src/badger/gui/components/create_process.py +++ b/src/badger/gui/components/create_process.py @@ -9,6 +9,12 @@ from badger.core_subprocess import run_routine_subprocess from badger.log import get_logging_manager from badger.settings import init_settings +from badger.types import ( + ArgumentQueueType, + DataQueueMessage, + DialogActionMessage, + ProcessWithArgs, +) logger = logging.getLogger(__name__) @@ -32,12 +38,14 @@ def create_subprocess(self) -> None: """ self.stop_event = Event() self.pause_event = Event() - self.args_queue = Queue() - self.data_queue = Queue() + self.args_queue: Queue[ArgumentQueueType] = Queue() + self.data_queue: Queue[DataQueueMessage] = Queue() self.evaluate_queue = Pipe() self.wait_event = Event() - self.dialog_action_queue = Queue() - config_path = init_settings()._instance.config_path + self.dialog_action_queue: Queue[DialogActionMessage] = Queue() + + config_instance = init_settings()._instance + config_path = config_instance.config_path if config_instance else None # Get the logging queue from the centralized manager logging_manager = get_logging_manager() @@ -53,22 +61,23 @@ def create_subprocess(self) -> None: self.stop_event, self.pause_event, self.wait_event, + self.dialog_action_queue, config_path, log_queue, - self.dialog_action_queue, ), ) new_process.start() - self.subprocess_prepared.emit( - { - "process": new_process, - "args_queue": self.args_queue, - "stop_event": self.stop_event, - "pause_event": self.pause_event, - "data_queue": self.data_queue, - "evaluate_queue": self.evaluate_queue, - "wait_event": self.wait_event, - "dialog_action_queue": self.dialog_action_queue, - } + + _p = ProcessWithArgs( + process=new_process, + args_queue=self.args_queue, + data_queue=self.data_queue, + evaluate_queue=self.evaluate_queue, + stop_event=self.stop_event, + pause_event=self.pause_event, + wait_event=self.wait_event, + dialog_action_queue=self.dialog_action_queue, ) + + self.subprocess_prepared.emit(_p) self.finished.emit() diff --git a/src/badger/gui/components/data_panel.py b/src/badger/gui/components/data_panel.py index e64a339b..9b062203 100644 --- a/src/badger/gui/components/data_panel.py +++ b/src/badger/gui/components/data_panel.py @@ -1,7 +1,11 @@ """Panel for viewing and managing pre-loaded optimization data. Lets users load data from archived runs or clear the buffer before starting.""" +import logging +from typing import overload + import pandas as pd +from gest_api.vocs import VOCS from PyQt5.QtCore import Qt from PyQt5.QtWidgets import ( QCheckBox, @@ -13,7 +17,6 @@ QVBoxLayout, QWidget, ) -from xopt.vocs import VOCS from badger.gui.components.data_table import ( TableWithCopy, @@ -27,6 +30,8 @@ ) from badger.routine import Routine +logger = logging.getLogger(__name__) + LABEL_WIDTH = 96 stylesheet_data = """ @@ -45,9 +50,8 @@ class BadgerDataPanel(QWidget): - def __init__(self, parent=None): + def __init__(self, parent: QWidget | None = None): super().__init__(parent) - self.parent = parent # Set up ui self.init_ui() @@ -55,12 +59,18 @@ def __init__(self, parent=None): # DataFrame to keep track of data in table including metadata self.table_data = pd.DataFrame() - self.selected_routine = None # Keeps track of VOCS in displayed data - self.env_vocs = None # Keeps track of selected VOCS from environment tab + self.selected_routine: Routine | None = ( + None # Keeps track of VOCS in displayed data + ) + self.env_vocs: VOCS | None = ( + None # Keeps track of selected VOCS from environment tab + ) # boolean indicating whether to show metadata in table self.info = False + self.tc_dialog: BadgerLoadDataFromRunDialog | None = None + def init_ui(self) -> None: """Initialize interface""" vbox = QVBoxLayout(self) @@ -77,7 +87,7 @@ def init_ui(self) -> None: padding: 4px; """ ) - title_label.setAlignment(Qt.AlignCenter) # Center-align the title + title_label.setAlignment(Qt.AlignmentFlag.AlignCenter) # Center-align the title vbox_table.addWidget(title_label, 0) vbox.addWidget(panel_table) @@ -105,7 +115,7 @@ def init_ui(self) -> None: self.btn_reset_table.setFixedSize(96, 24) hbox_load_data_options.addWidget(self.btn_load_data) hbox_load_data_options.addWidget(self.btn_reset_table) - hbox_load_data_options.setAlignment(Qt.AlignLeft) + hbox_load_data_options.setAlignment(Qt.AlignmentFlag.AlignLeft) hbox_load_data_options.addStretch() self.info_checkbox = QCheckBox("Display metadata ") self.info_checkbox.setToolTip( @@ -139,6 +149,8 @@ def init_ui(self) -> None: def show_metadata(self) -> None: self.info = self.info_checkbox.isChecked() if self.has_data: + if self.selected_routine is None: + raise ValueError("No routine selected.") self.update_table( self.data_table, self.table_data, vocs=self.selected_routine.vocs ) @@ -146,9 +158,11 @@ def show_metadata(self) -> None: headers = get_horizontal_header_as_list(self.data_table) if "live" in headers: i = headers.index("live") - self.data_table.horizontalHeaderItem(i).setToolTip( - "Indicates whether data was acquired \n(1) live during the current run, or \n(0) imported from previous optimizations" - ) + header_item = self.data_table.horizontalHeaderItem(i) + if header_item is not None: + header_item.setToolTip( + "Indicates whether data was acquired \n(1) live during the current run, or \n(0) imported from previous optimizations" + ) def update_vocs(self, vocs: VOCS) -> None: """ @@ -169,20 +183,20 @@ def load_from_dialog(self) -> None: """ Verify that variables and objectives have been selected, then open dialog to load data """ - vocs = self.env_vocs - if not vocs.variable_names or not vocs.objective_names: - dialog = QMessageBox( - text="Select Environment + VOCS before adding data!", - parent=self, - ) + if self.env_vocs is None: + raise ValueError("Environment VOCS not set.") + + if not self.env_vocs.variable_names or not self.env_vocs.objective_names: + dialog = QMessageBox(self) + dialog.setText("Select Environment + VOCS before adding data!") dialog.setIcon(QMessageBox.Information) dialog.setStandardButtons(QMessageBox.Ok) _ = dialog.exec_() return else: - _ = self.get_data_from_dialog() + self.get_data_from_dialog() def set_routine(self, routine: Routine) -> None: self.selected_routine = routine @@ -225,6 +239,8 @@ def get_data_from_dialog(self) -> None: """ Opens a dialog window for loading data into generator. """ + if self.env_vocs is None: + raise ValueError("Environment VOCS not set.") dlg = BadgerLoadDataFromRunDialog( parent=self, env_vocs=self.env_vocs.variable_names + self.env_vocs.output_names, @@ -256,7 +272,7 @@ def add_live_data(self, data: pd.DataFrame) -> None: self.update_table(self.data_table, all_data, vocs) def update_table( - self, table: TableWithCopy, data: pd.DataFrame = None, vocs: VOCS = None + self, table: TableWithCopy, data: pd.DataFrame, vocs: VOCS = None ) -> None: """Makes sure column order remains consistent and self.table_data stays updated to match displayed table. Call data_table's update_table method to update the table @@ -268,7 +284,7 @@ def update_table( def get_data(self) -> pd.DataFrame: return self.table_data - def get_data_as_dict(self) -> dict: + def get_data_as_dict(self) -> dict[str, list[str]]: data = get_table_content_as_dict(self.data_table) if not self.info: data = filter_metadata(data) @@ -276,7 +292,7 @@ def get_data_as_dict(self) -> dict: return data @property - def routine(self) -> Routine: + def routine(self) -> Routine | None: return self.selected_routine def load_data_from_dialog(self, routine: Routine) -> None: @@ -290,6 +306,9 @@ def load_data_from_dialog(self, routine: Routine) -> None: """ # Data from routine to load data = routine.data + if data is None: + logger.warning("No data found in the selected routine") + return # Create copy of data without metadata columns filtered_data = filter_metadata(data) data_keys = list(filtered_data.keys()) @@ -300,12 +319,10 @@ def load_data_from_dialog(self, routine: Routine) -> None: # Raise error if loaded data keys do not match selected vocs # This happens here if selected VOCS have been changed but old data is still in the table. if self.has_data and set(data_keys) != set(filtered_table_keys): - dialog = QMessageBox( - text=( - "Keys in loaded data do not match current table!" - "\nTry clearing the table before adding new data." - ), - parent=self, + dialog = QMessageBox(self) + dialog.setText( + "Keys in loaded data do not match current table!" + "\nTry clearing the table before adding new data." ) dialog.setIcon(QMessageBox.Warning) dialog.setStandardButtons(QMessageBox.Ok) @@ -329,6 +346,9 @@ def load_data(self, routine: Routine) -> None: """ self.set_routine(routine) data = routine.data + if data is None: + logger.warning("No data found in the selected routine") + return self.run_data_checkbox.setEnabled(True) @@ -348,7 +368,11 @@ def _reorder_cols(self, data: pd.DataFrame) -> pd.DataFrame: columns = list(data.columns) columns_set = set(columns) reordered_cols = [] - vocs = self.selected_routine.vocs + routine = self.selected_routine + if routine is None: + logger.warning("No routine selected") + return data + vocs = routine.vocs priority_groups = [ vocs.objective_names, @@ -383,9 +407,19 @@ def reset_data_table(self) -> None: self.run_data_checkbox.setEnabled(False) -def filter_metadata(data: dict) -> dict: +@overload +def filter_metadata(data: dict[str, list[str]]) -> dict[str, list[str]]: ... + + +@overload +def filter_metadata(data: pd.DataFrame) -> pd.DataFrame: ... + + +def filter_metadata( + data: dict[str, list[str]] | pd.DataFrame, +) -> dict[str, list[str]] | pd.DataFrame: """ - Remove metadata columns from data dictionary + Remove metadata columns from a data dict or DataFrame. """ data_copy = data.copy() diff --git a/src/badger/gui/components/data_table.py b/src/badger/gui/components/data_table.py index d962441c..5fcb2d73 100644 --- a/src/badger/gui/components/data_table.py +++ b/src/badger/gui/components/data_table.py @@ -2,10 +2,17 @@ optimization data (variables, objectives, constraints) with clipboard copy and alternating-row styling.""" +import logging +from typing import Any + +from gest_api.vocs import VOCS from pandas import DataFrame from PyQt5.QtCore import Qt +from PyQt5.QtGui import QKeyEvent from PyQt5.QtWidgets import QApplication, QTableWidget, QTableWidgetItem +logger = logging.getLogger(__name__) + stylesheet = """ QTableWidget { @@ -41,35 +48,40 @@ class TableWithCopy(QTableWidget): like google sheets, excel, or numbers """ - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) - - def keyPressEvent(self, event): + def keyPressEvent(self, event: QKeyEvent | None) -> None: super().keyPressEvent(event) - if event.key() == Qt.Key.Key_C and ( - event.modifiers() & Qt.KeyboardModifier.ControlModifier + if ( + event is not None + and event.key() == Qt.Key.Key_C + and (event.modifiers() & Qt.KeyboardModifier.ControlModifier) ): - copied_cells = sorted(self.selectedIndexes()) + copied_cells = sorted( + self.selectedIndexes(), key=lambda idx: (idx.row(), idx.column()) + ) copy_text = "" max_column = copied_cells[-1].column() for c in copied_cells: - copy_text += self.item(c.row(), c.column()).text() + item = self.item(c.row(), c.column()) + if item is not None: + copy_text += item.text() if c.column() == max_column: copy_text += "\n" else: copy_text += "\t" - QApplication.clipboard().setText(copy_text) + clipboard = QApplication.clipboard() + if clipboard is not None: + clipboard.setText(copy_text) - def set_uneditable(self): + def set_uneditable(self) -> None: self.setEditTriggers(QTableWidget.NoEditTriggers) - def set_editable(self): + def set_editable(self) -> None: self.setEditTriggers(QTableWidget.DoubleClicked) -def format_value(v): +def format_value(v: Any) -> str: try: f = f"{v:.6g}" except (ValueError, TypeError): @@ -77,18 +89,31 @@ def format_value(v): return f -def resize_columns_to_content_or_default(table): +def resize_columns_to_content_or_default(table: QTableWidget) -> None: """Resize columns to the maximum of default width and content width.""" - default_width = table.horizontalHeader().defaultSectionSize() + hheader = table.horizontalHeader() + if hheader is None: + logger.warning("Horizontal header is None, cannot resize columns.") + return + default_width = hheader.defaultSectionSize() table.resizeColumnsToContents() for col in range(table.columnCount()): current_width = table.columnWidth(col) table.setColumnWidth(col, max(default_width, current_width)) -def update_table(table, data=None, vocs=None, info=False): +def update_table( + table: TableWithCopy, + data: DataFrame | None = None, + vocs: VOCS | None = None, + info: bool = False, +) -> TableWithCopy: table.setRowCount(0) - table.horizontalHeader().setVisible(False) + hheader = table.horizontalHeader() + if hheader is None: + logger.warning("Horizontal header is None, cannot hide it.") + else: + hheader.setVisible(False) if data is None: return table @@ -117,25 +142,37 @@ def update_table(table, data=None, vocs=None, info=False): table.setVerticalHeaderLabels( list(map(str, _data.index)) ) # row index starts from 0 - table.horizontalHeader().setVisible(True) + hheader = table.horizontalHeader() + if hheader is None: + logger.warning("Horizontal header is None, cannot show it.") + else: + hheader.setVisible(True) resize_columns_to_content_or_default(table) return table -def reset_table(table, header): +def reset_table(table: TableWithCopy, header: list[str]) -> TableWithCopy: table.setRowCount(0) # Need to set col num or the old col num will be used for new data, # resulting in potential incomplete table table.setColumnCount(len(header)) - table.horizontalHeader().setVisible(False) + hheader = table.horizontalHeader() + if hheader is None: + logger.warning("Horizontal header is None, cannot hide it.") + else: + hheader.setVisible(False) table.setHorizontalHeaderLabels(header) - table.horizontalHeader().setVisible(True) + hheader = table.horizontalHeader() + if hheader is None: + logger.warning("Horizontal header is None, cannot show it.") + else: + hheader.setVisible(True) return table -def add_row(table, row): +def add_row(table: TableWithCopy, row: list[float | str]) -> TableWithCopy: r = table.rowCount() table.insertRow(r) for i, v in enumerate(row): @@ -145,7 +182,7 @@ def add_row(table, row): return table -def data_table(data=None): +def data_table(data: DataFrame | None = None) -> TableWithCopy: table = TableWithCopy() table.setAlternatingRowColors(True) table.setStyleSheet(stylesheet_data) @@ -153,7 +190,7 @@ def data_table(data=None): return update_table(table, data) -def init_data_table(variable_names=None): +def init_data_table(variable_names: list[str] | None = None) -> TableWithCopy: table = TableWithCopy() table.setAlternatingRowColors(True) table.setStyleSheet(stylesheet) @@ -164,24 +201,38 @@ def init_data_table(variable_names=None): return table table.setColumnCount(len(variable_names)) - table.horizontalHeader().setVisible(False) + hheader = table.horizontalHeader() + if hheader is None: + logger.warning("Horizontal header is None, cannot hide it.") + else: + hheader.setVisible(False) table.setHorizontalHeaderLabels(variable_names) - table.horizontalHeader().setVisible(True) + hheader = table.horizontalHeader() + if hheader is None: + logger.warning("Horizontal header is None, cannot show it.") + else: + hheader.setVisible(True) resize_columns_to_content_or_default(table) return table -def get_horizontal_header_as_list(table): +def get_horizontal_header_as_list(table: TableWithCopy) -> list[str]: header = table.horizontalHeader() + if header is None: + logger.warning("Horizontal header is None, cannot read labels.") + return [] + model = header.model() + if model is None: + logger.warning("Header model is None, cannot read labels.") + return [] header_labels = [ - header.model().headerData(i, header.orientation()) - for i in range(header.count()) + model.headerData(i, header.orientation()) for i in range(header.count()) ] return header_labels -def get_table_content_as_dict(table): +def get_table_content_as_dict(table: TableWithCopy) -> dict[str, list[str]]: table_content = {} header_labels = get_horizontal_header_as_list(table) @@ -201,13 +252,21 @@ def get_table_content_as_dict(table): return table_content -def update_init_data_table(table, variable_names): +def update_init_data_table(table: TableWithCopy, variable_names: list[str]) -> None: current_init_data = get_table_content_as_dict(table) table.setColumnCount(len(variable_names)) - table.horizontalHeader().setVisible(False) + hheader = table.horizontalHeader() + if hheader is None: + logger.warning("Horizontal header is None, cannot hide it.") + else: + hheader.setVisible(False) table.setHorizontalHeaderLabels(variable_names) - table.horizontalHeader().setVisible(True) + hheader = table.horizontalHeader() + if hheader is None: + logger.warning("Horizontal header is None, cannot show it.") + else: + hheader.setVisible(True) for col, name in enumerate(variable_names): if name in current_init_data: @@ -218,7 +277,7 @@ def update_init_data_table(table, variable_names): table.setItem(row, col, QTableWidgetItem("")) -def set_init_data_table(table, data: DataFrame): +def set_init_data_table(table: TableWithCopy, data: DataFrame) -> None: variable_names = get_horizontal_header_as_list(table) try: data_dict = data.to_dict("list") diff --git a/src/badger/gui/components/editable_table.py b/src/badger/gui/components/editable_table.py index 84397cf2..f34fbd59 100644 --- a/src/badger/gui/components/editable_table.py +++ b/src/badger/gui/components/editable_table.py @@ -23,9 +23,8 @@ import logging from collections.abc import Callable from functools import partial, wraps -from typing import Any, ParamSpec, cast +from typing import Any, Concatenate, ParamSpec, TypeVar, cast -from pyparsing import TypeVar from PyQt5.QtCore import QRegExp, Qt, pyqtSignal from PyQt5.QtGui import QColor, QDragEnterEvent, QDragMoveEvent, QDropEvent from PyQt5.QtWidgets import ( @@ -45,6 +44,8 @@ T = TypeVar("T") P = ParamSpec("P") +# Bound to the class so the wrapper can call self.blockSignals(...). +S = TypeVar("S", bound="EditableTable") class EditableTable(QTableWidget): @@ -87,7 +88,7 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: header.setSectionResizeMode(1, QHeaderView.ResizeMode.Stretch) self.setColumnWidth(0, 20) # width for checkboxes - self.data: list[dict[str, Any]] = [] + self.data: list[dict[str, tuple[str, float, bool]]] = [] self.status: dict[str, bool] = {} # track selection self.formulas: dict[str, dict[str, Any]] = {} # track formula item @@ -109,16 +110,16 @@ def update_vocs(self) -> None: logger.debug("Emitting data_changed signal from editable_table") self.data_changed.emit() - def default_info(self) -> list[Any]: + def default_info(self) -> tuple[str, float, bool]: """ Get the default information list for a new item. Returns ------- - list - A list containing default values for a new item. + tuple + A tuple containing default values for a new item. """ - return ["<", 0.0, False] + return ("<", 0.0, False) def new_item_prompt(self) -> str: """ @@ -141,7 +142,7 @@ def heads_up(self, name: str) -> None: f"Item {name} already exists!", ) - def create_cell_widgets(self, info: list[Any]) -> tuple[QWidget, ...]: + def create_cell_widgets(self, info: tuple[str, float, bool]) -> tuple[QWidget, ...]: # Relation relation_combo = QComboBox() relation_combo.setItemDelegate(QStyledItemDelegate()) @@ -166,7 +167,7 @@ def dragEnterEvent(self, e: QDragEnterEvent | None) -> None: Parameters ---------- - event : QDragEnterEvent + e : QDragEnterEvent | None The drag enter event. """ # Accept internal moves (reordering) or external text drops. @@ -190,7 +191,7 @@ def dragMoveEvent(self, e: QDragMoveEvent | None) -> None: Parameters ---------- - event : QDragMoveEvent + e : QDragMoveEvent | None The drag move event. """ @@ -220,7 +221,7 @@ def dropEvent(self, event: QDropEvent | None) -> None: Parameters ---------- - event : QDropEvent + event : QDropEvent | None The drop event. """ @@ -275,21 +276,25 @@ def item_names(self) -> list[str]: return [next(iter(item)) for item in self.data] @staticmethod - def block_signals(func: Callable[..., Any]) -> Callable[..., Any]: + def block_signals( + func: Callable[Concatenate[S, P], T], + ) -> Callable[Concatenate[S, P], T]: """ A decorator to block signals at the beginning of a function and unblock them at the end. """ @wraps(func) - def wrapper(self: "EditableTable", *args: P.args, **kwargs: P.kwargs) -> Any: + def wrapper(self: S, *args: P.args, **kwargs: P.kwargs) -> T: self.blockSignals(True) try: return func(self, *args, **kwargs) finally: self.blockSignals(False) - return wrapper + # wraps() returns a _Wrapped object; it is callable with wrapper's exact + # signature but is not nominally a Callable, so re-assert the type. + return cast(Callable[Concatenate[S, P], T], wrapper) def header_clicked(self, idx: int) -> None: """ @@ -306,10 +311,11 @@ def header_clicked(self, idx: int) -> None: all_checked = True visible_names: list[str] = [] for i in range(self.rowCount() - 1): # Exclude the last empty row - checkbox = self.cellWidget(i, 0) - if checkbox is None: + # PyQt5-stubs mistypes cellWidget as List[QWidget]; it returns one widget. + checkbox_widget = self.cellWidget(i, 0) + if checkbox_widget is None: raise ValueError("Checkbox widget is None") - checkbox = cast(QCheckBox, checkbox) + checkbox = cast(QCheckBox, checkbox_widget) name_item = self.item(i, 1) if name_item is None: @@ -328,7 +334,7 @@ def insert_item( self, row: int, name: str, - info: list[Any], + info: tuple[str, float, bool], selected: bool = False, ) -> None: """ @@ -340,8 +346,8 @@ def insert_item( The row index where the item should be inserted. name : str The name of the item. - info : list - A list containing additional information about the item. + info : tuple[str, float, bool] + A tuple containing additional information about the item. selected : bool, optional Whether the item is selected, default is False. """ @@ -359,19 +365,18 @@ def insert_item( name_item.setForeground(QColor("darkCyan")) else: # Make non-new PVs not editable - name_item.setFlags(name_item.flags() & ~Qt.ItemIsEditable) + name_item.setFlags( + name_item.flags() & ~Qt.ItemFlags(Qt.ItemFlag.ItemIsEditable) + ) self.setItem(row, 1, name_item) cell_widgets = self.create_cell_widgets(info) for i, widget in enumerate(cell_widgets): if isinstance(widget, QComboBox): - widget = cast(QComboBox, widget) widget.currentIndexChanged.connect(partial(self.update_info, i)) elif isinstance(widget, QDoubleSpinBox): - widget = cast(QDoubleSpinBox, widget) widget.valueChanged.connect(partial(self.update_info, i)) elif isinstance(widget, QCheckBox): - widget = cast(QCheckBox, widget) widget.stateChanged.connect(partial(self.update_info, i)) else: raise NotImplementedError(f"Unsupported widget type: {type(widget)}") @@ -428,7 +433,7 @@ def add_formula_item(self, formula_tuple: tuple[str, str, dict[str, str]]) -> No print(f"Error adding formula item: {e}") return - def add_plain_item(self, name: str): + def add_plain_item(self, name: str) -> None: self.add_formula_item((name, "", {})) def get_visible_items(self) -> list[str]: @@ -540,14 +545,14 @@ def on_edit_table_item(self, item: QTableWidgetItem) -> None: if self.keyword and QRegExp(self.keyword).indexIn(name, 0) == -1: self.removeRow(row) - def add_empty_row(self): + def add_empty_row(self) -> None: """ Add an empty row for entering a new constraint. """ row = self.rowCount() self.setRowCount(row + 1) item = QTableWidgetItem(self.new_item_prompt()) - item.setFlags(item.flags() | Qt.ItemIsEditable) + item.setFlags(item.flags() | Qt.ItemFlag.ItemIsEditable) item.setForeground(QColor("gray")) self.setItem(row, 1, item) @@ -594,13 +599,10 @@ def update_info(self, idx: int) -> None: if cell_widget is not None: item = self.get_item_by_name(name) if isinstance(cell_widget, QComboBox): - cell_widget = cast(QComboBox, cell_widget) - value = cell_widget.currentText() + value: str | float | bool = cell_widget.currentText() elif isinstance(cell_widget, QDoubleSpinBox): - cell_widget = cast(QDoubleSpinBox, cell_widget) value = cell_widget.value() elif isinstance(cell_widget, QCheckBox): - cell_widget = cast(QCheckBox, cell_widget) value = cell_widget.isChecked() else: raise NotImplementedError( @@ -635,7 +637,7 @@ def update_keyword(self, keyword: str) -> None: def update_items( self, - data: list[dict[str, Any]] | None = None, + data: list[dict[str, tuple[str, float, bool]]] | None = None, status: dict[str, bool] | None = None, formulas: dict[str, dict[str, Any]] | None = None, vocs_signal: bool = True, @@ -647,7 +649,7 @@ def update_items( @block_signals def update_items_wrapper( self, - data: list[dict[str, Any]] | None = None, + data: list[dict[str, tuple[str, float, bool]]] | None = None, status: dict[str, bool] | None = None, formulas: dict[str, dict[str, Any]] | None = None, ) -> None: @@ -689,16 +691,16 @@ def update_items_wrapper( # Add an empty row for new constraints self.add_empty_row() - def export_data(self) -> list[dict[str, Any]]: + def export_data(self) -> list[dict[str, tuple[str, float, bool]]]: """ - Export the items as a list of dictionaries. + Export the items as a list of dictionaries with tuple info for each item. Returns ------- - List[Dict[str, Any]] + List[Dict[str, tuple[str, float, bool]]] A list of items with their properties. """ - exported_items: list[dict[str, Any]] = [] + exported_items: list[dict[str, tuple[str, float, bool]]] = [] for item in self.data: if self.status.get(next(iter(item)), False): exported_items.append(item) diff --git a/src/badger/gui/components/eliding_label.py b/src/badger/gui/components/eliding_label.py index 1bd890b7..63b41951 100644 --- a/src/badger/gui/components/eliding_label.py +++ b/src/badger/gui/components/eliding_label.py @@ -1,10 +1,15 @@ """QLabel that truncates text with "..." when it doesn't fit. Used for long routine names and status messages in tight layouts.""" -from PyQt5 import QtCore, QtGui, QtWidgets +# https://stackoverflow.com/a/67628976/4263605 +from typing import Any +from PyQt5.QtCore import QPoint, Qt, pyqtSignal +from PyQt5.QtGui import QFontMetrics, QPainter, QPaintEvent, QResizeEvent, QTextLayout +from PyQt5.QtWidgets import QLabel, QSizePolicy, QWidget -class ElidingLabel(QtWidgets.QLabel): + +class ElidingLabel(QLabel): """Label with text elision. QLabel which will elide text too long to fit the widget. Based on: @@ -38,19 +43,25 @@ class ElidingLabel(QtWidgets.QLabel): """ - elision_changed = QtCore.pyqtSignal(bool) + elision_changed = pyqtSignal(bool) - def __init__(self, text="", mode=QtCore.Qt.ElideMiddle, **kwargs): + def __init__( + self, + text: str = "", + mode: Qt.TextElideMode = Qt.TextElideMode.ElideMiddle, + **kwargs: Any, + ) -> None: super().__init__(**kwargs) self._mode = mode self.is_elided = False - self.setSizePolicy( - QtWidgets.QSizePolicy.Expanding, QtWidgets.QSizePolicy.Preferred - ) + self.setSizePolicy(QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Preferred) self.setText(text) - def setText(self, text): + def setText(self, text: str | None) -> None: + if text is None: + text = "" + self._contents = text # This line set for testing. Its value is the return value of @@ -61,20 +72,20 @@ def setText(self, text): self.update() - def text(self): + def text(self) -> str: return self._contents - def paintEvent(self, event): + def paintEvent(self, event: QPaintEvent | None) -> None: super().paintEvent(event) did_elide = False - painter = QtGui.QPainter(self) + painter = QPainter(self) font_metrics = painter.fontMetrics() text_width = font_metrics.horizontalAdvance(self.text()) # layout phase - text_layout = QtGui.QTextLayout(self._contents, painter.font()) + text_layout = QTextLayout(self._contents, painter.font()) text_layout.beginLayout() while True: @@ -89,13 +100,11 @@ def paintEvent(self, event): self._elided_line = font_metrics.elidedText( self._contents, self._mode, self.width() ) - painter.drawText( - QtCore.QPoint(0, font_metrics.ascent()), self._elided_line - ) + painter.drawText(QPoint(0, font_metrics.ascent()), self._elided_line) did_elide = line.isValid() break else: - line.draw(painter, QtCore.QPoint(0, 0)) + line.draw(painter, QPoint(0, 0)) text_layout.endLayout() @@ -104,20 +113,22 @@ def paintEvent(self, event): self.elision_changed.emit(did_elide) -class SimpleElidedLabel(QtWidgets.QLabel): - def __init__(self, text="", parent=None): +class SimpleElidedLabel(QLabel): + def __init__(self, text: str | None = None, parent: QWidget | None = None) -> None: super().__init__(parent) self._text = text - def setText(self, text): + def setText(self, text: str | None) -> None: self._text = text super().setText(self.elidedText()) - def resizeEvent(self, event): + def resizeEvent(self, event: QResizeEvent | None) -> None: super().setText(self.elidedText()) super().resizeEvent(event) - def elidedText(self): - metrics = QtGui.QFontMetrics(self.font()) - elided = metrics.elidedText(self._text, QtCore.Qt.ElideRight, self.width()) + def elidedText(self) -> str: + metrics = QFontMetrics(self.font()) + elided = metrics.elidedText( + self._text, Qt.TextElideMode.ElideRight, self.width() + ) return elided diff --git a/src/badger/gui/components/env_cbox.py b/src/badger/gui/components/env_cbox.py index 8342d717..a77de27e 100644 --- a/src/badger/gui/components/env_cbox.py +++ b/src/badger/gui/components/env_cbox.py @@ -22,7 +22,7 @@ from importlib import resources from typing import Any -from gest_api.vocs import ContinuousVariable +from gest_api.vocs import VOCS, ContinuousVariable from pydantic_core import ValidationError from PyQt5.QtCore import QPropertyAnimation, QRegExp, pyqtSignal from PyQt5.QtGui import QFont, QIcon @@ -31,13 +31,13 @@ QHBoxLayout, QLabel, QLineEdit, + QListWidget, QPushButton, QSizePolicy, QStyledItemDelegate, QVBoxLayout, QWidget, ) -from xopt.vocs import VOCS from badger.errors import BadgerRoutineError from badger.gui.components.collapsible_box import CollapsibleBox @@ -167,7 +167,7 @@ def __init__( self.init_ui() self.config_logic() - def init_ui(self): + def init_ui(self) -> None: config_singleton = init_settings() icon_ref = resources.files(__package__) / "../images/import.png" @@ -520,8 +520,7 @@ def init_ui(self): cbox_more.setContentLayout(vbox_more) - def config_logic(self): - self.dict_con = {} + def config_logic(self) -> None: self.edit_var.textChanged.connect(self.filter_var) self.check_only_var.stateChanged.connect(self.toggle_var_show_mode) @@ -544,7 +543,7 @@ def update_vocs(self, origin: str) -> None: vocs, _ = self.compose_vocs() self.vocs_updated.emit(vocs) - def toggle_params(self, checked: bool): + def toggle_params(self, checked: bool) -> None: if not checked: self.animation.setStartValue(self.edit.sizeHint().height()) self.animation.setEndValue(0) @@ -557,18 +556,18 @@ def toggle_params(self, checked: bool): # Configure the animation self.animation.start() - def animation_finished(self): + def animation_finished(self) -> None: if self.edit.maximumHeight() == 0: self.edit.hide() - def toggle_var_show_mode(self, _): + def toggle_var_show_mode(self, _: Any) -> None: self.var_table.toggle_show_mode(self.check_only_var.isChecked()) - def add_var(self, name, lb, ub): + def add_var(self, name: str, lb: float, ub: float) -> None: self.var_table.add_variable(name, lb, ub) self.filter_var() - def filter_var(self): + def filter_var(self) -> None: keyword = self.edit_var.text() rx = QRegExp(keyword) @@ -580,33 +579,33 @@ def filter_var(self): self.var_table.update_variables(_variables, 1) - def toggle_obj_show_mode(self, _): + def toggle_obj_show_mode(self, _: Any) -> None: self.obj_table.update_show_selected_only(self.check_only_obj.isChecked()) - def filter_obj(self): + def filter_obj(self) -> None: self.obj_table.update_keyword(self.edit_obj.text()) - def toggle_con_show_mode(self, _): + def toggle_con_show_mode(self, _: Any) -> None: self.con_table.update_show_selected_only(self.check_only_con.isChecked()) - def filter_con(self): + def filter_con(self) -> None: self.con_table.update_keyword(self.edit_con.text()) - def toggle_sta_show_mode(self, _): + def toggle_sta_show_mode(self, _: Any) -> None: self.sta_table.update_show_selected_only(self.check_only_sta.isChecked()) - def filter_sta(self): + def filter_sta(self) -> None: self.sta_table.update_keyword(self.edit_sta.text()) - def _fit_content(self, list): + def _fit_content(self, list: QListWidget) -> None: height = list.sizeHintForRow(0) * list.count() + 2 * list.frameWidth() + 4 height = max(28, min(height, 192)) list.setFixedHeight(height) - def fit_content(self): + def fit_content(self) -> None: return - def switch_var_panel_style(self, auto=True): + def switch_var_panel_style(self, auto: bool = True) -> None: if auto: self.var_panel.setStyleSheet(stylesheet_auto) self.msg_auto.setStyleSheet(stylesheet_auto_msg) @@ -616,7 +615,7 @@ def switch_var_panel_style(self, auto=True): self.msg_auto.setStyleSheet(stylesheet_manual_msg) self.msg_auto.setText(self.MSG_MANUAL) - def update_stylesheets(self, environment=""): + def update_stylesheets(self, environment: str = "") -> None: if environment in self.env_dict: color_dict = self.env_dict[environment] stylesheet = f""" @@ -634,13 +633,13 @@ def compose_vocs(self) -> tuple[VOCS, list[str]]: # Compose the VOCS settings variables = self.var_table.export_variables() - objectives: dict[str, Any] = {} + objectives: dict[str, str] = {} for objective in self.obj_table.export_data(): obj_name = next(iter(objective)) - (rule,) = objective[obj_name] + rule = objective[obj_name][0] objectives[obj_name] = rule - constraints: dict[str, list[float]] = {} + constraints: dict[str, list[str | float]] = {} critical_constraints: list[str] = [] for constraint in self.con_table.export_data(): con_name = next(iter(constraint)) @@ -656,12 +655,12 @@ def compose_vocs(self) -> tuple[VOCS, list[str]]: try: # We want to ensure it's a dict of either lists or ContinuousVariables - variables = { + processed_variables: dict[str, list[float] | ContinuousVariable] = { k: list(v) if not isinstance(v, ContinuousVariable) else v for k, v in variables.items() } vocs = VOCS( - variables=variables, + variables=processed_variables, objectives=objectives, constraints=constraints, constants={}, diff --git a/src/badger/gui/components/filter_cbox.py b/src/badger/gui/components/filter_cbox.py index 6ed9857d..e616fcd6 100644 --- a/src/badger/gui/components/filter_cbox.py +++ b/src/badger/gui/components/filter_cbox.py @@ -14,12 +14,18 @@ class BadgerFilterBox(CollapsibleBox): - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) + def __init__( + self, + parent: QWidget | None = None, + title: str = "", + duration: int = 100, + tooltip: str = "", + ): + super().__init__(parent=parent, title=title, duration=duration, tooltip=tooltip) self.init_ui() - def init_ui(self): + def init_ui(self) -> None: vbox = QVBoxLayout() # Obj filter diff --git a/src/badger/gui/components/generator_cbox.py b/src/badger/gui/components/generator_cbox.py index a4bf3629..d73f8637 100644 --- a/src/badger/gui/components/generator_cbox.py +++ b/src/badger/gui/components/generator_cbox.py @@ -1,6 +1,7 @@ """Panel where users pick an optimization algorithm from the Xopt registry and configure its parameters via the Pydantic tree editor.""" +from gest_api.vocs import VOCS from PyQt5.QtWidgets import ( QCheckBox, QComboBox, @@ -12,7 +13,6 @@ QVBoxLayout, QWidget, ) -from xopt.vocs import VOCS from badger.gui.components.collapsible_box import CollapsibleBox from badger.gui.components.pydantic_editor import BadgerPydanticEditor @@ -29,10 +29,10 @@ class BadgerAlgoBox(QWidget): def __init__( self, - parent=None, - generators: list | None = None, - scaling_functions: list | None = None, - ): + parent: QWidget | None = None, + generators: list[str] | None = None, + scaling_functions: list[str] | None = None, + ) -> None: if generators is None: generators = [] if scaling_functions is None: @@ -44,7 +44,7 @@ def __init__( self.init_ui() - def init_ui(self): + def init_ui(self) -> None: config_singleton = init_settings() vbox = QVBoxLayout(self) @@ -145,5 +145,5 @@ def init_ui(self): if not strtobool(config_singleton.read_value("BADGER_ENABLE_ADVANCED")): cbox_misc.hide() - def update_vocs(self, vocs: VOCS): + def update_vocs(self, vocs: VOCS) -> None: self.edit.update_vocs(vocs) diff --git a/src/badger/gui/components/labeled_lineedit.py b/src/badger/gui/components/labeled_lineedit.py index 085b3f84..1f91df53 100644 --- a/src/badger/gui/components/labeled_lineedit.py +++ b/src/badger/gui/components/labeled_lineedit.py @@ -4,7 +4,13 @@ from PyQt5.QtWidgets import QHBoxLayout, QLabel, QLineEdit, QWidget -def labeled_lineedit(name, text, width_name=64, placeholder=None, readonly=True): +def labeled_lineedit( + name: str, + text: str, + width_name: int = 64, + placeholder: str | None = None, + readonly: bool = True, +) -> QWidget: widget = QWidget() hbox = QHBoxLayout(widget) hbox.setContentsMargins(0, 0, 0, 0) diff --git a/src/badger/gui/components/navigators.py b/src/badger/gui/components/navigators.py index e661684b..88fda3ef 100644 --- a/src/badger/gui/components/navigators.py +++ b/src/badger/gui/components/navigators.py @@ -3,8 +3,11 @@ TemplateNavigator shows routine template files. Both support right-click context menus (open, copy path, etc.).""" +import logging import os +from typing import cast +from PyQt5 import QtCore from PyQt5.QtCore import QDir, Qt, QTimer, QUrl from PyQt5.QtGui import QCursor, QDesktopServices, QFont from PyQt5.QtWidgets import ( @@ -24,6 +27,8 @@ from badger.settings import init_settings from badger.utils import run_names_to_dict +logger = logging.getLogger(__name__) + class FileContextMenuBase: """ @@ -34,7 +39,7 @@ class FileContextMenuBase: - Copy full file path to clipboard """ - def show_context_menu(self, widget, fullpath): + def show_context_menu(self, widget: QWidget, fullpath: str) -> None: menu = QMenu(widget) menu.setStyleSheet(""" QMenu { @@ -43,8 +48,11 @@ def show_context_menu(self, widget, fullpath): """) # Functions to execute the context-menu actions - def copy_fullpath_to_clipboard(): + def copy_fullpath_to_clipboard() -> None: clip = QApplication.clipboard() + if clip is None: + logger.warning("Clipboard not available.") + return clip.setText(fullpath) # Display a little popup box a bit after user's click QTimer.singleShot( @@ -56,15 +64,19 @@ def copy_fullpath_to_clipboard(): ), ) - def open_file(): + def open_file() -> None: QDesktopServices.openUrl(QUrl.fromLocalFile(fullpath)) - def open_file_location(): + def open_file_location() -> None: QDesktopServices.openUrl(QUrl.fromLocalFile(os.path.dirname(fullpath))) # Add actions to the context menu - menu.addAction("Open File").triggered.connect(open_file) - menu.addAction("Open File Directory").triggered.connect(open_file_location) + action = menu.addAction("Open File") + if action: + action.triggered.connect(open_file) + action = menu.addAction("Open File Directory") + if action: + action.triggered.connect(open_file_location) # Submenu for copying the full file path fullpath_item = QMenu("File Path", menu) @@ -82,8 +94,8 @@ class HistoryNavigator(QWidget, FileContextMenuBase): Navigate through the history files. """ - def __init__(self): - super().__init__() + def __init__(self, parent: QWidget | None = None) -> None: + super().__init__(parent) # Layout for the widget layout = QVBoxLayout(self) @@ -95,41 +107,51 @@ def __init__(self): # Set the font of the header to bold bold_font = QFont() bold_font.setBold(True) - header.setFont(bold_font) + if header: + header.setFont(bold_font) self.history_tree_widget.setMinimumHeight(256) - self.history_tree_widget.setContextMenuPolicy(Qt.CustomContextMenu) + self.history_tree_widget.setContextMenuPolicy( + Qt.ContextMenuPolicy.CustomContextMenu + ) self.history_tree_widget.customContextMenuRequested.connect( self.show_context_menu_history ) layout.addWidget(self.history_tree_widget) - self.runs = None # all runs to be shown in the tree widget + self.runs: list[str] | None = None # all runs to be shown in the tree widget self.setStyleSheet(""" QTreeWidget { background-color: #37414F; } """) - def _firstSelectableItem(self, parent=None): + def _firstSelectableItem( + self, parent: QTreeWidgetItem | None = None + ) -> QTreeWidgetItem | None: """ Internal recursive function for finding the first selectable item. """ if parent is None: parent = self.history_tree_widget.invisibleRootItem() + if parent is None: + logger.warning("No root item found in the history tree widget.") + return None for i in range(parent.childCount()): item = parent.child(i) - if item.flags() & Qt.ItemIsSelectable: + if not item: + continue + if item.flags() & Qt.ItemFlag.ItemIsSelectable: return item result = self._firstSelectableItem(item) if result: return result return None - def updateItems(self, runs=None): + def updateItems(self, runs: list[str] | None = None) -> None: self.history_tree_widget.clear() self.runs = runs # store the runs for navigation if runs is None: @@ -141,21 +163,36 @@ def updateItems(self, runs=None): for year, dict_year in runs_dict.items(): item_year = QTreeWidgetItem([year]) - item_year.setFlags(item_year.flags() & ~Qt.ItemIsSelectable) + item_year.setFlags( + cast( + Qt.ItemFlags, + item_year.flags() & ~Qt.ItemFlag.ItemIsSelectable, + ) + ) if flag_first_item: first_items.append(item_year) for month, dict_month in dict_year.items(): item_month = QTreeWidgetItem([month]) - item_month.setFlags(item_month.flags() & ~Qt.ItemIsSelectable) + item_month.setFlags( + cast( + Qt.ItemFlags, + item_month.flags() & ~Qt.ItemFlag.ItemIsSelectable, + ) + ) if flag_first_item: first_items.append(item_month) for day, list_day in dict_month.items(): item_day = QTreeWidgetItem([day]) - item_day.setFlags(item_day.flags() & ~Qt.ItemIsSelectable) + item_day.setFlags( + cast( + Qt.ItemFlags, + item_day.flags() & ~Qt.ItemFlag.ItemIsSelectable, + ) + ) if flag_first_item: first_items.append(item_day) @@ -172,47 +209,60 @@ def updateItems(self, runs=None): for item in first_items: item.setExpanded(True) - def selectNextItem(self): + def selectNextItem(self) -> None: run_curr = get_base_run_filename(self.currentText()) + if self.runs is None: + logger.warning("No runs available for navigation.") + return idx = self.runs.index(run_curr) if idx < len(self.runs) - 1: self._selectItemByRun(self.runs[idx + 1]) - def selectPreviousItem(self): + def selectPreviousItem(self) -> None: run_curr = get_base_run_filename(self.currentText()) + if self.runs is None: + return idx = self.runs.index(run_curr) if idx > 0: self._selectItemByRun(self.runs[idx - 1]) - def _selectItemByRun(self, run): + def _selectItemByRun(self, run: str) -> None: """ Internal function to select a tree widget item by run name. """ for i in range(self.history_tree_widget.topLevelItemCount()): year_item = self.history_tree_widget.topLevelItem(i) + if year_item is None: + continue for j in range(year_item.childCount()): month_item = year_item.child(j) + if month_item is None: + continue for k in range(month_item.childCount()): day_item = month_item.child(k) + if day_item is None: + continue for _l in range(day_item.childCount()): file_item = day_item.child(_l) + if file_item is None: + continue if get_base_run_filename(file_item.text(0)) == run: self.history_tree_widget.setCurrentItem(file_item) return - def currentText(self): + def currentText(self) -> str: current_item = self.history_tree_widget.currentItem() if current_item: return current_item.text(0) return "" - def count(self): + def count(self) -> int: if self.runs is None: return 0 return len(self.runs) - def show_context_menu_history(self, position): + def show_context_menu_history(self, position: QtCore.QPoint) -> None: """ Executed when items in the tree are right-clicked on, and then calls the base-class function to actually display the context-menu. @@ -236,7 +286,7 @@ def show_context_menu_history(self, position): self.show_context_menu(self.history_tree_widget, fullpath) - def find_run_by_name(self, runs, filename): + def find_run_by_name(self, runs: list[str], filename: str) -> str | None: """ Search in run_list (full paths) for a file matching filename. Returns the full path if found, else None. @@ -252,7 +302,7 @@ class TemplateNavigator(QWidget, FileContextMenuBase): Navigate through the templates """ - def __init__(self, parent=None): + def __init__(self, parent: QWidget | None = None) -> None: super().__init__(parent) # get template directory @@ -270,7 +320,9 @@ def __init__(self, parent=None): self.template_tree_view = QTreeView(self) layout.addWidget(self.template_tree_view) - self.template_tree_view.setContextMenuPolicy(Qt.CustomContextMenu) + self.template_tree_view.setContextMenuPolicy( + Qt.ContextMenuPolicy.CustomContextMenu + ) self.template_tree_view.customContextMenuRequested.connect( self.show_context_menu_template ) @@ -303,7 +355,7 @@ def __init__(self, parent=None): # Expand the root for quick glance self.template_tree_view.expand(self.file_sys_model.index(self.template_dir)) - def show_context_menu_template(self, position): + def show_context_menu_template(self, position: QtCore.QPoint) -> None: """ Executed when items in the tree are right-clicked on, and then calls the base-class function to actually display the context-menu. @@ -312,7 +364,7 @@ def show_context_menu_template(self, position): if not selected_index.isValid(): return # user didn't click on any menu item - fullpath = self.template_tree_view.model().filePath(selected_index) + fullpath = self.file_sys_model.filePath(selected_index) filename = os.path.basename(fullpath) if not filename.endswith( (".yaml", ".yml") diff --git a/src/badger/gui/components/obj_table.py b/src/badger/gui/components/obj_table.py index 107212ab..75b35518 100644 --- a/src/badger/gui/components/obj_table.py +++ b/src/badger/gui/components/obj_table.py @@ -43,18 +43,18 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: self.setColumnWidth(2, 192) self.setHorizontalHeaderLabels(["", "Name", "Rule"]) - def default_info(self) -> list[str]: + def default_info(self) -> tuple[str, float, bool]: """ Get the default information list for a new item. Returns ------- - list - A list containing default values for a new item. + tuple + A tuple containing default values for a new item. """ - return ["MINIMIZE"] + return ("MINIMIZE", 0.0, False) - def new_item_prompt(self): + def new_item_prompt(self) -> str: """ The prompt text to enter a new item. """ @@ -75,7 +75,7 @@ def heads_up(self, name: str) -> None: f"Objective {name} already exists!", ) - def create_cell_widgets(self, info: list[Any]) -> tuple[QWidget, ...]: + def create_cell_widgets(self, info: tuple[str, float, bool]) -> tuple[QWidget, ...]: # Rule cb_rule = QComboBox() cb_rule.setItemDelegate(QStyledItemDelegate()) @@ -83,5 +83,3 @@ def create_cell_widgets(self, info: list[Any]) -> tuple[QWidget, ...]: cb_rule.setCurrentIndex(0 if info[0] == "MINIMIZE" else 1) return (cb_rule,) - - return [cb_rule] diff --git a/src/badger/gui/components/obs_table.py b/src/badger/gui/components/obs_table.py index 8c8730c5..e26f68c2 100644 --- a/src/badger/gui/components/obs_table.py +++ b/src/badger/gui/components/obs_table.py @@ -39,18 +39,18 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: self.setHorizontalHeaderLabels(["", "Name"]) - def default_info(self) -> list[Any]: + def default_info(self) -> tuple[str, float, bool]: """ Get the default information list for a new item. Returns ------- - list - A list containing default values for a new item. + tuple + A tuple containing default values for a new item. """ - return [] + return ("", 0.0, False) - def new_item_prompt(self): + def new_item_prompt(self) -> str: """ The prompt text to enter a new item. """ @@ -71,5 +71,5 @@ def heads_up(self, name: str) -> None: f"Observable {name} already exists!", ) - def create_cell_widgets(self, info: list[Any]) -> tuple[Any, ...]: + def create_cell_widgets(self, info: tuple[str, float, bool]) -> tuple[Any, ...]: return () diff --git a/src/badger/gui/components/pf_viewer/pf_widget.py b/src/badger/gui/components/pf_viewer/pf_widget.py index 0af876bd..34e6e15d 100644 --- a/src/badger/gui/components/pf_viewer/pf_widget.py +++ b/src/badger/gui/components/pf_viewer/pf_widget.py @@ -65,9 +65,9 @@ } -class ParetoFrontWidget(AnalysisWidget): - generator: MOBOGenerator # type: ignore - parameters: ConfigurableOptions = DEFAULT_PARAMETERS # type: ignore +class ParetoFrontWidget(AnalysisWidget[ConfigurableOptions]): + generator: MOBOGenerator + parameters: ConfigurableOptions = DEFAULT_PARAMETERS.copy() hypervolume_history: pd.DataFrame = pd.DataFrame() pf_1: Tensor | None = None @@ -318,8 +318,8 @@ def update_pareto_front_plot( with MatplotlibFigureContext(fig_size=self.plot_size) as (fig, ax): try: fig, ax = self.create_pareto_plot(fig, ax) - canvas = FigureCanvas(fig) - toolbar = NavigationToolbar(canvas, self) + canvas = FigureCanvas(fig) # type: ignore[no-untyped-call] + toolbar = NavigationToolbar(canvas, self) # type: ignore[no-untyped-call] variables = self.parameters["variables"] @@ -344,14 +344,14 @@ def update_pareto_front_plot( plot_tab_widget.addTab(widget, "Variable Space") except ValueError: logger.error("No data points available for Variable Space") - blank_canvas = FigureCanvas(fig) + blank_canvas = FigureCanvas(fig) # type: ignore[no-untyped-call] plot_tab_widget.addTab(blank_canvas, "Variable Space") with MatplotlibFigureContext(fig_size=self.plot_size) as (fig, ax): try: fig, ax = self.create_pareto_plot(fig, ax) - canvas = FigureCanvas(fig) - toolbar = NavigationToolbar(canvas, self) + canvas = FigureCanvas(fig) # type: ignore[no-untyped-call] + toolbar = NavigationToolbar(canvas, self) # type: ignore[no-untyped-call] variables = self.parameters["objectives"] @@ -376,7 +376,7 @@ def update_pareto_front_plot( plot_tab_widget.addTab(widget, "Objective Space") except ValueError: logger.error("No data points available for Objective Space") - blank_canvas = FigureCanvas(fig) + blank_canvas = FigureCanvas(fig) # type: ignore[no-untyped-call] plot_tab_widget.addTab(blank_canvas, "Objective Space") plot_tab_widget.setCurrentIndex(self.parameters["plot_tab"]) @@ -397,11 +397,11 @@ def update_hypervolume_plot( with MatplotlibFigureContext(fig_size=self.plot_size) as (fig, ax): try: fig, ax = self.create_hypervolume_plot(fig, ax) - canvas = FigureCanvas(fig) + canvas = FigureCanvas(fig) # type: ignore[no-untyped-call] plot_hypervolume.addWidget(canvas) except ValueError: logger.error("No data points available for Hypervolume") - blank_canvas = FigureCanvas(fig) + blank_canvas = FigureCanvas(fig) # type: ignore[no-untyped-call] plot_hypervolume.addWidget(blank_canvas) def create_pareto_plot(self, fig: Figure, ax: Axes) -> tuple[Figure, Axes]: @@ -514,8 +514,8 @@ def create_hypervolume_plot(self, fig: Figure, ax: Axes) -> tuple[Figure, Axes]: ValueError, "No data points available for Hypervolume" ) # Extract the x and y coordinates from the data points - x = data_points["iteration"].values - y = data_points["hypervolume"].values + x = data_points["iteration"].to_numpy(dtype=float) + y = data_points["hypervolume"].to_numpy(dtype=float) # Create a scatter plot ax.plot( x, diff --git a/src/badger/gui/components/pf_viewer/types.py b/src/badger/gui/components/pf_viewer/types.py index 94f127c5..2b30f388 100644 --- a/src/badger/gui/components/pf_viewer/types.py +++ b/src/badger/gui/components/pf_viewer/types.py @@ -3,17 +3,66 @@ from typing import TypedDict +from PyQt5.QtWidgets import ( + QCheckBox, + QComboBox, + QGridLayout, + QHBoxLayout, + QPushButton, + QTabWidget, + QVBoxLayout, +) + +from badger.gui.components.types import InteractionParameters + class PlotOptions(TypedDict): show_only_pareto_front: bool -class ConfigurableOptions(TypedDict): +class ConfigurableOptions(InteractionParameters): plot_options: PlotOptions - variable_1: int - variable_2: int - variables: list[str] objectives: list[str] objective_1: int objective_2: int plot_tab: int + + +class PFOptionsUIWidgets(TypedDict): + show_only_pareto_front: QCheckBox + + +class PFVariablesUIWidgets(TypedDict): + variable_1: QComboBox + variable_2: QComboBox + + +class PFPlotUIWidgets(TypedDict): + pareto: QTabWidget + hypervolume: QVBoxLayout + + +class PFUIWidgets(TypedDict): + variables: PFVariablesUIWidgets + options: PFOptionsUIWidgets + update: QPushButton + plot: PFPlotUIWidgets + + +class PFVariablesLayouts(TypedDict): + variable_1: QVBoxLayout + variable_2: QVBoxLayout + + +class PFUILayouts(TypedDict): + main: QHBoxLayout + settings: QVBoxLayout + plot: QGridLayout + options: QVBoxLayout + variables: QVBoxLayout + update: QVBoxLayout + + +class PFUI(TypedDict): + components: PFUIWidgets + layouts: PFUILayouts diff --git a/src/badger/gui/components/plot_event_handlers.py b/src/badger/gui/components/plot_event_handlers.py index 73ac038b..e5a1cad4 100644 --- a/src/badger/gui/components/plot_event_handlers.py +++ b/src/badger/gui/components/plot_event_handlers.py @@ -2,20 +2,21 @@ Handles click, hover, scroll-zoom, and data-point annotation tooltips.""" import logging +from collections.abc import Callable from typing import cast +import numpy as np import pandas as pd from matplotlib.backend_bases import MouseButton, MouseEvent, PickEvent from matplotlib.backends.backend_qtagg import FigureCanvasQTAgg from matplotlib.collections import PathCollection from matplotlib.text import Annotation -from pyparsing import Callable from badger.gui.components.extension_utilities import ( HandledException, to_precision_float, ) -from badger.gui.components.types import ConfigurableOptions +from badger.gui.components.types import InteractionParameters from badger.routine import Routine logger = logging.getLogger(__name__) @@ -27,16 +28,10 @@ class MatplotlibInteractionHandler: This class is designed to be used with matplotlib figures and axes. """ - parameters: ConfigurableOptions - routine: Routine - callback: Callable[[Routine, bool], None] - moving: bool - start: dict[str, float] # Starting coordinates for movement - def __init__( self, canvas: FigureCanvasQTAgg, - parameters: ConfigurableOptions, + parameters: InteractionParameters, routine: Routine, variables: list[str], callback: Callable[[Routine, bool], None], @@ -45,46 +40,25 @@ def __init__( self.parameters = parameters self.variables = variables self.routine = routine - self.callback = callback + self.callback: Callable[[Routine, bool], None] = callback self.moving = False - self.start = {"x": 0, "y": 0} # Starting coordinates for movement + self.start: dict[str, float] = { + "x": 0, + "y": 0, + } # Starting coordinates for movement self.step = 0 self.tooltips: list[Annotation] = [] def connect_events(self) -> None: - self.canvas.mpl_connect( - "button_press_event", - lambda event: self.on_click( - event, # type: ignore[call-arg] - self.parameters, - self.routine, - self.callback, - ), - ) - # self.canvas.mpl_connect( - # "button_release_event", - # lambda event: self.on_release( - # event, # type: ignore[call-arg] - # ), - # ) - # self.canvas.mpl_connect( - # "motion_notify_event", - # lambda event: self.on_motion(event), # type: ignore[call-arg] - # ) - - self.canvas.mpl_connect( - "scroll_event", - lambda event: self.on_scroll(event), # type: ignore[call-arg] - ) - - self.canvas.mpl_connect( - "pick_event", - lambda event: self.on_pick(event), # type: ignore[call-arg] - ) + self.canvas.mpl_connect("button_press_event", self.on_click) + # self.canvas.mpl_connect("button_release_event", self.on_release) + # self.canvas.mpl_connect("motion_notify_event", self.on_motion) + self.canvas.mpl_connect("scroll_event", self.on_scroll) + self.canvas.mpl_connect("pick_event", self.on_pick) def update_reference_points( self, - parameters: ConfigurableOptions, + parameters: InteractionParameters, desired_coordinate: tuple[float, float], ) -> None: if ( @@ -119,13 +93,10 @@ def update_reference_points( logger.debug(f"Updated reference points: {parameters['reference_points']}") - def on_click( - self, - event: MouseEvent, - parameters: ConfigurableOptions, - routine: Routine, - callback: Callable[[Routine, bool], None], - ) -> None: + def on_click(self, event: MouseEvent) -> None: + parameters = self.parameters + routine = self.routine + callback = self.callback logger.debug(f"Clicked at {event.xdata}, {event.ydata}, button: {event.button}") if event.inaxes is None: logger.debug("Click outside axes, ignoring") @@ -267,7 +238,7 @@ def on_scroll(self, event: MouseEvent) -> None: def on_pick(self, event: PickEvent) -> None: logger.debug("on_pick event triggered") plot = event.artist - mouseevent = event.mouseevent + mouseevent = cast(MouseEvent, event.mouseevent) # type: ignore[redundant-cast] # mypy does not recognize MouseEvent from PickEvent if mouseevent.inaxes is None: logger.debug("Mouse event outside axes, ignoring") return @@ -287,18 +258,20 @@ def on_pick(self, event: PickEvent) -> None: if isinstance(plot, PathCollection): data = plot.get_offsets() - indexes = cast(list[int], event.ind) + indexes = cast(list[int], getattr(event, "ind", [])) if len(indexes) == 0: logger.debug("No indices in pick event, ignoring") return index = indexes[0] - if index < 0 or index >= len(data): - logger.debug(f"Index {index} out of bounds for data length {len(data)}") + if index < 0 or index >= len(np.asarray(data)): + logger.debug( + f"Index {index} out of bounds for data length {len(np.asarray(data))}" + ) return - point = cast(tuple[float, float], data[index]) + point = cast(tuple[float, float], np.asarray(data)[index]) logger.debug(f"Picked point: {point} at index {index}") # Routine data @@ -314,20 +287,17 @@ def on_pick(self, event: PickEvent) -> None: searched_row = pd.Series(dtype=float) if x_column and y_column: - # Find the true index in the routine data - # This is done by finding the row in the routine data that is closest to the picked point - searched_row = routine_data.loc[ # type: ignore - (routine_data[x_column] - point[0]).abs().idxmin() # type: ignore - & (routine_data[y_column] - point[1]).abs().idxmin() - ] + # Find the row in the routine data closest to the picked point in both dimensions + distances = (routine_data[x_column] - point[0]) ** 2 + ( + routine_data[y_column] - point[1] + ) ** 2 + searched_row = routine_data.iloc[int(distances.to_numpy().argmin())] elif x_column: - searched_row = routine_data.loc[ - (routine_data[x_column] - point[0]).abs().idxmin() - ] + distances = (routine_data[x_column] - point[0]).abs() + searched_row = routine_data.iloc[int(distances.to_numpy().argmin())] elif y_column: - searched_row = routine_data.loc[ - (routine_data[y_column] - point[1]).abs().idxmin() - ] + distances = (routine_data[y_column] - point[1]).abs() + searched_row = routine_data.iloc[int(distances.to_numpy().argmin())] true_index = searched_row.name @@ -375,7 +345,7 @@ def on_pick(self, event: PickEvent) -> None: event.canvas.draw_idle() return - def clear_tooltips(self): + def clear_tooltips(self) -> None: """ Clear all tooltips from the plot. This is useful to remove any existing tooltips before a new plot is drawn. @@ -384,4 +354,4 @@ def clear_tooltips(self): for tooltip in self.tooltips: tooltip.remove() self.tooltips.clear() - self.canvas.draw_idle() + self.canvas.draw_idle() # type: ignore[no-untyped-call] diff --git a/src/badger/gui/components/process_manager.py b/src/badger/gui/components/process_manager.py index 8beee914..551e05d6 100644 --- a/src/badger/gui/components/process_manager.py +++ b/src/badger/gui/components/process_manager.py @@ -1,8 +1,12 @@ """Keeps a pool of pre-spawned subprocesses ready to run optimizations. When one is consumed by a run, signals that a new one should be created.""" +from __future__ import annotations + from PyQt5.QtCore import QObject, pyqtSignal +from badger.types import ProcessWithArgs + class ProcessManager(QObject): """ @@ -14,26 +18,26 @@ class ProcessManager(QObject): def __init__(self) -> None: super().__init__() - self.processes_queue = [] + self.processes_queue: list[ProcessWithArgs] = [] - def add_to_queue(self, process_with_args: dict) -> None: + def add_to_queue(self, process_with_args: ProcessWithArgs) -> None: """ - Add to a dict contaitng a process and it's coresponding args to the processes_queue. + Add to a ProcessWithArgs object containing a process and its corresponding args to the processes_queue. Parameters ---------- - process_with_args: dict + process_with_args: ProcessWithArgs """ self.processes_queue.append(process_with_args) - def remove_from_queue(self) -> dict | None: + def remove_from_queue(self) -> ProcessWithArgs | None: """ - Removes and returns a process and it's coresponding args to the processes_queue. - If no process are in the processes_queue then the method returns None. + Removes and returns a ProcessWithArgs object from the processes_queue. + If no ProcessWithArgs objects are in the processes_queue then the method returns None. Returns ------- - process_with_args: dict | None + process_with_args: ProcessWithArgs | None """ if self.processes_queue: process_with_args = self.processes_queue.pop(0) @@ -42,17 +46,17 @@ def remove_from_queue(self) -> dict | None: return None - def close_proccesses(self) -> bool: + def close_processes(self) -> bool: """ - Closes the processes stored in the processes_queue. + Closes the ProcessWithArgs objects stored in the processes_queue. Returns ------- True: bool """ for i in range(len(self.processes_queue)): - process = self.processes_queue.pop(0) - process["process"].terminate() - process["process"].join() + _p = self.processes_queue.pop(0) + _p.process.terminate() + _p.process.join() return True diff --git a/src/badger/gui/components/pydantic_editor.py b/src/badger/gui/components/pydantic_editor.py index dd29ba52..0676644c 100644 --- a/src/badger/gui/components/pydantic_editor.py +++ b/src/badger/gui/components/pydantic_editor.py @@ -13,7 +13,7 @@ from collections.abc import Callable, Sequence from dataclasses import dataclass from inspect import isclass -from types import NoneType +from types import NoneType, UnionType from typing import ( Annotated, Any, @@ -29,6 +29,7 @@ import yaml from bax_algorithms.emittance import PathwiseMinimizeEmittance from bax_algorithms.solenoid_alignment import PathwiseSolenoidAlignment +from gest_api.vocs import VOCS from pydantic import BaseModel, Field, ValidationError, create_model from pydantic.fields import FieldInfo from pydantic_core import PydanticUndefined, PydanticUndefinedType @@ -57,7 +58,6 @@ from xopt.generators.bayesian.bayesian_generator import BayesianGenerator from xopt.generators.bayesian.turbo import TurboController from xopt.numerical_optimizer import NumericalOptimizer -from xopt.vocs import VOCS logger = logging.getLogger(__name__) @@ -94,9 +94,11 @@ def convert_to_type(value: Any, type: Callable[[Any], T]) -> T: def _set_value_for_basic_widget( - widget: QWidget, + widget: QWidget | None, value: str | float | bool | None, ) -> None: + if widget is None: + return nullable = bool(widget.property("badger_nullable")) if isinstance(widget, (QLabel, QLineEdit)): widget.setText("null" if value is None else str(value)) @@ -164,7 +166,8 @@ def resolve(cls, annotation: type[Any] | Any | None) -> "BadgerResolvedType": if origin == Annotated: return BadgerResolvedType.resolve(args[0]) - if origin == Union: + # `X | None` (PEP 604) resolves to types.UnionType, `Union[X, None]` to typing.Union. + if origin is Union or origin is UnionType: if NoneType in args: origin = Optional args = tuple(arg for arg in args if arg != NoneType) @@ -181,7 +184,7 @@ def resolve(cls, annotation: type[Any] | Any | None) -> "BadgerResolvedType": elif len(args) > 1: primary = BadgerResolvedType.find_primary(args) return BadgerResolvedType( - main=origin, + main=cast("type[Any] | None", origin), nullable=nullable, subtype=primary, ) @@ -206,7 +209,7 @@ def resolve(cls, annotation: type[Any] | Any | None) -> "BadgerResolvedType": return BadgerResolvedType(main=str) return BadgerResolvedType( - main=origin, + main=cast("type[Any] | None", origin), nullable=nullable, subtype=BadgerResolvedType.find_primary(args), ) @@ -284,19 +287,20 @@ def resolve_qt( ) if isinstance(resolved_type.subtype, list): primary_type = resolved_type.subtype[0] - secondary_type = ( + list_secondary_type: BadgerResolvedType | None = ( resolved_type.subtype[1] if len(resolved_type.subtype) > 1 else None ) else: primary_type = resolved_type.subtype - secondary_type = None + list_secondary_type = None if primary_type.main is None: raise ValueError( f"Property name {property_name}: List subtype must be a basic type" ) widget = BadgerListEditor( - primary_type.main, secondary_type.main if secondary_type else None + primary_type.main, + list_secondary_type.main if list_secondary_type else None, ) if default is not None and isinstance(default, list): @@ -520,7 +524,9 @@ def _default_for_new_row( class BadgerListItem(QWidget): - def __init__(self, editor: "BadgerListEditor", parent: QWidget | None = None): + def __init__( + self, editor: "BadgerListEditor", parent: QWidget | None = None + ) -> None: super().__init__(parent) self.editor = editor layout = QHBoxLayout(self) @@ -582,7 +588,7 @@ def __init__( widget_type: type[Any], widget_type2: type[Any] | None = None, parent: QWidget | None = None, - ): + ) -> None: super().__init__(parent) self.widget_type = widget_type self.widget_type2 = widget_type2 @@ -643,8 +649,8 @@ def get_parameters_yaml(self) -> str | None: + "}" ) - def get_parameters_dict(self) -> dict[str, Any] | None: - child_values: list[str | None] = [ + def get_parameters_dict(self) -> dict[str, Any] | list[Any] | None: + child_values: list[Any] = [ _qt_widget_to_value(child.parameter_value) for child in self.list_container.children() if isinstance(child, BadgerListItem) @@ -660,6 +666,8 @@ def get_parameters_dict(self) -> dict[str, Any] | None: for child in self.list_container.children() if isinstance(child, BadgerListItem) ] + # TODO: handle key in dictionary being None + return {k: v for k, v in zip(child_values, child_values2)} @@ -1201,8 +1209,14 @@ def update_after_validate(self, defaults: dict[str, Any]) -> None: # Rebuilding the tree resets the scrollbars, making the view jump back # to the top on every edit. Capture the current scroll positions so we # can restore them once the tree has been repopulated. - h_scroll = self.horizontalScrollBar().value() - v_scroll = self.verticalScrollBar().value() + h_scroll_value = 0 + h_scroll = self.horizontalScrollBar() + if h_scroll is not None: + h_scroll_value = h_scroll.value() + v_scroll = self.verticalScrollBar() + v_scroll_value = 0 + if v_scroll is not None: + v_scroll_value = v_scroll.value() self.clear() @@ -1242,8 +1256,12 @@ def update_after_validate(self, defaults: dict[str, Any]) -> None: # next event-loop iteration so the restore runs after the tree has laid # out its (re)created items and updated the scrollbar ranges. def restore_scroll() -> None: - self.horizontalScrollBar().setValue(h_scroll) - self.verticalScrollBar().setValue(v_scroll) + h_scroll = self.horizontalScrollBar() + if h_scroll is not None: + h_scroll.setValue(h_scroll_value) + v_scroll = self.verticalScrollBar() + if v_scroll is not None: + v_scroll.setValue(v_scroll_value) QTimer.singleShot(0, restore_scroll) @@ -1362,9 +1380,10 @@ def update_error_styles(self, loc: tuple[int | str, ...], msg: str) -> None: else: error_widget = self if error_widget: - if type(error_widget) is QTreeWidgetItem: - error_widget.setBackground(0, Qt.GlobalColor.red) - widget = self.itemWidget(error_widget, self.value_col) + if isinstance(error_widget, QTreeWidgetItem): + item: QTreeWidgetItem = error_widget + item.setBackground(0, Qt.GlobalColor.red) + widget = self.itemWidget(item, self.value_col) if widget is not None: widget.setProperty("error", True) widget.setStyleSheet('*[error="true"] { border: 2px dashed red }') @@ -1374,9 +1393,8 @@ def update_error_styles(self, loc: tuple[int | str, ...], msg: str) -> None: '*[error="true"] { border: 2px dashed red }' ) - error_widget.setToolTip(0, msg) + item.setToolTip(0, msg) else: - error_widget = cast(QTreeWidget, error_widget) error_widget.setProperty("error", True) error_widget.setStyleSheet('*[error="true"] { border: 2px dashed red }') error_widget.setToolTip(msg) diff --git a/src/badger/gui/components/reorderable_table.py b/src/badger/gui/components/reorderable_table.py index a76be5fd..595c8018 100644 --- a/src/badger/gui/components/reorderable_table.py +++ b/src/badger/gui/components/reorderable_table.py @@ -8,26 +8,52 @@ # https://creativecommons.org/publicdomain/zero/1.0/ # https://creativecommons.org/publicdomain/zero/1.0/legalcode -from PyQt5 import QtGui, QtWidgets +from PyQt5.QtCore import QMimeData, QModelIndex, Qt +from PyQt5.QtGui import QPainter, QStandardItem, QStandardItemModel +from PyQt5.QtWidgets import ( + QHeaderView, + QProxyStyle, + QStyle, + QStyleOption, + QTableView, + QWidget, +) -class MyModel(QtGui.QStandardItemModel): - def dropMimeData(self, data, action, row, col, parent): +class MyModel(QStandardItemModel): + def dropMimeData( + self, + data: QMimeData | None, + action: Qt.DropAction, + row: int, + col: int, + parent: QModelIndex, + ) -> bool: """ Always move the entire row, and don't allow column "shifting" """ return super().dropMimeData(data, action, row, 0, parent) -class MyStyle(QtWidgets.QProxyStyle): - def drawPrimitive(self, element, option, painter, widget=None): +class MyStyle(QProxyStyle): + def drawPrimitive( + self, + element: QStyle.PrimitiveElement, + option: QStyleOption | None, + painter: QPainter | None, + widget: QWidget | None = None, + ) -> None: """ Draw a line across the entire row rather than just the column we're hovering over. This may not always work depending on global style - for instance I think it won't work on OSX. """ - if element == self.PE_IndicatorItemViewItemDrop and not option.rect.isNull(): - option_new = QtWidgets.QStyleOption(option) + if ( + element == QStyle.PrimitiveElement.PE_IndicatorItemViewItemDrop + and option is not None + and not option.rect.isNull() + ): + option_new = QStyleOption(option) option_new.rect.setLeft(0) if widget: option_new.rect.setRight(widget.width()) @@ -35,12 +61,16 @@ def drawPrimitive(self, element, option, painter, widget=None): super().drawPrimitive(element, option, painter, widget) -class MyTableView(QtWidgets.QTableView): - def __init__(self, parent): +class MyTableView(QTableView): + def __init__(self, parent: QWidget | None) -> None: super().__init__(parent) - self.verticalHeader().hide() - self.horizontalHeader().hide() - self.horizontalHeader().setSectionResizeMode(QtWidgets.QHeaderView.Stretch) + vheader = self.verticalHeader() + if vheader: + vheader.hide() + hheader = self.horizontalHeader() + if hheader: + hheader.hide() + hheader.setSectionResizeMode(QHeaderView.Stretch) self.setSelectionBehavior(self.SelectRows) self.setSelectionMode(self.SingleSelection) self.setShowGrid(False) @@ -51,16 +81,16 @@ def __init__(self, parent): self.setStyle(MyStyle()) # Set our custom model - this prevents row "shifting" - self.model = MyModel() - self.setModel(self.model) + self._model = MyModel() + self.setModel(self._model) for idx, data in enumerate(["foo", "bar", "baz"]): - item_1 = QtGui.QStandardItem(f"Item {idx}") + item_1 = QStandardItem(f"Item {idx}") item_1.setEditable(False) item_1.setDropEnabled(False) - item_2 = QtGui.QStandardItem(data) + item_2 = QStandardItem(data) item_2.setEditable(False) item_2.setDropEnabled(False) - self.model.appendRow([item_1, item_2]) + self._model.appendRow([item_1, item_2]) diff --git a/src/badger/gui/components/robust_spinbox.py b/src/badger/gui/components/robust_spinbox.py index 706c1988..bbbdd42c 100644 --- a/src/badger/gui/components/robust_spinbox.py +++ b/src/badger/gui/components/robust_spinbox.py @@ -3,48 +3,25 @@ thresholds.""" from PyQt5.QtCore import Qt -from PyQt5.QtWidgets import QAbstractSpinBox, QDoubleSpinBox +from PyQt5.QtWidgets import QAbstractSpinBox, QDoubleSpinBox, QWidget from badger.gui.utils import MouseWheelWidgetAdjustmentGuard class RobustSpinBox(QDoubleSpinBox): - def __init__(self, *args, **kwargs): - try: - decimals = kwargs["decimals"] - del kwargs["decimals"] - except KeyError: - decimals = 6 - - try: - lb = kwargs["lower_bound"] - del kwargs["lower_bound"] - if lb is None: - lb = -1e3 - except KeyError: - lb = -1e3 - - try: - ub = kwargs["upper_bound"] - del kwargs["upper_bound"] - if ub is None: - ub = 1e3 - except KeyError: - ub = 1e3 - - try: - default_value = kwargs["default_value"] - del kwargs["default_value"] - if default_value is None: - default_value = 0 - except KeyError: - default_value = 0 - - super().__init__(*args, **kwargs) + def __init__( + self, + parent: QWidget | None = None, + decimals: int = 6, + lower_bound: float = -1e3, + upper_bound: float = 1e3, + default_value: float = 0, + ) -> None: + super().__init__(parent) self.setDecimals(decimals) self.setFocusPolicy(Qt.FocusPolicy.StrongFocus) self.installEventFilter(MouseWheelWidgetAdjustmentGuard(self)) - self.setRange(lb, ub) + self.setRange(lower_bound, upper_bound) self.setStepType(QAbstractSpinBox.StepType.AdaptiveDecimalStepType) self.setValue(default_value) diff --git a/src/badger/gui/components/routine_editor.py b/src/badger/gui/components/routine_editor.py index 7482e4df..d7e679d9 100644 --- a/src/badger/gui/components/routine_editor.py +++ b/src/badger/gui/components/routine_editor.py @@ -22,13 +22,13 @@ class BadgerRoutineEditor(QWidget): sig_canceled = pyqtSignal() sig_deleted = pyqtSignal() - def __init__(self): - super().__init__() + def __init__(self, parent: QWidget | None = None) -> None: + super().__init__(parent) self.init_ui() self.config_logic() - def init_ui(self): + def init_ui(self) -> None: vbox = QVBoxLayout(self) vbox.setContentsMargins(0, 0, 0, 0) @@ -79,33 +79,33 @@ def init_ui(self): hbox_action.addWidget(btn_save) vbox.addWidget(action_bar) - def config_logic(self): + def config_logic(self) -> None: self.btn_cancel.clicked.connect(self.cancel_create_routine) self.btn_save.clicked.connect(self.save_routine) - def set_routine(self, routine: Routine, silent=False): + def set_routine(self, routine: Routine, silent: bool = False) -> None: try: self.routine_edit.setText(routine.yaml()) except AttributeError: self.routine_edit.setText("") self.routine_page.refresh_ui(routine, silent=silent) - def edit_routine(self): + def edit_routine(self) -> None: self.stacks.setCurrentIndex(1) - def cancel_create_routine(self): + def cancel_create_routine(self) -> None: self.sig_canceled.emit() - def save_routine(self): + def save_routine(self) -> None: # here save() is not a property/attribute # it's a method that also calls _compose_routine() if self.routine_page.save() == 0: self.sig_saved.emit() - def clear(self): + def clear(self) -> None: self.routine_edit.clear() - def switch_mode(self, mode): + def switch_mode(self, mode: str) -> None: if mode == "regular": self.btn_cancel.hide() elif mode == "new routine": diff --git a/src/badger/gui/components/routine_item.py b/src/badger/gui/components/routine_item.py index 2343b267..751a7827 100644 --- a/src/badger/gui/components/routine_item.py +++ b/src/badger/gui/components/routine_item.py @@ -3,7 +3,7 @@ from datetime import datetime -from PyQt5.QtCore import Qt, pyqtSignal +from PyQt5.QtCore import QEvent, Qt, pyqtSignal from PyQt5.QtGui import QFont from PyQt5.QtWidgets import ( QHBoxLayout, @@ -76,8 +76,15 @@ class BadgerRoutineItem(QWidget): sig_del = pyqtSignal(str) def __init__( - self, id, name, timestamp, environment, env_dict, description="", parent=None - ): + self, + id: str, + name: str, + timestamp: str, + environment: str, + env_dict: dict[str, dict[str, str]], + description: str = "", + parent: QWidget | None = None, + ) -> None: super().__init__(parent) self.activated = False @@ -113,8 +120,8 @@ def __init__( self.init_ui() self.config_logic() - def init_ui(self): - self.setAttribute(Qt.WA_StyledBackground) + def init_ui(self) -> None: + self.setAttribute(Qt.WidgetAttribute.WA_StyledBackground) self.setStyleSheet(self.stylesheet_normal) cool_font = QFont() @@ -144,13 +151,17 @@ def init_ui(self): # Routine tools self.btn_fav = btn_fav = create_button( - "star.png", "Favorite routine", stylesheet_fav, size=None + "star.png", + "Favorite routine", + stylesheet_fav, ) btn_fav.setSizePolicy(QSizePolicy.Preferred, QSizePolicy.Expanding) btn_fav.setFixedWidth(32) btn_fav.hide() # hide it for now self.btn_del = btn_del = create_button( - "trash.png", "Delete routine", stylesheet_del, size=None + "trash.png", + "Delete routine", + stylesheet_del, ) btn_del.setSizePolicy(QSizePolicy.Preferred, QSizePolicy.Expanding) btn_del.setFixedWidth(32) @@ -160,25 +171,25 @@ def init_ui(self): self.update_tooltip() - def config_logic(self): + def config_logic(self) -> None: self.btn_del.clicked.connect(self.delete_routine) # self.btn_fav.clicked.connect(self.favorite_routine) - def activate(self): + def activate(self) -> None: self.activated = True if self.hover: self.setStyleSheet(self.stylesheet_activate_hover) else: self.setStyleSheet(self.stylesheet_activate) - def deactivate(self): + def deactivate(self) -> None: self.activated = False if self.hover: self.setStyleSheet(self.stylesheet_normal_hover) else: self.setStyleSheet(self.stylesheet_normal) - def enterEvent(self, event): + def enterEvent(self, event: QEvent | None) -> None: self.hover = True # self.btn_fav.show() # self.btn_del.show() @@ -187,7 +198,7 @@ def enterEvent(self, event): else: self.setStyleSheet(self.stylesheet_normal_hover) - def leaveEvent(self, event): + def leaveEvent(self, event: QEvent | None) -> None: self.hover = False # self.btn_fav.hide() # self.btn_del.hide() @@ -196,9 +207,9 @@ def leaveEvent(self, event): else: self.setStyleSheet(self.stylesheet_normal) - def delete_routine(self): + def delete_routine(self) -> None: reply = QMessageBox.question( - self.parent(), + self.parentWidget(), "Delete routine", f"Are you sure you want to delete routine {self.name}?", QMessageBox.Yes | QMessageBox.No, @@ -209,13 +220,13 @@ def delete_routine(self): self.sig_del.emit(self.id) - def update_tooltip(self): + def update_tooltip(self) -> None: _timestamp = datetime.fromisoformat(self.timestamp) time_str = _timestamp.strftime("%m/%d/%Y, %H:%M:%S") self.setToolTip( f"name: {self.name}\ncreated at: {time_str}\ndescription:\n{self.description}" ) - def update_description(self, descr): + def update_description(self, descr: str) -> None: self.description = descr self.update_tooltip() diff --git a/src/badger/gui/components/routine_page.py b/src/badger/gui/components/routine_page.py index 024f272e..95ef2255 100644 --- a/src/badger/gui/components/routine_page.py +++ b/src/badger/gui/components/routine_page.py @@ -39,6 +39,7 @@ import yaml from coolname import generate_slug from gest_api.vocs import ( + VOCS, BaseConstraint, BaseObjective, ContinuousVariable, @@ -64,7 +65,6 @@ QVBoxLayout, QWidget, ) -from xopt import VOCS from xopt.generators import ( all_generator_names, get_generator_defaults, @@ -129,7 +129,9 @@ def format_validation_error(e: ValidationError) -> str: return "\n".join(messages) -def extract_constraint_symbol_and_value(constraint: BaseConstraint) -> str: +def extract_constraint_symbol_and_value( + constraint: BaseConstraint, +) -> tuple[str, float]: """ Extract symbol and value from gest-api constraint objects generator standard library [gest-api](https://github.com/campa-consortium/gest-api) @@ -205,7 +207,7 @@ def __init__(self) -> None: # 2: not initialized, 1: apply to all, 0: apply to only visible self.lim_apply_to_vars = 2 - def init_ui(self): + def init_ui(self) -> None: logger.info("Initializing UI for BadgerRoutinePage.") config_singleton = init_settings() @@ -232,7 +234,7 @@ def init_ui(self): edit_save.setPlaceholderText(generate_slug(2)) hbox_name.addWidget(label) hbox_name.addWidget(edit_save, 1) - vbox_meta.addWidget(name, alignment=Qt.AlignTop) + vbox_meta.addWidget(name, alignment=Qt.AlignmentFlag.AlignTop) # Description descr = QWidget() @@ -274,15 +276,15 @@ def init_ui(self): "Save as Template" ) save_template_button.setFixedSize(128, 24) - hbox_name.addWidget(save_template_button, alignment=Qt.AlignRight) - vbox_meta.addWidget(template_button, alignment=Qt.AlignBottom) + hbox_name.addWidget(save_template_button, alignment=Qt.AlignmentFlag.AlignRight) + vbox_meta.addWidget(template_button, alignment=Qt.AlignmentFlag.AlignBottom) template_button.show() # Tags self.cbox_tags = cbox_tags = BadgerFilterBox(title=" Tags") if not strtobool(config_singleton.read_value("BADGER_ENABLE_ADVANCED")): cbox_tags.hide() - vbox_meta.addWidget(cbox_tags, alignment=Qt.AlignTop) + vbox_meta.addWidget(cbox_tags, alignment=Qt.AlignmentFlag.AlignTop) # vbox_meta.addStretch() # vbox.addWidget(group_meta) @@ -342,7 +344,7 @@ def init_ui(self): except KeyError: self.template_dir = os.path.join(self.BADGER_PLUGIN_ROOT, "templates") - def config_logic(self): + def config_logic(self) -> None: logger.info("Configuring logic for BadgerRoutinePage.") self.btn_descr_update.clicked.connect(self.update_description) self.env_box.load_template_button.clicked.connect(self.load_template_yaml) @@ -375,9 +377,7 @@ def config_logic(self): # lambda: logger.debug("Selection changed") # ) # for debugging - def load_template_yaml( - self, checked_state, template_path: str | None = None - ) -> None: + def load_template_yaml(self, template_path: str | None = None) -> None: logger.info("Loading template YAML.") """ Load data from template .yaml into template_dict dictionary. @@ -411,7 +411,7 @@ def load_template_yaml( print(f"Error loading template: {e}") return - def set_options_from_template(self, template_dict: dict[str, Any]): + def set_options_from_template(self, template_dict: dict[str, Any]) -> None: logger.info( f"Setting options from template: {template_dict.get('name', 'unknown')}" ) diff --git a/src/badger/gui/components/routine_runner.py b/src/badger/gui/components/routine_runner.py index 2efbecb7..a9315f47 100644 --- a/src/badger/gui/components/routine_runner.py +++ b/src/badger/gui/components/routine_runner.py @@ -10,6 +10,9 @@ import logging import time import traceback +from multiprocessing import Process, Queue +from multiprocessing.synchronize import Event +from typing import Any, cast import pandas as pd from PyQt5.QtCore import QObject, QTimer, pyqtSignal @@ -18,12 +21,8 @@ from badger.errors import ( MEASUREMENT_ACTION_ABORT, MEASUREMENT_ACTION_RETRY, - MEASUREMENT_ACTION_TYPE, - MEASUREMENT_ERROR_TYPE, TERMINATION_ACTION_CONTINUE, TERMINATION_ACTION_END, - TERMINATION_ACTION_TYPE, - TERMINATION_REACHED_TYPE, BadgerError, BadgerRunTerminated, ) @@ -35,6 +34,15 @@ from badger.routine import Routine, calculate_initial_points, calculate_variable_bounds from badger.settings import init_settings from badger.tests.utils import get_current_vars +from badger.types import ( + ArgumentQueueType, + DataQueueMessage, + MeasurementActionMessage, + MeasurementErrorMessage, + TerminationActionMessage, + TerminationConditionConfig, + TerminationReachedMessage, +) logger = logging.getLogger(__name__) @@ -58,7 +66,7 @@ class BadgerRoutineSubprocess: def __init__( self, process_manager: ProcessManager, - routine: Routine = None, + routine: Routine | None = None, routine_filename: str | None = None, save: bool = False, verbose: int = 2, @@ -68,12 +76,12 @@ def __init__( """ Parameters ---------- - routine: Routine - Defined routine for runner + routine: Routine | None + Defined routine for runner, or None if not yet defined save: bool Flag to enable saving to database verbose: int, default: 2 - Verbostiy level (higher is more output) + Verbosity level (higher is more output) use_full_ts: bool If true use full time stamp info when dumping to database """ @@ -84,33 +92,37 @@ def __init__( self.process_manager = process_manager self.routine = routine self.routine_filename = routine_filename - self.run_filename = None + self.run_filename: str | None = None self.states = None # system states to be saved at start of a run self.save = save self.verbose = verbose self.use_full_ts = use_full_ts - self.termination_condition = ( + self.termination_condition: TerminationConditionConfig | None = ( None # additional option to control the optimization flow ) - self.start_time = None # track the time cost of the run - self.last_dump_time = None # track the time the run data got dumped - self.data_and_error_queue = None - self.stop_event = None - self.pause_event = None - self.routine_process = None + self.start_time: float | None = None # track the time cost of the run + self.last_dump_time: float | None = ( + None # track the time the run data got dumped + ) + self.data_and_error_queue: Queue[DataQueueMessage] | None = None + self.stop_event: Event | None = None + self.pause_event: Event | None = None + self.routine_process: Process | None = None self.is_killed = False self.interval = 100 self.testing = testing self.config_singleton = init_settings() - def set_termination_condition(self, termination_condition: dict) -> None: + def set_termination_condition( + self, termination_condition: TerminationConditionConfig + ) -> None: logger.info(f"Setting termination condition: {termination_condition}") """ Setter method for the termination condition. Parameters ---------- - termination_condition : dict + termination_condition : TerminationConditionConfig """ self.termination_condition = termination_condition @@ -125,6 +137,10 @@ def run(self, run_data_flag: bool = False, init_points_flag: bool = False) -> No self.start_time = time.time() self.last_dump_time = None # reset the timer + if self.routine is None: + logger.error("No routine defined for the runner.") + return + # Patch for converting dtype str to torch object try: dtype = self.routine.generator.turbo_controller.tkwargs["dtype"] @@ -136,7 +152,7 @@ def run(self, run_data_flag: bool = False, init_points_flag: bool = False) -> No except TypeError: pass - if not run_data_flag: + if not run_data_flag and self.routine is not None: self.routine.data = None # reset data # Recalculate the bounds and initial points if asked @@ -144,8 +160,15 @@ def run(self, run_data_flag: bool = False, init_points_flag: bool = False) -> No self.config_singleton.read_value("AUTO_REFRESH") and self.routine.relative_to_current ): + limit_options = self.routine.vrange_limit_options + init_actions = self.routine.initial_point_actions + if limit_options is None or init_actions is None: + raise BadgerError( + "relative_to_current requires variable range options and " + "initial point actions to be set" + ) variables_updated = calculate_variable_bounds( - self.routine.vrange_limit_options, + limit_options, self.routine.vocs, self.routine.environment, ) @@ -153,43 +176,47 @@ def run(self, run_data_flag: bool = False, init_points_flag: bool = False) -> No self.routine.vocs.variables = variables_updated init_points = calculate_initial_points( - self.routine.initial_point_actions, + init_actions, self.routine.vocs, self.routine.environment, ) try: - init_points = pd.DataFrame(init_points) + init_points_df = pd.DataFrame(init_points) except IndexError: - init_points = pd.DataFrame(init_points, index=[0]) + init_points_df = pd.DataFrame(init_points, index=[0]) - self.routine.initial_points = init_points + self.routine.initial_points = init_points_df try: self.save_init_vars() process_with_args = self.process_manager.remove_from_queue() - self.routine_process = process_with_args["process"] - self.stop_event = process_with_args["stop_event"] - self.pause_event = process_with_args["pause_event"] - self.args_queue = process_with_args["args_queue"] - self.data_and_error_queue = process_with_args["data_queue"] - self.evaluate_queue = process_with_args["evaluate_queue"] - self.wait_event = process_with_args["wait_event"] - self.dialog_action_queue = process_with_args["dialog_action_queue"] - - arg_dict = { - "routine_id": self.routine.id, - "routine_filename": self.routine_filename, - "routine_name": self.routine.name, - "variable_ranges": self.routine.vocs.variables, - "initial_points": self.routine.initial_points, - "evaluate": True, - "archive": self.save, - "termination_condition": self.termination_condition, - "start_time": self.start_time, - "testing": self.testing, - "run_data": run_data_flag, - "init_points": init_points_flag, - } + if process_with_args is None: + logger.error("No available process in the process manager.") + return + + self.routine_process = process_with_args.process + self.stop_event = process_with_args.stop_event + self.pause_event = process_with_args.pause_event + self.args_queue = process_with_args.args_queue + self.data_and_error_queue = process_with_args.data_queue + self.evaluate_queue = process_with_args.evaluate_queue + self.wait_event = process_with_args.wait_event + self.dialog_action_queue = process_with_args.dialog_action_queue + + arg_dict = ArgumentQueueType( + routine_id=self.routine.id, + routine_filename=self.routine_filename, + routine_name=self.routine.name, + variable_ranges=self.routine.vocs.variables, + initial_points=self.routine.initial_points, + evaluate=True, + archive=self.save, + termination_condition=self.termination_condition, + start_time=self.start_time, + testing=self.testing, + run_data=run_data_flag, + init_points=init_points_flag, + ) self.args_queue.put(arg_dict) self.wait_event.set() @@ -205,8 +232,15 @@ def run(self, run_data_flag: bool = False, init_points_flag: bool = False) -> No self.config_singleton.read_value("AUTO_REFRESH") and self.routine.relative_to_current ): + limit_options = self.routine.vrange_limit_options + init_actions = self.routine.initial_point_actions + if limit_options is None or init_actions is None: + raise BadgerError( + "relative_to_current requires variable range options and " + "initial point actions to be set" + ) variables_updated = calculate_variable_bounds( - self.routine.vrange_limit_options, + limit_options, self.routine.vocs, self.routine.environment, ) @@ -214,23 +248,24 @@ def run(self, run_data_flag: bool = False, init_points_flag: bool = False) -> No self.routine.vocs.variables = variables_updated init_points = calculate_initial_points( - self.routine.initial_point_actions, + init_actions, self.routine.vocs, self.routine.environment, ) try: - init_points = pd.DataFrame(init_points) + init_points_df = pd.DataFrame(init_points) except IndexError: - init_points = pd.DataFrame(init_points, index=[0]) + init_points_df = pd.DataFrame(init_points, index=[0]) - self.routine.initial_points = init_points + self.routine.initial_points = init_points_df except BadgerRunTerminated as e: self.signals.finished.emit() self.signals.info.emit(str(e)) except Exception as e: # noqa: BLE001 - run worker boundary traceback_info = traceback.format_exc() - e._details = traceback_info + # Attach traceback as a dynamic attribute read back by the run monitor. + cast(Any, e)._details = traceback_info self.signals.finished.emit() self.signals.error.emit(e) @@ -252,6 +287,12 @@ def check_queue(self) -> None: It also checks the self.data_and_error_queue to see if an exception was thrown during the routine. It is called by a QTimer every 100 miliseconds. """ + if ( + self.routine is None + or self.data_and_error_queue is None + or self.routine_process is None + ): + return if self.evaluate_queue[1].poll(): while self.evaluate_queue[1].poll(): results = self.evaluate_queue[1].recv() @@ -261,28 +302,18 @@ def check_queue(self) -> None: if not self.data_and_error_queue.empty(): try: msg = self.data_and_error_queue.get() - if isinstance(msg, dict) and msg.get("type") == MEASUREMENT_ERROR_TYPE: + if isinstance(msg, MeasurementErrorMessage): action = self.handle_measurement_error(msg) self.dialog_action_queue.put( - { - "type": MEASUREMENT_ACTION_TYPE, - "action": action, - } + MeasurementActionMessage(action=action) ) - elif ( - isinstance(msg, dict) - and msg.get("type") == TERMINATION_REACHED_TYPE - ): + elif isinstance(msg, TerminationReachedMessage): action = self.handle_termination_reached(msg) self.dialog_action_queue.put( - { - "type": TERMINATION_ACTION_TYPE, - "action": action, - } + TerminationActionMessage(action=action) ) else: - error_title, error_traceback = msg - BadgerError(error_title, error_traceback) + BadgerError(msg.title, msg.traceback) except ValueError: # seems to only occur in tests pass @@ -290,34 +321,36 @@ def check_queue(self) -> None: self.close() self.evaluate_queue[1].close() - def handle_measurement_error(self, msg: dict) -> str: + def handle_measurement_error(self, msg: MeasurementErrorMessage) -> str: dialog = BadgerMeasurementRetryDialog( - text=msg.get("title", "Measurement failed."), - detailedText=msg.get("traceback", ""), + text=msg.title, + detailedText=msg.traceback, ) result = dialog.exec_() - if result == QDialog.Accepted: + if result == QDialog.DialogCode.Accepted: return MEASUREMENT_ACTION_RETRY return MEASUREMENT_ACTION_ABORT - def handle_termination_reached(self, msg: dict) -> str: + def handle_termination_reached(self, msg: TerminationReachedMessage) -> str: # update status - tc_condition = msg.get("tc_condition") + tc_condition = msg.tc_condition status_str = self._format_tc_status_str(tc_condition) self.signals.sig_status.emit(status_str) # launch dialog dialog = BadgerTerminationReachedDialog( tc_condition=tc_condition, - text=msg.get("title"), + text=msg.title, ) result = dialog.exec_() - if result == QDialog.Accepted: - self.signals.sig_status.emit(f"Running routine {self.routine.name}...") + if result == QDialog.DialogCode.Accepted: + if self.routine is not None: + self.signals.sig_status.emit(f"Running routine {self.routine.name}...") return TERMINATION_ACTION_CONTINUE return TERMINATION_ACTION_END - def _format_tc_status_str(self, tc_condition: dict) -> str: + def _format_tc_status_str(self, tc_condition: dict[str, Any]) -> str: + routine_name = self.routine.name if self.routine is not None else "unknown" tc_type = tc_condition["type"] if tc_type == "max_eval": tc_type_text = "N iterations" @@ -325,7 +358,9 @@ def _format_tc_status_str(self, tc_condition: dict) -> str: else: tc_type_text = "timeout" state = f"{tc_condition['state']:.2f} s" - return f"Routine {self.routine.name} paused: Condition {tc_type_text} = {state} reached" + return ( + f"Routine {routine_name} paused: Condition {tc_type_text} = {state} reached" + ) def after_evaluate(self, results: pd.DataFrame) -> None: logger.debug("Received evaluation results from subprocess.") @@ -344,6 +379,8 @@ def save_init_vars(self) -> None: """ Emits the intital variables in the env_ready signal. """ + if self.routine is None: + return init_vars = get_current_vars(self.routine) self.signals.env_ready.emit(init_vars) @@ -355,6 +392,8 @@ def stop_routine(self) -> None: then the method will terminate the process. The method then emits a signal that the process has been stopped. """ + if self.stop_event is None or self.routine_process is None: + return self.stop_event.set() self.routine_process.join(timeout=0.1) @@ -374,6 +413,8 @@ def ctrl_routine(self, pause: bool) -> None: ---------- pause : bool """ + if self.routine is None or self.pause_event is None: + return if pause: self.signals.sig_status.emit(f"Routine {self.routine.name} paused") self.pause_event.clear() diff --git a/src/badger/gui/components/run_monitor.py b/src/badger/gui/components/run_monitor.py index 9a707f3c..f532558c 100644 --- a/src/badger/gui/components/run_monitor.py +++ b/src/badger/gui/components/run_monitor.py @@ -10,12 +10,14 @@ import logging import os import traceback +from collections.abc import Callable from importlib import resources -from typing import TYPE_CHECKING +from typing import Any, Literal, cast import numpy as np import pandas as pd import pyqtgraph as pg +from gest_api.vocs import VOCS from PyQt5.QtCore import pyqtSignal from PyQt5.QtGui import QIcon from PyQt5.QtWidgets import ( @@ -31,11 +33,12 @@ QWidget, ) from pyqtgraph.Qt import QtCore, QtGui -from xopt.vocs import VOCS, normalize_inputs, select_best +from xopt.vocs import normalize_inputs, select_best from badger.archive import BADGER_ARCHIVE_ROOT, archive_run from badger.gui.components.analysis_extensions import AnalysisExtension from badger.gui.components.extensions_palette import ExtensionsPalette +from badger.gui.components.process_manager import ProcessManager from badger.gui.components.pydantic_editor import BadgerPydanticEditor from badger.gui.components.routine_runner import BadgerRoutineSubprocess from badger.gui.windows.message_dialog import BadgerScrollableMessageBox @@ -44,9 +47,7 @@ from badger.logbook import BADGER_LOGBOOK_ROOT, send_to_logbook from badger.routine import Routine from badger.tests.utils import get_current_vars - -if TYPE_CHECKING: - from badger.gui.components.process_manager import ProcessManager +from badger.types import TerminationConditionConfig logger = logging.getLogger(__name__) @@ -73,17 +74,17 @@ class BadgerOptMonitor(QWidget): sig_toggle_other = pyqtSignal(bool) sig_env_ready = pyqtSignal() - def __init__(self, process_manager: "ProcessManager | None" = None): + def __init__(self, process_manager: ProcessManager): super().__init__() # self.setAttribute(Qt.WA_DeleteOnClose, True) # For plot type switching - self.x_plot_y_axis = 0 # 0: raw, 1: normalized - self.plot_x_axis = 0 # 0: iteration, 1: time + self.x_plot_y_axis: Literal[0, 1] = 0 # 0: raw, 1: normalized + self.plot_x_axis: Literal[0, 1] = 0 # 0: iteration, 1: time self.x_plot_relative = True # Routine info - self.routine = None - self.routine_filename = None + self.routine: Routine | None = None + self.routine_filename: str | None = None self.process_manager = process_manager # Curves in the monitor @@ -97,7 +98,7 @@ def __init__(self, process_manager: "ProcessManager | None" = None): self.running = False # Termination condition for the run - self.termination_condition = None + self.termination_condition: TerminationConditionConfig | None = None self.checkpoint_data: dict[str, float] | None = None @@ -106,7 +107,7 @@ def __init__(self, process_manager: "ProcessManager | None" = None): self.testing = False self.tc_dialog = None - self.post_run_actions = [] + self.post_run_actions: list[Callable[[], None]] = [] self.init_ui() self.config_logic() @@ -114,7 +115,7 @@ def __init__(self, process_manager: "ProcessManager | None" = None): @property def vocs(self) -> VOCS: - return self.routine.vocs + return self.routine.vocs if self.routine is not None else None def states(self, new_states: dict) -> None: self._states = new_states @@ -161,7 +162,8 @@ def init_ui(self) -> None: # Set up the monitor # Don't set show=True or there will be a blank window flashing once - self.monitor = monitor = pg.GraphicsLayoutWidget() + # Cast the GraphicsLayoutWidget to GraphicsLayout for type consistency + self.monitor = monitor = cast(pg.GraphicsLayout, pg.GraphicsLayoutWidget()) pg.setConfigOptions(antialias=True) # monitor.ci.setBorder((50, 50, 100)) # monitor.resize(1000, 600) @@ -240,7 +242,7 @@ def config_logic(self) -> None: self.check_relative.stateChanged.connect(self.toggle_x_plot_y_axis_relative) def init_plots( - self, routine: Routine = None, run_filename: str | None = None + self, routine: Routine | None = None, run_filename: str | None = None ) -> None: """ Initialize and configure the plots and related components in the application. @@ -251,11 +253,11 @@ def init_plots( Parameters ---------- - routine : Routine, + routine : Routine | None, The routine to use for configuring the plots. If not provided, the method will use the previously set routine. - run_filename : str, optional + run_filename : str | None, optional The filename of the run, used to determine the state of the application's UI elements. @@ -406,10 +408,12 @@ def init_plots( self.sig_toggle_other.emit(False) - def _configure_plot(self, plot_object, inspector, names: list[str]) -> dict: + def _configure_plot( + self, plot_object: pg.PlotItem, inspector: pg.InfiniteLine, names: list[str] + ) -> pg.PlotDataItem: plot_object.clear() plot_object.addItem(inspector) - curves = {} + curves: dict[str, pg.PlotDataItem] = {} for i, name in enumerate(names): color = self.colors[i % len(self.colors)] # symbol = self.symbols[i % len(self.colors)] @@ -452,7 +456,7 @@ def init_routine_runner(self) -> None: routine_runner.signals.env_ready.connect(self.env_ready) routine_runner.signals.finished.connect(self.routine_finished) - routine_runner.signals.progress.connect(self.update) + routine_runner.signals.progress.connect(self.update_plots) routine_runner.signals.error.connect(self.on_error) routine_runner.signals.info.connect(self.on_info) routine_runner.signals.states.connect(self.states) @@ -488,7 +492,7 @@ def start( self.sig_run_started.emit() self.sig_lock.emit(True) - def save_termination_condition(self, tc) -> None: + def save_termination_condition(self, tc: TerminationConditionConfig) -> None: self.termination_condition = tc def enable_auto_range(self) -> None: @@ -519,7 +523,7 @@ def extract_timestamp(self, data: pd.DataFrame | None = None) -> np.ndarray: return data["timestamp"].to_numpy(copy=True) - def update(self, results: pd.DataFrame) -> None: + def update_plots(self, results: pd.DataFrame) -> None: """Update plots in main window as well as any active extensions and the extensions palette @@ -659,6 +663,9 @@ def routine_finished(self) -> None: try: # TODO: fill in the states # TODO: replace self.testing with with another processes for having a testing mode + if self.routine is None: + raise ValueError("Routine is not set") + if not self.testing: run = archive_run(self.routine, states=self._states) self.routine_runner.run_filename = run["filename"] @@ -666,15 +673,15 @@ def routine_finished(self) -> None: path = run["path"] filename = run["filename"][:-4] + "pickle" - try: - env.interface.stop_recording(os.path.join(path, filename)) - except AttributeError: # recording was not enabled - logger.debug("Recording was not enabled") + try: + env.interface.stop_recording(os.path.join(path, filename)) + except AttributeError: # recording was not enabled + logger.debug("Recording was not enabled") - self.sig_run_name.emit(run["filename"]) - self.sig_status.emit( - f"Archive success: Run data archived to {BADGER_ARCHIVE_ROOT}" - ) + self.sig_run_name.emit(run["filename"]) + self.sig_status.emit( + f"Archive success: Run data archived to {BADGER_ARCHIVE_ROOT}" + ) # if not self.testing: # QMessageBox.information( # self, 'Success!', @@ -694,6 +701,9 @@ def routine_finished(self) -> None: def destroy_unused_env(self) -> None: if not self.running: + if self.routine_runner is None: + logger.debug("Routine runner is not set") + return try: del self.routine_runner.routine.environment except AttributeError: # env already destroyed @@ -1073,7 +1083,13 @@ def register_post_run_action(self, action) -> None: self.post_run_actions.append(action) -def add_axes(monitor, ylabel, title, cursor_line, **kwargs) -> pg.PlotItem: +def add_axes( + monitor: pg.GraphicsLayout, + ylabel: str, + title: str, + cursor_line: pg.InfiniteLine, + **kwargs: Any, +) -> pg.PlotItem: plot_obj = monitor.addPlot(title=title, **kwargs) plot_obj.setLabel("left", ylabel) plot_obj.setLabel("bottom", "iterations") @@ -1100,7 +1116,13 @@ def create_cursor_line() -> pg.InfiniteLine: ) -def set_data(names: list[str], curves: dict, data: pd.DataFrame, ts=None) -> None: +def set_data( + names: list[str], + curves: dict[str, pg.PlotDataItem], + data: pd.DataFrame, + ts=None, +) -> None: + # Split data into live and not live live_mask = data["live"].astype(bool) live_data = data.loc[live_mask] diff --git a/src/badger/gui/components/search_bar.py b/src/badger/gui/components/search_bar.py index cc2f66b5..008b9c3d 100644 --- a/src/badger/gui/components/search_bar.py +++ b/src/badger/gui/components/search_bar.py @@ -4,7 +4,7 @@ from PyQt5.QtWidgets import QLineEdit -def search_bar(): +def search_bar() -> QLineEdit: # completer = QCompleter(word_list) line_edit = QLineEdit() diff --git a/src/badger/gui/components/state_item.py b/src/badger/gui/components/state_item.py index f0901a34..3a536a9b 100644 --- a/src/badger/gui/components/state_item.py +++ b/src/badger/gui/components/state_item.py @@ -1,12 +1,15 @@ """Creates a single observable-state row: a combo box for the observable name and a remove button.""" +from collections.abc import Callable + +from PyQt5.QtCore import QEvent from PyQt5.QtWidgets import QHBoxLayout, QPushButton, QStyledItemDelegate, QWidget from badger.gui.utils import NoHoverFocusComboBox -def state_item(options, remove_item, name=None): +def state_item(options: list[str], remove_item: Callable, name: str = "") -> QWidget: widget = QWidget() hbox = QHBoxLayout(widget) hbox.setContentsMargins(2, 2, 2, 2) @@ -30,10 +33,10 @@ def state_item(options, remove_item, name=None): btn_del.clicked.connect(remove_item) - def show_button(event): + def show_button(event: QEvent | None) -> None: btn_del.show() - def hide_button(event): + def hide_button(event: QEvent | None) -> None: btn_del.hide() widget.enterEvent = show_button diff --git a/src/badger/gui/components/status_bar.py b/src/badger/gui/components/status_bar.py index 2c5751a5..f4f45549 100644 --- a/src/badger/gui/components/status_bar.py +++ b/src/badger/gui/components/status_bar.py @@ -11,13 +11,13 @@ class BadgerStatusBar(QWidget): - def __init__(self): + def __init__(self) -> None: super().__init__() self.init_ui() self.config_logic() - def init_ui(self): + def init_ui(self) -> None: icon_ref = resources.files(__package__) / "../images/gear.png" with resources.as_file(icon_ref) as icon_path: self.icon_settings = QIcon(str(icon_path)) @@ -33,7 +33,7 @@ def init_ui(self): self.summary = summary = SimpleElidedLabel() summary.setObjectName("StatusBar") # summary.setStyleSheet("background-color: orange;") # for debugging - summary.setAlignment(Qt.AlignCenter) + summary.setAlignment(Qt.AlignmentFlag.AlignCenter) self.btn_settings = btn_settings = QPushButton() btn_settings.setStyleSheet("background-color: transparent; border: none;") @@ -55,13 +55,13 @@ def init_ui(self): } """) - def config_logic(self): + def config_logic(self) -> None: self.btn_settings.clicked.connect(self.go_settings) - def set_summary(self, text): + def set_summary(self, text: str) -> None: self.summary.setText(text) self.setToolTip(text) - def go_settings(self): + def go_settings(self) -> None: dlg = BadgerSettingsDialog(self) dlg.exec() diff --git a/src/badger/gui/components/syntax.py b/src/badger/gui/components/syntax.py index 1c0c8885..0a959ada 100644 --- a/src/badger/gui/components/syntax.py +++ b/src/badger/gui/components/syntax.py @@ -3,18 +3,25 @@ from typing import ClassVar -from PyQt5 import QtCore, QtGui +from PyQt5.QtCore import QRegExp +from PyQt5.QtGui import ( + QColor, + QFont, + QSyntaxHighlighter, + QTextCharFormat, + QTextDocument, +) -def format(color, style=""): +def format(color: str, style: str = "") -> QTextCharFormat: """Return a QTextCharFormat with the given attributes.""" - _color = QtGui.QColor() + _color = QColor() _color.setNamedColor(color) - _format = QtGui.QTextCharFormat() + _format = QTextCharFormat() _format.setForeground(_color) if "bold" in style: - _format.setFontWeight(QtGui.QFont.Bold) + _format.setFontWeight(QFont.Bold) if "italic" in style: _format.setFontItalic(True) @@ -35,7 +42,7 @@ def format(color, style=""): } -class PythonHighlighter(QtGui.QSyntaxHighlighter): +class PythonHighlighter(QSyntaxHighlighter): """Syntax highlighter for the Python language.""" # Python keywords @@ -118,12 +125,12 @@ class PythonHighlighter(QtGui.QSyntaxHighlighter): "]", ] - def __init__(self, parent: QtGui.QTextDocument) -> None: + def __init__(self, parent: QTextDocument) -> None: super().__init__(parent) # Multi-line strings (expression, flag, style) - self.tri_single = (QtCore.QRegExp("'''"), 1, STYLES["string2"]) - self.tri_double = (QtCore.QRegExp('"""'), 2, STYLES["string2"]) + self.tri_single = (QRegExp("'''"), 1, STYLES["string2"]) + self.tri_double = (QRegExp('"""'), 2, STYLES["string2"]) rules = [] @@ -155,11 +162,11 @@ def __init__(self, parent: QtGui.QTextDocument) -> None: ] # Build a QRegExp for each pattern - self.rules = [(QtCore.QRegExp(pat), index, fmt) for (pat, index, fmt) in rules] + self.rules = [(QRegExp(pat), index, fmt) for (pat, index, fmt) in rules] - def highlightBlock(self, text): + def highlightBlock(self, text: str | None) -> None: """Apply syntax highlighting to the given block of text.""" - self.tripleQuoutesWithinStrings = [] + self.tripleQuoutesWithinStrings: list[int] = [] # Do other syntax formatting for expression, nth, format in self.rules: index = expression.indexIn(text, 0) @@ -193,12 +200,21 @@ def highlightBlock(self, text): self.setCurrentBlockState(0) + if text is None: + return + # Do multi-line strings in_multiline = self.match_multiline(text, *self.tri_single) if not in_multiline: in_multiline = self.match_multiline(text, *self.tri_double) - def match_multiline(self, text, delimiter, in_state, style): + def match_multiline( + self, + text: str, + delimiter: QRegExp, + in_state: int, + style: QTextCharFormat, + ) -> bool: """Do highlighting of multi-line strings. ``delimiter`` should be a ``QRegExp`` for triple-single-quotes or triple-double-quotes, and ``in_state`` should be a unique integer to represent the corresponding diff --git a/src/badger/gui/components/types.py b/src/badger/gui/components/types.py index 9c11268d..fb591d86 100644 --- a/src/badger/gui/components/types.py +++ b/src/badger/gui/components/types.py @@ -1,14 +1,21 @@ -"""TypedDict definitions shared by the analysis extensions — plot options, -variable selections, and reference point settings.""" +"""Shared TypedDict contract for the analysis extensions. -from typing import TypedDict +``InteractionParameters`` is the base set of fields consumed by +``MatplotlibInteractionHandler``. Every extension's own options TypedDict should +inherit from it so the handler works with any extension, current or future.""" +from typing import NotRequired, TypedDict + + +class InteractionParameters(TypedDict): + """Fields the shared plot interaction handler reads and writes. + + ``variable_1``/``variable_2`` index into ``variables`` to identify the plotted + axes. The reference-point fields are only present for extensions that support + reference points, so they are optional.""" -class ConfigurableOptions(TypedDict, total=False): variable_1: int variable_2: int variables: list[str] - reference_points: dict[str, float] - reference_points_range: dict[str, tuple[float, float]] - include_variable_2: bool - plot_options: dict[str, bool | int] + reference_points: NotRequired[dict[str, float]] + reference_points_range: NotRequired[dict[str, tuple[float, float]]] diff --git a/src/badger/gui/components/var_table.py b/src/badger/gui/components/var_table.py index 67247205..4e2159cf 100644 --- a/src/badger/gui/components/var_table.py +++ b/src/badger/gui/components/var_table.py @@ -25,10 +25,10 @@ import traceback from functools import partial from importlib import resources -from typing import Any, cast +from typing import TYPE_CHECKING, Any, cast from gest_api.vocs import ContinuousVariable -from PyQt5.QtCore import QPoint, QSize, Qt, pyqtSignal +from PyQt5.QtCore import QModelIndex, QPoint, QSize, Qt, pyqtSignal from PyQt5.QtGui import ( QColor, QDragEnterEvent, @@ -58,6 +58,9 @@ from badger.gui.components.robust_spinbox import RobustSpinBox from badger.gui.windows.expandable_message_box import ExpandableMessageBox +if TYPE_CHECKING: + from badger.factory import BadgerPluginConfig + logger = logging.getLogger(__name__) @@ -121,26 +124,26 @@ def __init__(self, *args: Any, **kwargs: Any): self.bounds: dict[str, tuple[float, float]] = {} # track var bounds self.checked_only = False self.bounds_locked = False - self.addtl_vars = [] # track variables added on the fly + self.addtl_vars: list[str] = [] # track variables added on the fly self.env_class: type[Environment] | None = ( None # needed to get bounds on the fly ) - self.env = None # needed to get bounds on the fly - self.configs = None # needed to get bounds on the fly + self.env: Environment | None = None # needed to get bounds on the fly + self.configs: BadgerPluginConfig | None = None # for bounds on the fly self.previous_values: dict[ tuple[int, int], str ] = {} # to track changes in table self.config_logic() self.customContextMenuRequested.connect(self.display_context_menu) - def update_vocs(self): + def update_vocs(self) -> None: """ Emit the data_changed signal to notify that the VOCS has been updated. """ logger.debug("Emitting data_changed signal from VariableTable") self.data_changed.emit() - def config_logic(self): + def config_logic(self) -> None: hheader = self.horizontalHeader() if hheader: hheader.sectionClicked.connect(self.header_clicked) @@ -155,7 +158,7 @@ def setItem(self, row: int, column: int, item: QTableWidgetItem | None) -> None: self.previous_values[(row, column)] = item.text() super().setItem(row, column, item) - def is_all_checked(self): + def is_all_checked(self) -> bool: for i in range(self.rowCount() - 1): item = self.cellWidget(i, 0) if not item: @@ -166,7 +169,7 @@ def is_all_checked(self): return True - def header_clicked(self, idx: int): + def header_clicked(self, idx: int) -> None: if idx: return @@ -183,7 +186,7 @@ def header_clicked(self, idx: int): item.blockSignals(False) self.update_selected() - def update_bounds(self): + def update_bounds(self) -> None: for i in range(self.rowCount() - 1): widget = self.item(i, 1) if not widget: @@ -204,7 +207,7 @@ def update_bounds(self): self.validate_row(i) # Validate the row after updating bounds self.data_changed.emit() - def validate_row(self, row: int): + def validate_row(self, row: int) -> None: """ Validate the bounds for a given row. If invalid, apply a red border to the row. """ @@ -226,7 +229,7 @@ def validate_row(self, row: int): def set_bounds( self, variables: dict[str, tuple[float, float]], signal: bool = True - ): + ) -> None: for name, bounds in variables.items(): self.bounds[name] = bounds @@ -237,7 +240,7 @@ def set_bounds( def refresh_variable( self, name: str, bounds: tuple[float, float], hard_bounds: tuple[float, float] - ): + ) -> None: self.bounds[name] = bounds # Update the variable in self.variables @@ -254,7 +257,7 @@ def refresh_variable( self.update_variables(self.variables, 2) - def update_selected(self): + def update_selected(self) -> None: for i in range(self.rowCount() - 1): _cb = self.cellWidget(i, 0) if not _cb: @@ -277,21 +280,21 @@ def update_selected(self): if self.checked_only: self.show_checked_only() - def set_selected(self, variable_names: list[str]): - self.selected: dict[str, bool] = {} + def set_selected(self, variable_names: list[str]) -> None: + self.selected = {} for vname in variable_names: self.selected[vname] = True self.update_variables(self.variables, 3) - def toggle_show_mode(self, checked_only: bool): + def toggle_show_mode(self, checked_only: bool) -> None: self.checked_only = checked_only if checked_only: self.show_checked_only() else: self.show_all() - def show_checked_only(self): + def show_checked_only(self) -> None: checked_variables: list[dict[str, tuple[float, float]]] = [] for var in self.variables: name = next(iter(var)) @@ -299,7 +302,7 @@ def show_checked_only(self): checked_variables.append(var) self.update_variables(checked_variables, 3) - def show_all(self): + def show_all(self) -> None: self.update_variables(self.variables, 3) def is_checked(self, name: str) -> bool: @@ -331,13 +334,13 @@ def _convert_bounds_to_tuple(self, bounds: Any) -> dict[str, tuple[float, float] def _set_spinbox_style( self, spinbox: RobustSpinBox, is_selected: bool, has_error: bool = False - ): + ) -> None: """Set the style of a spinbox based on selection and error status.""" color = "lightgray" if is_selected else "gray" border = "border: 1px solid red;" if has_error else "" spinbox.setStyleSheet(f"QDoubleSpinBox {{ color: {color}; {border} }}") - def _set_spinbox_text_color(self, row: int, is_selected: bool): + def _set_spinbox_text_color(self, row: int, is_selected: bool) -> None: """Set the text color of the spinboxes in a given row based on selection status.""" sb_lower = self.cellWidget(row, 2) sb_upper = self.cellWidget(row, 3) @@ -355,7 +358,7 @@ def update_variables( self, variables: list[dict[str, tuple[float, float] | ContinuousVariable]], filtered: int = 0, - ): + ) -> None: # filtered = 0: completely refresh # filtered = 1: filtered by keyword # filtered = 2: just rerender based on check status @@ -370,7 +373,7 @@ def update_variables( self.variables = self.all_variables[:] # make a copy self.selected = {} self.bounds = {} - self.addtl_vars: list[str] = [] + self.addtl_vars = [] for var in self.variables: name = next(iter(var)) self.bounds[name] = var[name] @@ -405,7 +408,9 @@ def update_variables( item.setForeground(QColor("darkCyan")) else: # Make non-new PVs not editable - item.setFlags(item.flags() & ~Qt.ItemIsEditable) + item.setFlags( + cast(Qt.ItemFlags, item.flags() & ~Qt.ItemFlag.ItemIsEditable) + ) item.setForeground(QColor("lightgray" if _cb.isChecked() else "gray")) self.setItem(i, 1, item) @@ -458,7 +463,7 @@ def update_variables( # Make extra editable row item = QTableWidgetItem(self.PLACEHOLDER_TEXT) - item.setFlags(item.flags() | Qt.ItemIsEditable) + item.setFlags(item.flags() | Qt.ItemFlag.ItemIsEditable) item.setForeground(QColor("gray")) self.setItem(n - 1, 1, item) @@ -469,10 +474,10 @@ def update_variables( self.sig_sel_changed.emit() self.data_changed.emit() - def handle_config_button(self, var_name: str): + def handle_config_button(self, var_name: str) -> None: self.sig_var_config.emit(var_name) - def add_additional_variable(self, item: QTableWidgetItem): + def add_additional_variable(self, item: QTableWidgetItem) -> None: row = item.row() column = item.column() name = item.text() @@ -498,7 +503,7 @@ def add_additional_variable(self, item: QTableWidgetItem): if row == self.rowCount() - 1 and column == 1 and name != self.PLACEHOLDER_TEXT: self.try_insert_variable(name) - def try_insert_variable(self, name: str): + def try_insert_variable(self, name: str) -> None: idx = self.rowCount() - 1 # Check that variables doesn't already exist in table @@ -577,13 +582,15 @@ def try_insert_variable(self, name: str): def get_bounds(self, name: str) -> tuple[Any, Any]: # TODO: move elsewhere? + if self.env_class is None or self.configs is None: + raise RuntimeError("Environment is not configured for fetching bounds!") self.env = instantiate_env(self.env_class, self.configs) value = self.env.get_variable(name) bound = self.env.get_bounds([name])[name] return value, bound - def add_variable(self, name: str, lb: float, ub: float): + def add_variable(self, name: str, lb: float, ub: float) -> None: var = {name: (lb, ub)} self.all_variables.append(var) @@ -599,15 +606,15 @@ def export_variables(self) -> dict[str, tuple[float, float]]: return variables_exported - def lock_bounds(self): + def lock_bounds(self) -> None: self.bounds_locked = True self.toggle_show_mode(self.checked_only) - def unlock_bounds(self): + def unlock_bounds(self) -> None: self.bounds_locked = False self.toggle_show_mode(self.checked_only) - def dragEnterEvent(self, e: QDragEnterEvent | None): + def dragEnterEvent(self, e: QDragEnterEvent | None) -> None: if e is None: return @@ -621,7 +628,7 @@ def dragEnterEvent(self, e: QDragEnterEvent | None): else: e.ignore() - def dragMoveEvent(self, e: QDragMoveEvent | None): + def dragMoveEvent(self, e: QDragMoveEvent | None) -> None: if e is None: return @@ -635,7 +642,7 @@ def dragMoveEvent(self, e: QDragMoveEvent | None): else: e.ignore() - def dropEvent(self, event: QDropEvent | None): + def dropEvent(self, event: QDropEvent | None) -> None: if event is None: return @@ -661,19 +668,23 @@ def dropEvent(self, event: QDropEvent | None): else: event.ignore() - def flags(self, index): + def flags(self, index: QModelIndex) -> Qt.ItemFlags: if not index.isValid(): - return Qt.ItemIsEnabled - flags = Qt.ItemIsSelectable | Qt.ItemIsEnabled + return cast(Qt.ItemFlags, Qt.ItemFlag.ItemIsEnabled) + flags = Qt.ItemFlag.ItemIsSelectable | Qt.ItemFlag.ItemIsEnabled if index.column() == 1: - flags |= Qt.ItemIsEditable | Qt.ItemIsDropEnabled - return flags + flags |= Qt.ItemFlag.ItemIsEditable | Qt.ItemFlag.ItemIsDropEnabled + return cast(Qt.ItemFlags, flags) - def display_info(self, item: QTableWidgetItem | None): + def display_info(self, item: QTableWidgetItem | None) -> None: """ Opens a message box displaying status info from the underlying interface about a variable. """ + if item is None: + return if not self.env: + if self.env_class is None or self.configs is None: + raise RuntimeError("Environment is not configured for fetching info!") self.env = instantiate_env(self.env_class, self.configs) status = self.env.get_info([item.text()]) @@ -684,7 +695,7 @@ def display_info(self, item: QTableWidgetItem | None): mb.setWindowTitle("Variable Info") layout = QGridLayout(mb) row = 0 - for k, v in status["vars"][item.text()].items(): + for k, v in cast(Any, status["vars"])[item.text()].items(): layout.addWidget(QLabel(text=f"{k}:", parent=mb), row, 0) layout.addWidget(QLabel(text=v, parent=mb), row, 1) row += 1 @@ -698,17 +709,17 @@ def display_info(self, item: QTableWidgetItem | None): mb.exec() @staticmethod - def _handle_close(widget: QPushButton): + def _handle_close(widget: QPushButton) -> None: widget.close() - def copy_name(self, item: QTableWidgetItem | None): + def copy_name(self, item: QTableWidgetItem | None) -> None: if item is None: return clipboard = QGuiApplication.clipboard() if clipboard is not None: clipboard.setText(item.text()) - def display_context_menu(self, pt: QPoint): + def display_context_menu(self, pt: QPoint) -> None: menu = QMenu(self) item = self.itemAt(pt) if item is None or item.text() == self.PLACEHOLDER_TEXT or item.column() != 1: diff --git a/src/badger/gui/mini/__init__.py b/src/badger/gui/mini/__init__.py index 797b633a..8074ec59 100644 --- a/src/badger/gui/mini/__init__.py +++ b/src/badger/gui/mini/__init__.py @@ -8,7 +8,7 @@ import traceback from importlib import resources from types import TracebackType -from typing import NoReturn +from typing import Any, NoReturn from PyQt5 import QtCore from PyQt5.QtGui import QFont, QIcon @@ -27,10 +27,10 @@ if hasattr(QtCore.Qt, "AA_UseHighDpiPixmaps"): QApplication.setAttribute(QtCore.Qt.AA_UseHighDpiPixmaps, True) -TIMER = {"time": None} +TIMER: dict[str, float | None] = {"time": None} -def on_exit(*args): +def on_exit(*args: Any) -> None: if TIMER["time"] is None: print("Press Ctrl/Cmd + C again within 1s to quit Badger") TIMER["time"] = time.time() @@ -41,7 +41,7 @@ def on_exit(*args): QApplication.quit() -def on_timeout(): +def on_timeout() -> None: if TIMER["time"] is None: return @@ -52,7 +52,7 @@ def on_timeout(): def error_handler( - etype: type[BaseException], value: BaseException, tb: TracebackType + etype: type[BaseException], value: BaseException, tb: TracebackType | None ) -> NoReturn: """ Custom exception handler that formats uncaught exceptions and raises a BadgerError. @@ -63,8 +63,8 @@ def error_handler( The class of the exception that was raised. value : BaseException The exception instance. - tb : TracebackType - The traceback object associated with the exception. + tb : TracebackType | None + The traceback object associated with the exception, or None if not available. Raises ------ @@ -76,14 +76,17 @@ def error_handler( raise BadgerError(error_title, error_msg) -def launch_gui(config_path=None, template_filename: str | None = None): +def launch_mini_gui( + config_path: str | None = None, template_filename: str | None = None +) -> None: + sys.excepthook = error_handler app = QApplication(sys.argv) # Set app metainfo app.setApplicationName("Badger") - icon_ref = resources.files(__name__).parent / "images" / "icon.png" + icon_ref = resources.files(__name__).joinpath("images", "icon.png") with resources.as_file(icon_ref) as icon_path: app.setWindowIcon(QIcon(str(icon_path))) diff --git a/src/badger/gui/mini/components/var_table.py b/src/badger/gui/mini/components/var_table.py index af885087..acb7e8c3 100644 --- a/src/badger/gui/mini/components/var_table.py +++ b/src/badger/gui/mini/components/var_table.py @@ -68,7 +68,7 @@ def __init__(self, value: float | None): self.label = QLabel(self._format(value)) self.label.setAlignment( - Qt.AlignmentFlag.AlignLeft | Qt.AlignmentFlag.AlignVCenter + Qt.AlignmentFlag(Qt.AlignmentFlag.AlignLeft | Qt.AlignmentFlag.AlignVCenter) ) layout.addWidget(self.label) @@ -83,7 +83,7 @@ def update_style( is_selected: bool, alert_level: int, selected_color: str, - ): + ) -> None: """ Update the value cell visual style. @@ -126,7 +126,9 @@ class SavedValueCell(ValueCell): def __init__(self, value: float | None): super().__init__(value) - def update_style(self, is_selected: bool, alert_level: int): + def update_style( + self, is_selected: bool, alert_level: int, selected_color: str + ) -> None: super().update_style( is_selected, 0 if alert_level <= 1 else 2, self.SELECTED_COLOR ) @@ -142,7 +144,9 @@ class CurrentValueCell(ValueCell): def __init__(self, value: float | None): super().__init__(value) - def update_style(self, is_selected: bool, alert_level: int): + def update_style( + self, is_selected: bool, alert_level: int, selected_color: str + ) -> None: super().update_style(is_selected, alert_level, self.SELECTED_COLOR) diff --git a/src/badger/gui/mini/pages/home_page.py b/src/badger/gui/mini/pages/home_page.py index 350b9d06..7f9cf726 100644 --- a/src/badger/gui/mini/pages/home_page.py +++ b/src/badger/gui/mini/pages/home_page.py @@ -13,6 +13,7 @@ from importlib import resources import numpy as np +from gest_api.vocs import VOCS from pandas import DataFrame from PyQt5.QtCore import QModelIndex, Qt, pyqtSignal from PyQt5.QtGui import QIcon @@ -41,6 +42,7 @@ update_table, ) from badger.gui.components.navigators import TemplateNavigator +from badger.gui.components.process_manager import ProcessManager from badger.gui.components.run_monitor import BadgerOptMonitor from badger.gui.components.status_bar import BadgerStatusBar from badger.gui.mini.pages.routine_page import BadgerRoutinePage @@ -52,6 +54,7 @@ BadgerTerminationConditionDialog, ) from badger.settings import init_settings +from badger.types import TerminationConditionConfig from badger.utils import get_header logger = logging.getLogger(__name__) @@ -76,7 +79,7 @@ class BadgerHomePage(QWidget): sig_routine_activated = pyqtSignal(bool) sig_routine_invalid = pyqtSignal() - def __init__(self, process_manager=None): + def __init__(self, process_manager: ProcessManager | None = None) -> None: logger.info("Initializing BadgerHomePage.") super().__init__() @@ -90,7 +93,7 @@ def __init__(self, process_manager=None): self.load_all_runs() self.init_home_page() - def init_ui(self): + def init_ui(self) -> None: logger.info("Initializing UI for BadgerHomePage.") self.config_singleton = init_settings() @@ -119,7 +122,7 @@ def init_ui(self): template_browser.setFixedWidth(360) # Splitter - splitter = QSplitter(Qt.Horizontal) + splitter = QSplitter(Qt.Orientation.Horizontal) splitter.setStretchFactor(0, 0) splitter.setStretchFactor(1, 1) @@ -149,7 +152,7 @@ def init_ui(self): font-weight: bold; padding: 4px; """) - title_label.setAlignment(Qt.AlignCenter) # Center-align the title + title_label.setAlignment(Qt.AlignmentFlag.AlignCenter) # Center-align the title vbox_table.addWidget(title_label, 0) self.run_table = run_table = data_table() run_table.set_uneditable() # should not be editable @@ -175,13 +178,13 @@ def init_ui(self): vbox_run.setContentsMargins(0, 0, 0, 0) vbox_run.setSpacing(0) - splitter_run = QSplitter(Qt.Horizontal) + splitter_run = QSplitter(Qt.Orientation.Horizontal) splitter_run.setStretchFactor(0, 1) splitter_run.setStretchFactor(1, 1) vbox_run.addWidget(splitter_run, 1) # splitter_run.hide() - splitter_data = QSplitter(Qt.Vertical) + splitter_data = QSplitter(Qt.Orientation.Vertical) splitter_data.setStretchFactor(0, 1) splitter_data.setStretchFactor(1, 0) splitter_data.addWidget(panel_monitor) @@ -212,7 +215,7 @@ def init_ui(self): status_bar.set_summary("Badger is ready!") vbox.addWidget(status_bar) - def config_logic(self): + def config_logic(self) -> None: logger.info("Configuring logic for BadgerHomePage.") self.colors = ["c", "g", "m", "y", "b", "r", "w"] self.symbols = ["o", "t", "t1", "s", "p", "h", "d"] @@ -291,17 +294,19 @@ def config_logic(self): self._configure_default_run_action() - def _configure_default_run_action(self): + def _configure_default_run_action(self) -> None: """Set the default run action as run_until_action""" self.run_action_bar.btn_stop.setDefaultAction( self.run_action_bar.run_until_action ) # configure default to max_eval (tc_idx=0), 50 iterations - initial_tc = {"tc_idx": 0, "max_eval": 100, "max_time": 300, "ftol": 0} + initial_tc = TerminationConditionConfig( + tc_idx=0, max_eval=100, max_time=300, ftol=0 + ) self.run_monitor.save_termination_condition(initial_tc) self.run_action_bar.update_run_tooltip(initial_tc) - def update_saved_values_from_monitor(self): + def update_saved_values_from_monitor(self) -> None: """ Sync Saved column values to match run monitor reset_env targets. @@ -312,12 +317,12 @@ def update_saved_values_from_monitor(self): self.routine_editor.set_saved_values_from_init_vars(variable_names, init_vars) - def load_all_runs(self): + def load_all_runs(self) -> None: logger.info("Loading all runs into history browser.") runs = get_runs() self.history_browser.updateItems(runs) - def init_home_page(self): + def init_home_page(self) -> None: logger.info("Initializing home page.") # Load the default generator # self.routine_editor.env_box.algo_cb.setCurrentIndex(-1) @@ -325,7 +330,7 @@ def init_home_page(self): self.routine_editor.env_box.algo_cb.setCurrentIndex(idx) self.routine_editor.select_generator(idx) - def go_run(self, i: int | None = None): + def go_run(self, i: int | None = None) -> None: logger.info(f"Activating run: {i}") gc.collect() @@ -382,7 +387,7 @@ def go_run(self, i: int | None = None): self.run_monitor.update_analysis_extensions() - def go_template(self, index: QModelIndex): + def go_template(self, index: QModelIndex) -> None: path = self.template_browser.file_sys_model.filePath(index) # if directory, expand it. if os.path.isdir(path): @@ -395,15 +400,15 @@ def go_template(self, index: QModelIndex): self.status_bar.set_summary(f"Current template {path}") return - def inspect_solution(self, idx): + def inspect_solution(self, idx: int) -> None: logger.info(f"Inspecting solution at index: {idx}") self.run_table.selectRow(idx) - def solution_selected(self, r, c): + def solution_selected(self, r: int, c: int) -> None: logger.info(f"Solution selected at row {r}, column {c}") self.run_monitor.jump_to_solution(r) - def table_selection_changed(self): + def table_selection_changed(self) -> None: logger.info("Table selection changed.") indices = self.run_table.selectedIndexes() if len(indices) == 1: # let other method handles it @@ -426,7 +431,7 @@ def table_selection_changed(self): self.run_monitor.jump_to_solution(row) - def toggle_lock(self, lock, lock_tab=1): + def toggle_lock(self, lock: bool, lock_tab: int = 1) -> None: logger.info(f"Toggling lock: {lock}, tab: {lock_tab}") if lock: self.history_browser.setDisabled(True) @@ -435,7 +440,7 @@ def toggle_lock(self, lock, lock_tab=1): self.uncover_page() - def validate_loaded_data_keys(self, vocs, open_dialog: bool = True): + def validate_loaded_data_keys(self, vocs: VOCS, open_dialog: bool = True) -> None: """ This function is called when adding historical data to a new routine. It makes sure that the keys of data to be loaded match the @@ -473,12 +478,14 @@ def validate_loaded_data_keys(self, vocs, open_dialog: bool = True): if open_dialog: # Notify user that data has been added to the routine dialog = QMessageBox( - text=str( + parent=self, + ) + dialog.setText( + str( "Data loaded into routine for the following VOCS:\n\n" + f"{list(data_keys)}\n\n" + "Click OK to continue!" - ), - parent=self, + ) ) dialog.setIcon(QMessageBox.Information) dialog.setWindowTitle("Data added to routine") @@ -489,7 +496,9 @@ def validate_loaded_data_keys(self, vocs, open_dialog: bool = True): self.run_action_bar.routine_finished() # Reset action bar raise BadgerRoutineError("Routine initialization cancelled by user.") - def prepare_run(self, data=None, init_points_flag=True): + def prepare_run( + self, data: DataFrame | None = None, init_points_flag: bool = True + ) -> None: """ Prepares the run by composing the routine, validating data if present, saving created routine to a yaml file, and passing the routine to @@ -560,7 +569,7 @@ def prepare_run(self, data=None, init_points_flag=True): def start_run( self, use_termination_condition: bool = False, load_displayed_data: bool = False - ): + ) -> None: """ Prepares and starts optimization run with provided options. - Termination Condition is provided when called via BadgerTerminationConditionDialog @@ -607,7 +616,7 @@ def start_run( def load_data_from_run(self): return self.current_routine.sorted_data - def start_run_until(self, dialog: bool = True): + def start_run_until(self, dialog: bool = True) -> None: """ Starts run with termination condition. @@ -637,7 +646,7 @@ def start_run_until(self, dialog: bool = True): else: self.start_run(use_termination_condition=True) - def new_run(self): + def new_run(self) -> None: logger.info("Creating new run.") self.cover_page() @@ -655,7 +664,7 @@ def new_run(self): header = get_header(self.current_routine) reset_table(self.run_table, header) - def run_name(self, name): + def run_name(self, name: str) -> None: logger.info(f"Updating run name: {name}") runs = get_runs() # block signals on update after routine finished, since the selected run is already diplayed @@ -664,11 +673,14 @@ def run_name(self, name): self.history_browser._selectItemByRun(name) self.history_browser.history_tree_widget.blockSignals(False) - def update_status(self, info): + def update_status(self, info: str) -> None: logger.info(f"Updating status: {info}") self.status_bar.set_summary(info) - def progress(self, solution: DataFrame): + def progress(self, solution: DataFrame) -> None: + if self.current_routine is None: + logger.warning("Current routine is None, cannot update progress.") + return vocs = self.current_routine.vocs vars = list(solution[vocs.variable_names].to_numpy()[0]) objs = list(solution[vocs.objective_names].to_numpy()[0]) @@ -680,7 +692,7 @@ def progress(self, solution: DataFrame): vocs.variable_names, list(solution[vocs.variable_names].to_numpy()[0]) ) - def delete_run(self): + def delete_run(self) -> None: logger.info("Deleting run.") run_name = get_base_run_filename(self.history_browser.currentText()) @@ -701,7 +713,7 @@ def delete_run(self): self.history_browser.history_tree_widget.blockSignals(False) self.go_run(-1) - def cover_page(self): + def cover_page(self) -> None: logger.info("Covering page with overlay.") # try: @@ -715,7 +727,7 @@ def cover_page(self): # self.overlay = ModalOverlay(main_window) # self.overlay.show() - def uncover_page(self): + def uncover_page(self) -> None: logger.info("Uncovering page overlay.") return # disable overlay for now diff --git a/src/badger/gui/mini/pages/routine_page.py b/src/badger/gui/mini/pages/routine_page.py index ebf52b69..5c08a4c7 100644 --- a/src/badger/gui/mini/pages/routine_page.py +++ b/src/badger/gui/mini/pages/routine_page.py @@ -22,6 +22,8 @@ import yaml from coolname import generate_slug from gest_api.vocs import ( + VOCS, + BaseConstraint, BaseObjective, ContinuousVariable, GreaterThanConstraint, @@ -44,7 +46,6 @@ QVBoxLayout, QWidget, ) -from xopt import VOCS from xopt.generators import ( all_generator_names, get_generator_defaults, @@ -105,7 +106,9 @@ def format_validation_error(e: ValidationError) -> str: return "\n".join(messages) -def extract_constraint_symbol_and_value(constraint): +def extract_constraint_symbol_and_value( + constraint: BaseConstraint, +) -> tuple[str, float]: if isinstance(constraint, GreaterThanConstraint): return ">", constraint.value if isinstance(constraint, LessThanConstraint): @@ -134,7 +137,7 @@ class BadgerRoutinePage(QWidget): sig_select_env = pyqtSignal(str) sig_status = pyqtSignal(str) - def __init__(self): + def __init__(self) -> None: logger.info("Initializing BadgerRoutinePage.") super().__init__() @@ -179,7 +182,7 @@ def __init__(self): # 2: not initialized, 1: apply to all, 0: apply to only visible self.lim_apply_to_vars = 2 - def init_ui(self): + def init_ui(self) -> None: logger.info("Initializing UI for BadgerRoutinePage.") self.config_singleton = config_singleton = init_settings() @@ -190,7 +193,9 @@ def init_ui(self): self.tabs = tabs = QTabWidget() vbox.addWidget(tabs) - tabs.tabBar().setExpanding(False) # keep tabs at content width + tab_bar = tabs.tabBar() + if tab_bar: + tab_bar.setExpanding(False) # keep tabs at content width tabs.setStyleSheet("QTabWidget::tab-bar { alignment: center; }") self.history_browser = HistoryNavigator() @@ -246,7 +251,7 @@ def init_ui(self): except KeyError: self.template_dir = os.path.join(self.BADGER_PLUGIN_ROOT, "templates") - def config_logic(self): + def config_logic(self) -> None: logger.info("Configuring logic for BadgerRoutinePage.") # self.btn_descr_update.clicked.connect(self.update_description) self.env_box.load_template_button.clicked.connect(self.load_template_yaml) @@ -303,7 +308,7 @@ def set_saved_values_from_init_vars( self.env_box.var_table.set_saved_values(values_by_name) def load_template_yaml( - self, checked_state=None, template_path: str | None = None + self, checked_state: bool | None = None, template_path: str | None = None ) -> None: logger.info("Loading template YAML.") """ @@ -344,7 +349,7 @@ def load_template_yaml( print(f"Error loading template: {e}") return - def set_options_from_template(self, template_dict: dict[str, Any]): + def set_options_from_template(self, template_dict: dict[str, Any]) -> None: logger.info( f"Setting options from template: {template_dict.get('name', 'unknown')}" ) @@ -668,7 +673,7 @@ def _filter_generator_params(self, generator_name: str, generator_config: dict): return generator_config - def save_template_yaml(self): + def save_template_yaml(self) -> None: logger.info("Saving routine as template YAML.") """ Save the current routine as a template .yaml file @@ -703,7 +708,7 @@ def save_template_yaml(self): logger.error(f"Error saving template: {e}") return - def refresh_ui(self, routine: Routine | None = None, silent: bool = False): + def refresh_ui(self, routine: Routine | None = None, silent: bool = False) -> None: logger.info( f"Refreshing UI for routine: {getattr(routine, 'name', None)} (silent={silent})" ) @@ -1075,7 +1080,7 @@ def refresh_params_generator(self): QMessageBox.warning(self, "Invalid script!", str(e)) @with_busy_cursor - def select_env(self, i: int): + def select_env(self, i: int) -> None: logger.info(f"Environment selected: {self.env_box.env_name} (index={i})") self.sig_status.emit("Loading variables...") diff --git a/src/badger/gui/mini/windows/main_window.py b/src/badger/gui/mini/windows/main_window.py index 6f338945..ed7987d9 100644 --- a/src/badger/gui/mini/windows/main_window.py +++ b/src/badger/gui/mini/windows/main_window.py @@ -4,11 +4,13 @@ from importlib import metadata from PyQt5.QtCore import QThread +from PyQt5.QtGui import QCloseEvent from PyQt5.QtWidgets import QDesktopWidget, QMainWindow, QMessageBox, QStackedWidget from badger.gui.components.create_process import CreateProcess from badger.gui.components.process_manager import ProcessManager from badger.gui.mini.pages.home_page import BadgerHomePage +from badger.types import ProcessWithArgs logger = logging.getLogger(__name__) @@ -18,7 +20,7 @@ def __init__(self) -> None: logger.info("Initializing BadgerMiniWindow.") super().__init__() - self.thread_list = [] + self.thread_list: list[QThread] = [] self.process_manager = ProcessManager() self.process_manager.processQueueUpdated.connect(self.addSubprocess) self.addSubprocess() @@ -30,19 +32,19 @@ def addSubprocess(self) -> None: Adds a subprocess to the subprocess queue. This method builds the subprocess on a QThread so as to not disrupt the main process. """ - self.thread = QThread() + self._thread = QThread() self.worker = CreateProcess() - self.worker.moveToThread(self.thread) + self.worker.moveToThread(self._thread) - self.thread.started.connect(self.worker.create_subprocess) + self._thread.started.connect(self.worker.create_subprocess) self.worker.subprocess_prepared.connect(self.storeSubprocess) - self.worker.finished.connect(self.thread.quit) + self.worker.finished.connect(self._thread.quit) self.worker.finished.connect(self.worker.deleteLater) - self.thread.finished.connect(self.thread.deleteLater) - self.thread.finished.connect(self.cleanupThread) + self._thread.finished.connect(self._thread.deleteLater) + self._thread.finished.connect(self.cleanupThread) - self.thread_list.append(self.thread) - self.thread.start() + self.thread_list.append(self._thread) + self._thread.start() def cleanupThread(self) -> None: logger.info("Cleaning up finished thread.") @@ -56,13 +58,13 @@ def cleanupThread(self) -> None: if thread in self.thread_list: self.thread_list.remove(thread) - def storeSubprocess(self, process_with_args: dict) -> None: + def storeSubprocess(self, process_with_args: ProcessWithArgs) -> None: logger.info(f"Storing prepared subprocess: {process_with_args}") """ Store the prepared subprocess for later use. Parameters: - process_with_args: Dict + process_with_args: ProcessWithArgs """ self.process_manager.add_to_queue(process_with_args) @@ -94,7 +96,7 @@ def center(self) -> None: qr.moveCenter(cp) self.move(qr.topLeft()) - def closeEvent(self, event) -> None: + def closeEvent(self, event: QCloseEvent | None) -> None: logger.info("Main window close event triggered.") if ( hasattr(self.home_page.routine_editor, "archive_search") @@ -104,7 +106,7 @@ def closeEvent(self, event) -> None: monitor = self.home_page.run_monitor if not monitor.running: - self.process_manager.close_proccesses() + self.process_manager.close_processes() monitor.destroy_unused_env() return @@ -119,13 +121,13 @@ def closeEvent(self, event) -> None: if reply == QMessageBox.Yes: - def close_window(): + def close_window() -> None: monitor.destroy_unused_env() self.close() monitor.register_post_run_action(close_window) monitor.testing = True # suppress the archive pop-ups - monitor.routine_runner.stop_routine() - event.ignore() - else: + if monitor.routine_runner: + monitor.routine_runner.stop_routine() + if event: event.ignore() diff --git a/src/badger/gui/pages/home_page.py b/src/badger/gui/pages/home_page.py index 8cda47ee..d769fdab 100644 --- a/src/badger/gui/pages/home_page.py +++ b/src/badger/gui/pages/home_page.py @@ -15,9 +15,10 @@ import os import traceback from importlib import resources -from typing import TYPE_CHECKING, Optional +from typing import Optional import numpy as np +from gest_api.vocs import VOCS from pandas import DataFrame from PyQt5.QtCore import QModelIndex, Qt, pyqtSignal from PyQt5.QtGui import QIcon, QKeySequence @@ -48,6 +49,7 @@ update_table, ) from badger.gui.components.navigators import HistoryNavigator, TemplateNavigator +from badger.gui.components.process_manager import ProcessManager from badger.gui.components.routine_page import BadgerRoutinePage from badger.gui.components.run_monitor import BadgerOptMonitor from badger.gui.components.status_bar import BadgerStatusBar @@ -58,14 +60,10 @@ from badger.gui.windows.terminition_condition_dialog import ( BadgerTerminationConditionDialog, ) +from badger.routine import Routine from badger.settings import init_settings from badger.utils import get_header -if TYPE_CHECKING: - from badger.gui.components.process_manager import ProcessManager - from badger.routine import VOCS - - logger = logging.getLogger(__name__) stylesheet = """ @@ -88,14 +86,14 @@ class BadgerHomePage(QWidget): sig_routine_activated = pyqtSignal(bool) sig_routine_invalid = pyqtSignal() - def __init__(self, process_manager: "ProcessManager | None" = None): + def __init__(self, process_manager: ProcessManager): logger.info("Initializing BadgerHomePage.") super().__init__() self.mode = "regular" # home page mode self.splitter_state = None # store the run splitter state self.process_manager = process_manager - self.current_routine = None # current routine + self.current_routine: Routine | None = None # current routine self.go_run_failed = False # flag to indicate go_run failed self.init_ui() self.config_logic() @@ -131,7 +129,7 @@ def init_ui(self) -> None: template_browser.setFixedWidth(360) # Splitter - splitter = QSplitter(Qt.Horizontal) + splitter = QSplitter(Qt.Orientation.Horizontal) splitter.setStretchFactor(0, 0) splitter.setStretchFactor(1, 1) @@ -161,7 +159,7 @@ def init_ui(self) -> None: font-weight: bold; padding: 4px; """) - title_label.setAlignment(Qt.AlignCenter) # Center-align the title + title_label.setAlignment(Qt.AlignmentFlag.AlignCenter) # Center-align the title vbox_table.addWidget(title_label, 0) self.run_table = run_table = data_table() run_table.set_uneditable() # should not be editable @@ -187,12 +185,12 @@ def init_ui(self) -> None: vbox_run.setContentsMargins(0, 0, 0, 0) vbox_run.setSpacing(0) - splitter_run = QSplitter(Qt.Horizontal) + splitter_run = QSplitter(Qt.Orientation.Horizontal) splitter_run.setStretchFactor(0, 1) splitter_run.setStretchFactor(1, 1) vbox_run.addWidget(splitter_run, 1) - splitter_data = QSplitter(Qt.Vertical) + splitter_data = QSplitter(Qt.Orientation.Vertical) splitter_data.setStretchFactor(0, 1) splitter_data.setStretchFactor(1, 0) splitter_data.addWidget(panel_monitor) @@ -380,7 +378,7 @@ def go_template(self, index: QModelIndex) -> None: return # otherwise, open the template - self.routine_editor.load_template_yaml(False, template_path=path) + self.routine_editor.load_template_yaml(template_path=path) self.status_bar.set_summary(f"Current template {path}") return @@ -462,13 +460,13 @@ def validate_loaded_data_keys(self, vocs: "VOCS") -> None: # Notify user that data has been added to the routine dialog = QMessageBox( - text=str( - "Data loaded into routine for the following VOCS:\n\n" - + f"{list(data_keys)}\n\n" - + "Click OK to continue!" - ), parent=self, ) + dialog.setText( + "Data loaded into routine for the following VOCS:\n\n" + + f"{list(data_keys)}\n\n" + + "Click OK to continue!" + ) dialog.setIcon(QMessageBox.Information) dialog.setWindowTitle("Data added to routine") dialog.setStandardButtons(QMessageBox.Ok | QMessageBox.Cancel) diff --git a/src/badger/gui/utils.py b/src/badger/gui/utils.py index d9e276ef..020dae6e 100644 --- a/src/badger/gui/utils.py +++ b/src/badger/gui/utils.py @@ -8,12 +8,11 @@ from collections.abc import Callable from functools import wraps from importlib import resources -from typing import Any +from typing import Any, ParamSpec, TypeVar from PyQt5.QtCore import QEvent, QObject, QSize, Qt from PyQt5.QtGui import QIcon from PyQt5.QtWidgets import ( - QAbstractSpinBox, QApplication, QComboBox, QDialog, @@ -34,36 +33,41 @@ if BADGER_TEMP_DIRECTORY is None: raise BadgerConfigError("Please set the BADGER_TEMP_DIRECTORY env var!") +from PyQt5.QtGui import QFocusEvent, QMouseEvent +from PyQt5.QtWidgets import ( + QWidget, +) -def preventAnnoyingSpinboxScrollBehaviour(self, control: QAbstractSpinBox) -> None: - control.setFocusPolicy(Qt.StrongFocus) - control.installEventFilter(self.MouseWheelWidgetAdjustmentGuard(control)) +T = TypeVar("T") +P = ParamSpec("P") class MouseWheelWidgetAdjustmentGuard(QObject): def __init__(self, parent: QObject): super().__init__(parent) - def eventFilter(self, o: QObject, e: QEvent) -> bool: + def eventFilter(self, a0: QObject, a1: QEvent) -> bool: # Ignore mouse wheel events for widget - if e.type() == QEvent.Wheel: - e.ignore() + if a1.type() == QEvent.Wheel: + a1.ignore() return True - return super().eventFilter(o, e) + return super().eventFilter(a0, a1) def create_button( - icon_file, - tooltip, - stylesheet=None, - size=(32, 32), - icon_size=None, - tool_button=False, -): + icon_file: str, + tooltip: str, + stylesheet: str | None = None, + size: tuple[int, int] = (32, 32), + icon_size: tuple[int, int] | None = None, + tool_button: bool = False, +) -> QPushButton | QToolButton: icon_ref = resources.files(__package__) / f"./images/{icon_file}" with resources.as_file(icon_ref) as icon_path: icon = QIcon(str(icon_path)) + btn: QPushButton | QToolButton + if tool_button: btn = QToolButton() btn.setToolButtonStyle(Qt.ToolButtonStyle.ToolButtonIconOnly) @@ -139,28 +143,29 @@ def filter_generator_config(name: str, config: dict[str, Any]) -> dict[str, Any] class NoHoverFocusComboBox(QComboBox): - def focusInEvent(self, event): + def focusInEvent(self, event: QFocusEvent) -> None: # Prevent focus if it's from a hover event if event.reason() == Qt.MouseFocusReason: event.ignore() else: super().focusInEvent(event) - def mousePressEvent(self, event): + def mousePressEvent(self, event: QMouseEvent) -> None: # Ensure the combo box behaves normally when clicked self.setFocus() super().mousePressEvent(event) - def wheelEvent(self, event): + def wheelEvent(self, event: QEvent) -> None: event.ignore() class ModalOverlay(QDialog): - def __init__(self, parent=None, info=None): + def __init__(self, parent: QWidget | None = None, info: str | None = None) -> None: super().__init__(parent) self.setWindowModality(Qt.ApplicationModal) self.setWindowFlags(Qt.FramelessWindowHint | Qt.Dialog) - self.setGeometry(parent.geometry()) + if parent is not None: + self.setGeometry(parent.geometry()) layout = QVBoxLayout() if info is None: @@ -182,9 +187,9 @@ def unset_busy_cursor() -> None: QApplication.restoreOverrideCursor() -def with_busy_cursor(func: Callable) -> Callable: +def with_busy_cursor(func: Callable[P, T]) -> Callable[P, T]: @wraps(func) - def wrapped(*args, **kwargs): + def wrapped(*args: P.args, **kwargs: P.kwargs) -> T: set_busy_cursor() try: return func(*args, **kwargs) diff --git a/src/badger/gui/windows/add_random_dialog.py b/src/badger/gui/windows/add_random_dialog.py index fcc73bb2..91f77c1d 100644 --- a/src/badger/gui/windows/add_random_dialog.py +++ b/src/badger/gui/windows/add_random_dialog.py @@ -3,8 +3,10 @@ the current operating point before starting an optimization run.""" from copy import deepcopy +from dataclasses import dataclass from PyQt5.QtCore import Qt +from PyQt5.QtGui import QKeyEvent from PyQt5.QtWidgets import ( QComboBox, QDialog, @@ -22,24 +24,37 @@ from badger.gui.components.robust_spinbox import RobustSpinBox +@dataclass +class BadgerAddRandomConfig: + method: int + n_points: int + fraction: float + + class BadgerAddRandomDialog(QDialog): - def __init__(self, parent, add_points, save_config, configs=None): + def __init__( + self, + parent: QWidget | None = None, + add_points=None, + save_config=None, + configs: dict | None = None, + ) -> None: super().__init__(parent) self.add_points = add_points self.save_config = save_config self.configs = deepcopy(configs) if configs is None: - self.configs = { - "method": 0, - "n_points": 3, - "fraction": 0.1, - } + self.configs = BadgerAddRandomConfig( + method=0, + n_points=3, + fraction=0.1, + ) self.init_ui() self.config_logic() - def init_ui(self): + def init_ui(self) -> None: self.setWindowTitle("Add random points") self.setMinimumWidth(360) @@ -121,32 +136,34 @@ def init_ui(self): vbox.addWidget(group_config, 1) vbox.addWidget(button_set) - def config_logic(self): + def config_logic(self) -> None: self.cb.currentIndexChanged.connect(self.method_changed) self.btn_cancel.clicked.connect(self.close) self.btn_add.clicked.connect(self.add) self.sb_np.valueChanged.connect(self.n_points_changed) self.sb_frac.valueChanged.connect(self.frac_changed) - def n_points_changed(self, n_points): + def n_points_changed(self, n_points: int) -> None: self.configs["n_points"] = n_points - def frac_changed(self, frac): + def frac_changed(self, frac: float) -> None: self.configs["fraction"] = frac - def add(self): + def add(self) -> None: self.save_config(self.configs) self.add_points() self.close() - def method_changed(self, i): + def method_changed(self, i: int) -> None: self.stacks.setCurrentIndex(i) # Update configs self.configs["method"] = i - def keyPressEvent(self, event): - if event.key() == Qt.Key_Return or event.key() == Qt.Key_Enter: + def keyPressEvent(self, event: QKeyEvent | None) -> None: + if event is None: + return + if event.key() == Qt.Key.Key_Return or event.key() == Qt.Key.Key_Enter: self.add() - elif event.key() == Qt.Key_Escape: + elif event.key() == Qt.Key.Key_Escape: self.close() diff --git a/src/badger/gui/windows/expandable_message_box.py b/src/badger/gui/windows/expandable_message_box.py index 08f1e3b8..c98d1416 100644 --- a/src/badger/gui/windows/expandable_message_box.py +++ b/src/badger/gui/windows/expandable_message_box.py @@ -18,8 +18,13 @@ class ExpandableMessageBox(QDialog): def __init__( - self, icon=None, title="Message", text="", detailedText="", parent=None - ): + self, + icon: QMessageBox.Icon | None = None, + title: str = "Message", + text: str = "", + detailedText: str = "", + parent: QDialog | None = None, + ) -> None: super().__init__(parent) self.setWindowFlags(self.windowFlags() | Qt.WindowMaximizeButtonHint) @@ -78,7 +83,7 @@ def __init__( self.setWindowTitle(title) self.resize(420, 250) - def toggle_details(self): + def toggle_details(self) -> None: if self.detailedTextWidget.isVisible(): self.detailedTextWidget.setVisible(False) self.toggleButton.setText("Show Details") @@ -86,13 +91,13 @@ def toggle_details(self): self.detailedTextWidget.setVisible(True) self.toggleButton.setText("Hide Details") - def setText(self, text): + def setText(self, text: str) -> None: self.textLabel.setText(text) - def setDetailedText(self, detailedText): + def setDetailedText(self, detailedText: str) -> None: self.detailedTextWidget.setText(detailedText) - def setIcon(self, icon: QMessageBox.Icon): + def setIcon(self, icon: QMessageBox.Icon) -> None: # This maps the QMessageBox icons to the QDialog iconMap = { QMessageBox.Icon.Information: QMessageBox.standardIcon( diff --git a/src/badger/gui/windows/load_data_from_run_dialog.py b/src/badger/gui/windows/load_data_from_run_dialog.py index e9346c2a..a855c0e4 100644 --- a/src/badger/gui/windows/load_data_from_run_dialog.py +++ b/src/badger/gui/windows/load_data_from_run_dialog.py @@ -7,7 +7,9 @@ import numpy as np import pandas as pd import pyqtgraph as pg +from gest_api.vocs import VOCS from PyQt5.QtCore import Qt +from PyQt5.QtGui import QCloseEvent from PyQt5.QtWidgets import ( QDialog, QFileDialog, @@ -19,7 +21,6 @@ QWidget, ) from pyqtgraph.Qt import QtCore, QtGui -from xopt.vocs import VOCS from badger.archive import ( get_base_run_filename, @@ -57,22 +58,23 @@ class BadgerLoadDataFromRunDialog(QDialog): def __init__( self, parent: QWidget, - env_vocs: VOCS = None, - on_set: Callable[[Routine], None] | None = None, + env_vocs: VOCS, + on_set: Callable[[Routine], None], ): """ Initialize the dialog. Args: parent (QWidget): The parent widget. - data_table (QTableWidget, optional): The data table to update with loaded run data. + env_vocs (VOCS): The environment VOCS object containing variable and objective names. + on_set (Callable[[Routine], None] | None, optional): Function to call when loading data. Attributes: data_table (QTableWidget): The data table to update with loaded run data. selected_routine (Optional[Routine]): The currently selected routine. """ super().__init__(parent) - self.setWindowFlags(self.windowFlags() | Qt.WindowMaximizeButtonHint) + self.setWindowFlags(self.windowFlags() | Qt.WindowType.WindowMaximizeButtonHint) self.selected_routine = None self.env_vocs = env_vocs @@ -146,7 +148,7 @@ def init_ui(self) -> None: vbox.addWidget(content_widget) vbox.addWidget(button_set) - def preview_run(self, routine: Routine = None) -> None: + def preview_run(self, routine: Routine | None = None) -> None: """ Add data to plot to preview the selected run. """ @@ -158,9 +160,6 @@ def preview_run(self, routine: Routine = None) -> None: get_base_run_filename(self.history_browser.currentText()) ) - if routine is None: - return - # Configure plots curves_objective = self._configure_plot( self.plot_obj, @@ -233,24 +232,24 @@ def load_from_file(self) -> None: except Exception as e: raise BadgerRoutineError(f"{e}") from e - def load_data(self, run_filename: str) -> None: + def load_data(self, run_filename: str) -> Routine | None: """ Load data from the selected run. Returns: - Optional[Routine]: The loaded routine or None if loading fails. + Routine | None: The loaded routine or None if loading fails. """ if not run_filename: - return + return None try: routine = load_run(run_filename) return routine except IndexError: - return + return None except Exception: # noqa: BLE001 - run load boundary - return + return None - def init_plots(self) -> None: + def init_plots(self) -> pg.GraphicsLayoutWidget: """ Initialize the plots for data preview. These are static plots styled to match the plots on the main GUI from BadgerOptMonitor @@ -301,17 +300,17 @@ def init_plots(self) -> None: def _configure_plot( self, plot_object: pg.PlotItem, names: list[str] - ) -> dict[str : pg.PlotCurveItem]: + ) -> dict[str, pg.PlotCurveItem]: """ Configure the plot with the given data names. Adapted from BadgerOptMonitor._configure_plot Args: plot_object (pg.PlotItem): The plot object to configure. - names (List[str]): The names of the data series. + names (list[str]): The names of the data series. Returns: - dict: A dictionary mapping data names to plot curves. + dict[str, pg.PlotCurveItem]: A dictionary mapping data names to plot curves. """ plot_object.clear() curves = {} @@ -344,7 +343,7 @@ def _configure_plot( return curves def _set_plot_data( - self, names: list[str], curves: dict, data: pd.DataFrame + self, names: list[str], curves: dict[str, pg.PlotCurveItem], data: pd.DataFrame ) -> None: """ Set data for the plot curves. @@ -378,7 +377,7 @@ def _set_plot_data( hist_x, not_live_data[name].to_numpy(dtype=np.double) ) - def show_vocs_mismatch_dialog(self, list1: list[str], list2: list[str]): + def show_vocs_mismatch_dialog(self, list1: list[str], list2: list[str]) -> None: """ Display a helpful dialog notifying the user that the data they are trying to load does not have the same variables and objectives as they have selected in the GUI. @@ -397,8 +396,9 @@ def show_vocs_mismatch_dialog(self, list1: list[str], list2: list[str]): dialog.setStandardButtons(QMessageBox.Ok) dialog.exec_() - def cancel_changes(self): + def cancel_changes(self) -> None: self.close() - def closeEvent(self, event): - event.accept() + def closeEvent(self, event: QCloseEvent | None) -> None: + if event: + event.accept() diff --git a/src/badger/gui/windows/main_window.py b/src/badger/gui/windows/main_window.py index 2199edcc..502ee0a9 100644 --- a/src/badger/gui/windows/main_window.py +++ b/src/badger/gui/windows/main_window.py @@ -6,11 +6,13 @@ from importlib import metadata from PyQt5.QtCore import QThread +from PyQt5.QtGui import QCloseEvent from PyQt5.QtWidgets import QDesktopWidget, QMainWindow, QMessageBox, QStackedWidget from badger.gui.components.create_process import CreateProcess from badger.gui.components.process_manager import ProcessManager from badger.gui.pages.home_page import BadgerHomePage +from badger.types import ProcessWithArgs logger = logging.getLogger(__name__) @@ -19,7 +21,7 @@ class BadgerMainWindow(QMainWindow): def __init__(self) -> None: logger.info("Initializing BadgerMainWindow.") super().__init__() - self.thread_list = [] + self.thread_list: list[QThread] = [] self.process_manager = ProcessManager() self.process_manager.processQueueUpdated.connect(self.addSubprocess) self.addSubprocess() @@ -32,19 +34,19 @@ def addSubprocess(self) -> None: Adds a subprocess to the subprocess queue. This method builds the subprocess on a QThread so as to not disrupt the main process. """ - self.thread = QThread() + self._thread = QThread() self.worker = CreateProcess() - self.worker.moveToThread(self.thread) + self.worker.moveToThread(self._thread) - self.thread.started.connect(self.worker.create_subprocess) + self._thread.started.connect(self.worker.create_subprocess) self.worker.subprocess_prepared.connect(self.storeSubprocess) - self.worker.finished.connect(self.thread.quit) + self.worker.finished.connect(self._thread.quit) self.worker.finished.connect(self.worker.deleteLater) - self.thread.finished.connect(self.thread.deleteLater) - self.thread.finished.connect(self.cleanupThread) + self._thread.finished.connect(self._thread.deleteLater) + self._thread.finished.connect(self.cleanupThread) - self.thread_list.append(self.thread) - self.thread.start() + self.thread_list.append(self._thread) + self._thread.start() def cleanupThread(self) -> None: logger.info("Cleaning up finished thread.") @@ -52,19 +54,20 @@ def cleanupThread(self) -> None: Method to remove threads no longer active from the thread list. Parameters: - thread: QThread + None + The thread to be cleaned up is inferred from the sender of the signal. """ thread = self.sender() if thread in self.thread_list: self.thread_list.remove(thread) - def storeSubprocess(self, process_with_args: dict) -> None: + def storeSubprocess(self, process_with_args: ProcessWithArgs) -> None: logger.info(f"Storing prepared subprocess: {process_with_args}") """ Store the prepared subprocess for later use. Parameters: - process_with_args: Dict + process_with_args: ProcessWithArgs """ self.process_manager.add_to_queue(process_with_args) @@ -105,7 +108,7 @@ def center(self) -> None: def config_logic(self) -> None: logger.info("Configuring logic.") - def closeEvent(self, event) -> None: + def closeEvent(self, event: QCloseEvent | None) -> None: logger.info("Main window close event triggered.") if ( hasattr(self.home_page.routine_editor, "archive_search") @@ -115,7 +118,7 @@ def closeEvent(self, event) -> None: monitor = self.home_page.run_monitor if not monitor.running: - self.process_manager.close_proccesses() + self.process_manager.close_processes() monitor.destroy_unused_env() return @@ -130,13 +133,14 @@ def closeEvent(self, event) -> None: if reply == QMessageBox.Yes: - def close_window(): + def close_window() -> None: monitor.destroy_unused_env() self.close() monitor.register_post_run_action(close_window) monitor.testing = True # suppress the archive pop-ups - monitor.routine_runner.stop_routine() - event.ignore() - else: + if monitor.routine_runner is not None: + monitor.routine_runner.stop_routine() + + if event is not None: event.ignore() diff --git a/src/badger/gui/windows/measurement_retry_dialog.py b/src/badger/gui/windows/measurement_retry_dialog.py index b5bed059..ec40cff5 100644 --- a/src/badger/gui/windows/measurement_retry_dialog.py +++ b/src/badger/gui/windows/measurement_retry_dialog.py @@ -1,13 +1,15 @@ """Error dialog that allows users to retry measurements after errors in setting variable values / getting observables from the environment. Based on the ExpandableMessageBox dialog.""" -from PyQt5.QtWidgets import QDialogButtonBox, QMessageBox +from PyQt5.QtWidgets import QDialog, QDialogButtonBox, QMessageBox from .expandable_message_box import ExpandableMessageBox class BadgerMeasurementRetryDialog(ExpandableMessageBox): - def __init__(self, text="", detailedText="", parent=None): + def __init__( + self, text: str = "", detailedText: str = "", parent: QDialog | None = None + ) -> None: full_text = ( "There was an error setting variables or getting observables.\n\n" + text ) @@ -33,9 +35,11 @@ def __init__(self, text="", detailedText="", parent=None): QDialogButtonBox.Ok | QDialogButtonBox.Cancel ) self.retryButton = self.dialogButtonBox.button(QDialogButtonBox.Ok) - self.retryButton.setText("Retry Measurement") + if self.retryButton is not None: + self.retryButton.setText("Retry Measurement") self.stopButton = self.dialogButtonBox.button(QDialogButtonBox.Cancel) - self.stopButton.setText("Stop Run") + if self.stopButton is not None: + self.stopButton.setText("Stop Run") self.dialogButtonBox.accepted.connect(self.accept) self.dialogButtonBox.rejected.connect(self.reject) diff --git a/src/badger/gui/windows/message_dialog.py b/src/badger/gui/windows/message_dialog.py index 7a73155c..e5fcd947 100644 --- a/src/badger/gui/windows/message_dialog.py +++ b/src/badger/gui/windows/message_dialog.py @@ -3,7 +3,7 @@ used for showing verbose error information or optimization summaries.""" from PyQt5.QtCore import Qt -from PyQt5.QtGui import QFont, QFontDatabase, QTextOption +from PyQt5.QtGui import QFont, QFontDatabase, QPixmap, QTextOption from PyQt5.QtWidgets import ( QDialog, QHBoxLayout, @@ -20,10 +20,15 @@ class BadgerScrollableMessageBox(QDialog): def __init__( - self, icon=None, title="Message", text="", detailedText="", parent=None - ): + self, + icon: QPixmap | None = None, + title: str = "Message", + text: str = "", + detailedText: str = "", + parent: QDialog | None = None, + ) -> None: super().__init__(parent) - self.setWindowFlags(self.windowFlags() | Qt.WindowMaximizeButtonHint) + self.setWindowFlags(self.windowFlags() | Qt.WindowType.WindowMaximizeButtonHint) # Main layout mainLayout = QVBoxLayout(self) @@ -66,19 +71,27 @@ def __init__( self.setWindowTitle(title) self.resize(420, 300) - def setText(self, text): + def setText(self, text: str) -> None: self.textLabel.setText(text) - def setDetailedText(self, detailedText): + def setDetailedText(self, detailedText: str) -> None: self.detailedTextWidget.setText(detailedText) - def setIcon(self, icon): + def setIcon(self, icon: QPixmap) -> None: # This maps the QMessageBox icons to the QDialog iconMap = { - QMessageBox.Information: QMessageBox.standardIcon(QMessageBox.Information), - QMessageBox.Warning: QMessageBox.standardIcon(QMessageBox.Warning), - QMessageBox.Critical: QMessageBox.standardIcon(QMessageBox.Critical), - QMessageBox.Question: QMessageBox.standardIcon(QMessageBox.Question), + QMessageBox.Icon.Information: QMessageBox.standardIcon( + QMessageBox.Icon.Information + ), + QMessageBox.Icon.Warning: QMessageBox.standardIcon( + QMessageBox.Icon.Warning + ), + QMessageBox.Icon.Critical: QMessageBox.standardIcon( + QMessageBox.Icon.Critical + ), + QMessageBox.Icon.Question: QMessageBox.standardIcon( + QMessageBox.Icon.Question + ), } standardIcon = iconMap.get(icon) if standardIcon: diff --git a/src/badger/gui/windows/review_dialog.py b/src/badger/gui/windows/review_dialog.py index ee6367a7..71012b76 100644 --- a/src/badger/gui/windows/review_dialog.py +++ b/src/badger/gui/windows/review_dialog.py @@ -2,14 +2,20 @@ preview of a fully-configured routine before it is saved or run, allowing the user to confirm the settings.""" -from PyQt5.QtWidgets import QDialog, QDialogButtonBox, QTextBrowser, QVBoxLayout +from PyQt5.QtWidgets import ( + QDialog, + QDialogButtonBox, + QTextBrowser, + QVBoxLayout, + QWidget, +) from badger.routine import Routine from badger.utils import get_yaml_string class BadgerReviewDialog(QDialog): - def __init__(self, parent, routine: Routine): + def __init__(self, parent: QWidget | None, routine: Routine) -> None: super().__init__(parent) self.routine = routine @@ -17,7 +23,7 @@ def __init__(self, parent, routine: Routine): self.init_ui() self.config_logic() - def init_ui(self): + def init_ui(self) -> None: name = self.routine.name self.setWindowTitle(f"Review routine {name}") @@ -31,5 +37,5 @@ def init_ui(self): vbox.addWidget(brow_routine) vbox.addWidget(btn_ok) - def config_logic(self): + def config_logic(self) -> None: self.btn_ok.accepted.connect(self.accept) diff --git a/src/badger/gui/windows/settings_dialog.py b/src/badger/gui/windows/settings_dialog.py index 57c61be0..5ff25738 100644 --- a/src/badger/gui/windows/settings_dialog.py +++ b/src/badger/gui/windows/settings_dialog.py @@ -4,7 +4,7 @@ import logging import os -from typing import ClassVar +from typing import ClassVar, cast from PyQt5.QtCore import Qt @@ -37,10 +37,10 @@ class BadgerSettingsDialog(QDialog): "dark": 2, } - def __init__(self, parent): + def __init__(self, parent: QWidget | None = None) -> None: logger.info("Initializing BadgerSettingsDialog.") super().__init__(parent) - self.setWindowFlags(self.windowFlags() | Qt.WindowMaximizeButtonHint) + self.setWindowFlags(self.windowFlags() | Qt.WindowType.WindowMaximizeButtonHint) self.config_singleton = init_settings() @@ -51,7 +51,7 @@ def __init__(self, parent): self.init_ui() self.config_logic() - def init_ui(self): + def init_ui(self) -> None: logger.info("Initializing settings dialog UI.") self.setWindowTitle("Badger settings") self.setMinimumWidth(480) @@ -191,20 +191,24 @@ def init_ui(self): vbox.addWidget(widget_settings) - self.btns = QDialogButtonBox(QDialogButtonBox.Ok | QDialogButtonBox.Cancel) # type: ignore + self.btns = QDialogButtonBox(QDialogButtonBox.Ok | QDialogButtonBox.Cancel) vbox.addStretch(1) vbox.addWidget(self.btns) - def config_logic(self): + def config_logic(self) -> None: logger.info("Configuring logic for settings dialog.") # self.cb_theme.currentIndexChanged.connect(self.select_theme) self.btns.accepted.connect(self.apply_settings) self.btns.rejected.connect(self.restore_settings) - def set_theme(self, theme): + def set_theme(self, theme: str) -> None: logger.info(f"Setting theme: {theme}") app = QApplication.instance() + if app is None: + logger.warning("No QApplication instance found. Cannot set theme.") + return + app = cast(QApplication, app) if theme == "dark": app.setStyleSheet(load_stylesheet(palette=DarkPalette)) elif theme == "light": @@ -212,14 +216,14 @@ def set_theme(self, theme): else: app.setStyleSheet("") - def select_theme(self, i): + def select_theme(self, i: int) -> None: logger.info(f"Theme selected: {self.theme_list[i]}") theme = self.theme_list[i] self.set_theme(theme) # Update the internal settings self.config_singleton.write_value("BADGER_THEME", theme) - def apply_settings(self): + def apply_settings(self) -> None: logger.info("Applying settings.") self.accept() self.config_singleton.write_value( @@ -270,7 +274,7 @@ def apply_settings(self): logging_manager.update_logfile_path(new_logfile_path) logger.info(f"Runtime: Logs now writing to {new_logfile_path}") - def restore_settings(self): + def restore_settings(self) -> None: logger.info("Restoring previous settings.") # Reset theme if needed theme_curr = self.config_singleton.read_value("BADGER_THEME") diff --git a/src/badger/gui/windows/termination_reached_dialog.py b/src/badger/gui/windows/termination_reached_dialog.py index 56c3cc43..b9697a24 100644 --- a/src/badger/gui/windows/termination_reached_dialog.py +++ b/src/badger/gui/windows/termination_reached_dialog.py @@ -4,6 +4,8 @@ after the run is paused by a termination condition. """ +from typing import Any, cast + from PyQt5.QtCore import Qt from PyQt5.QtWidgets import ( QDialog, @@ -46,7 +48,12 @@ class BadgerTerminationReachedDialog(QDialog): - def __init__(self, tc_condition=None, text="", parent=None): + def __init__( + self, + tc_condition: dict[str, Any] | None = None, + text: str = "", + parent: QDialog | None = None, + ) -> None: super().__init__(parent) self.setWindowTitle("Termination Condition Reached") @@ -56,6 +63,7 @@ def __init__(self, tc_condition=None, text="", parent=None): layout.setContentsMargins(14, 14, 14, 14) layout.setSpacing(8) + tc_condition = cast(dict[str, Any], tc_condition) tc_type = tc_condition["type"] if tc_type == "max_eval": tc_type_text = "N iterations" @@ -72,12 +80,18 @@ def __init__(self, tc_condition=None, text="", parent=None): summary_label = QLabel(f"{tc_type_text}: {state} / {tc_condition['config']}") summary_label.setWordWrap(True) - summary_label.setAlignment(Qt.AlignLeft | Qt.AlignVCenter) + # PyQt5 stubs type flag OR as int, so cast back to the flag type. + summary_label.setAlignment( + cast( + Qt.AlignmentFlag, + Qt.AlignmentFlag.AlignLeft | Qt.AlignmentFlag.AlignVCenter, + ) + ) # summary_label.setStyleSheet("color: #8A949E;") text_column.addWidget(summary_label) body_label = QLabel("Badger optimization paused") - body_label.setAlignment(Qt.AlignLeft) + body_label.setAlignment(Qt.AlignmentFlag.AlignLeft) text_column.addWidget(body_label) content_row.addLayout(text_column) @@ -91,12 +105,14 @@ def __init__(self, tc_condition=None, text="", parent=None): font.setPointSize(12) self.setFont(font) - self.continueButton.setText("Continue") - # self.continueButton.setStyleSheet(stylesheet_run) - self.continueButton.setFixedSize(96, 24) - self.endButton.setText("End Run") - self.endButton.setStyleSheet(stylesheet_stop) - self.endButton.setFixedSize(96, 24) + if self.continueButton is not None: + self.continueButton.setText("Continue") + # self.continueButton.setStyleSheet(stylesheet_run) + self.continueButton.setFixedSize(96, 24) + if self.endButton is not None: + self.endButton.setText("End Run") + self.endButton.setStyleSheet(stylesheet_stop) + self.endButton.setFixedSize(96, 24) button_box.accepted.connect(self.accept) button_box.rejected.connect(self.reject) diff --git a/src/badger/gui/windows/terminition_condition_dialog.py b/src/badger/gui/windows/terminition_condition_dialog.py index d0e8e640..750d19b9 100644 --- a/src/badger/gui/windows/terminition_condition_dialog.py +++ b/src/badger/gui/windows/terminition_condition_dialog.py @@ -2,6 +2,9 @@ when an optimization run should automatically stop — either after a maximum number of evaluations or after a maximum elapsed time.""" +from collections.abc import Callable + +from PyQt5.QtGui import QCloseEvent from PyQt5.QtWidgets import ( QComboBox, QDialog, @@ -17,6 +20,8 @@ QWidget, ) +from badger.types import TerminationConditionConfig + stylesheet_run = """ QPushButton:hover:pressed { @@ -35,19 +40,28 @@ class BadgerTerminationConditionDialog(QDialog): - def __init__(self, parent, run_opt, save_config, configs=None): + def __init__( + self, + run_opt: Callable[[bool, bool], None], + save_config: Callable[[TerminationConditionConfig], None], + parent: QDialog | None = None, + configs: TerminationConditionConfig | None = None, + ): super().__init__(parent) self.run_opt = run_opt self.save_config = save_config - self.configs = configs if configs is None: - self.configs = {"tc_idx": 0, "max_eval": 42, "max_time": 600, "ftol": 0} + self.configs = TerminationConditionConfig( + tc_idx=0, max_eval=42, max_time=600, ftol=0 + ) + else: + self.configs = configs self.init_ui() self.config_logic() - def init_ui(self): + def init_ui(self) -> None: self.setWindowTitle("Optimization termination condition") self.setMinimumWidth(360) @@ -69,7 +83,7 @@ def init_ui(self): # 'optimization converged', ] ) - cb.setCurrentIndex(self.configs["tc_idx"]) + cb.setCurrentIndex(self.configs.tc_idx) hbox.addWidget(lbl) hbox.addWidget(cb, 1) @@ -87,7 +101,7 @@ def init_ui(self): self.sb_max_eval = sb_max_eval = QSpinBox() sb_max_eval.setMinimum(1) sb_max_eval.setMaximum(100000) - sb_max_eval.setValue(self.configs["max_eval"]) + sb_max_eval.setValue(self.configs.max_eval) sb_max_eval.setSingleStep(1) hbox_max_eval.addWidget(lbl) hbox_max_eval.addWidget(sb_max_eval, 1) @@ -99,7 +113,7 @@ def init_ui(self): self.sb_max_time = sb_max_time = QDoubleSpinBox() sb_max_time.setMinimum(0.0) sb_max_time.setMaximum(86400) - sb_max_time.setValue(self.configs["max_time"]) + sb_max_time.setValue(self.configs.max_time) sb_max_time.setDecimals(2) sb_max_time.setSingleStep(0.1) hbox_max_time.addWidget(lbl) @@ -112,7 +126,7 @@ def init_ui(self): self.sb_tol = sb_tol = QDoubleSpinBox() sb_tol.setMinimum(0.0) sb_tol.setMaximum(1000) - sb_tol.setValue(self.configs["ftol"]) + sb_tol.setValue(self.configs.ftol) sb_tol.setDecimals(6) sb_tol.setSingleStep(0.001) hbox_tol.addWidget(lbl) @@ -122,7 +136,7 @@ def init_ui(self): stacks.addWidget(max_time_config) stacks.addWidget(tol_config) - stacks.setCurrentIndex(self.configs["tc_idx"]) + stacks.setCurrentIndex(self.configs.tc_idx) vbox_config.addWidget(stacks) vbox_config.addStretch(1) @@ -143,35 +157,36 @@ def init_ui(self): vbox.addWidget(group_config, 1) vbox.addWidget(button_set) - def config_logic(self): + def config_logic(self) -> None: self.cb.currentIndexChanged.connect(self.terminition_condition_changed) - self.btn_cancel.clicked.connect(self.close) + self.btn_cancel.clicked.connect(self.closeEvent) self.btn_run.clicked.connect(self.run) self.sb_max_eval.valueChanged.connect(self.max_eval_changed) self.sb_max_time.valueChanged.connect(self.max_time_changed) self.sb_tol.valueChanged.connect(self.ftol_changed) - def max_eval_changed(self, max_eval): - self.configs["max_eval"] = max_eval + def max_eval_changed(self, max_eval: int) -> None: + self.configs.max_eval = max_eval - def max_time_changed(self, max_time): - self.configs["max_time"] = max_time + def max_time_changed(self, max_time: float) -> None: + self.configs.max_time = max_time - def ftol_changed(self, ftol): - self.configs["ftol"] = ftol + def ftol_changed(self, ftol: float) -> None: + self.configs.ftol = ftol - def run(self): + def run(self) -> None: self.save_config(self.configs) - self.run_opt(True) + self.run_opt(True, False) self.close() - def terminition_condition_changed(self, i): + def terminition_condition_changed(self, i: int) -> None: self.stacks.setCurrentIndex(i) # Update configs - self.configs["tc_idx"] = i + self.configs.tc_idx = i - def closeEvent(self, event): + def closeEvent(self, event: QCloseEvent | None) -> None: self.save_config(self.configs) - event.accept() + if event is not None: + event.accept() diff --git a/src/badger/gui/windows/var_dialog.py b/src/badger/gui/windows/var_dialog.py index a69f3cf5..ddb7a02f 100644 --- a/src/badger/gui/windows/var_dialog.py +++ b/src/badger/gui/windows/var_dialog.py @@ -3,6 +3,7 @@ plugin to verify the variable exists and retrieve its current value.""" import logging +from collections.abc import Callable from PyQt5.QtWidgets import ( QDialog, @@ -15,14 +16,21 @@ QWidget, ) -from badger.environment import instantiate_env +from badger.environment import Environment, instantiate_env +from badger.factory import BadgerPluginConfig from badger.gui.components.labeled_lineedit import labeled_lineedit logger = logging.getLogger(__name__) class BadgerVariableDialog(QDialog): - def __init__(self, parent, env_class, configs, callback): + def __init__( + self, + env_class: type[Environment], + configs: BadgerPluginConfig, + callback: Callable[[str, float, float], int], + parent: QDialog | None = None, + ) -> None: super().__init__(parent) self.env = instantiate_env(env_class, configs) @@ -31,7 +39,7 @@ def __init__(self, parent, env_class, configs, callback): self.init_ui() self.config_logic() - def init_ui(self): + def init_ui(self) -> None: self.setWindowTitle("Add variable") self.setMinimumWidth(360) @@ -78,13 +86,13 @@ def init_ui(self): vbox.addStretch() vbox.addWidget(button_set) - def config_logic(self): + def config_logic(self) -> None: self.edit_name.textChanged.connect(self.var_changed) - self.btn_cancel.clicked.connect(self.close) + self.btn_cancel.clicked.connect(self.closeEvent) self.btn_check.clicked.connect(self.check_var) self.btn_add.clicked.connect(self.add_var) - def var_changed(self, text): + def var_changed(self, text: str) -> None: self.btn_add.setDisabled(True) if text: @@ -92,7 +100,7 @@ def var_changed(self, text): else: self.btn_check.setDisabled(True) - def check_var(self): + def check_var(self) -> None: name = self.edit_name.text() try: value = self.env.get_variables([name])[name] @@ -112,7 +120,7 @@ def check_var(self): self.btn_add.setDisabled(True) - def add_var(self): + def add_var(self) -> None: name = self.edit_name.text() min = float(self.edit_min.edit.text()) max = float(self.edit_max.edit.text()) diff --git a/src/badger/interface.py b/src/badger/interface.py index 902965fd..43abff78 100644 --- a/src/badger/interface.py +++ b/src/badger/interface.py @@ -4,6 +4,7 @@ import pickle from abc import ABC, abstractmethod +from collections.abc import Callable, Mapping from typing import Any, ClassVar, TypedDict from pydantic import BaseModel @@ -11,8 +12,8 @@ from badger.utils import curr_ts -def log(func): - def func_log(*args, **kwargs): +def log(func: Callable[..., Any]) -> Callable[..., Any]: + def func_log(*args: Any, **kwargs: Any) -> Any: if func.__name__ == "set_values": if "channel_inputs" in kwargs: channel_inputs = kwargs["channel_inputs"] @@ -55,12 +56,12 @@ class Interface(BaseModel, ABC): # params: float = Field(..., description='Example intf parameter') # Private variables - _logs: list[dict] = [] # TODO: Add a property for it? + _logs: list[dict[str, Any]] = [] # TODO: Add a property for it? - def start_recording(self): + def start_recording(self) -> None: self._logs = [] - def stop_recording(self, filename): + def stop_recording(self, filename: str) -> None: # Do not create the pickle file if log is empty # This usually happens when the log decorator is not used in the # specific interface @@ -72,7 +73,7 @@ def stop_recording(self, filename): self._logs = [] - def dump_recording(self, filename): + def dump_recording(self, filename: str) -> None: # Dump the logs to disk w/o cleaning up the log history if not self._logs: return @@ -82,27 +83,27 @@ def dump_recording(self, filename): # Environment should only call this method to get channels @abstractmethod - def get_values(self, channel_names: list[str]) -> dict[str, Any]: - pass + def get_values(self, channel_names: list[str]) -> dict[str, float | list[float]]: + raise NotImplementedError() # Environment should only call this method to set channels @abstractmethod - def set_values(self, channel_inputs: dict[str, Any]): - pass + def set_values(self, channel_inputs: Mapping[str, float | list[float]]) -> None: + raise NotImplementedError() - def reset_interface(self): + def reset_interface(self) -> None: """ Called after the application forks (i.e. after spawning a new multiprocess.Process) - Subclasses should use this to reset any undesirable process wide state + Subclasses may override this to reset any undesirable process wide state. """ - def get_value(self, channel_name: str, **kwargs) -> Any: + def get_value(self, channel_name: str, **kwargs: Any) -> Any: return self.get_values([channel_name], **kwargs)[channel_name] - def set_value(self, channel_name: str, channel_value, **kwargs): + def set_value(self, channel_name: str, channel_value: Any, **kwargs: Any) -> None: return self.set_values({channel_name: channel_value}, **kwargs) - def get_info(self, channels: list[str]) -> InterfaceInfo: + def get_info(self, channels: list[str]) -> InterfaceInfo | None: """ Optional; Returns information about the channels and environment for display diff --git a/src/badger/log.py b/src/badger/log.py index 6f76062f..9879a236 100644 --- a/src/badger/log.py +++ b/src/badger/log.py @@ -14,27 +14,38 @@ import atexit import logging import os +from argparse import Namespace from datetime import UTC, datetime from logging.handlers import QueueHandler, QueueListener -from multiprocessing import Queue +from multiprocessing import Queue, current_process from badger.settings import get_user_config_folder, init_settings logger = logging.getLogger(__name__) +def _get_level_name(log_level: str | int) -> str: + """Non-deprecated replacement for logging.getLevelName (int -> level name).""" + if isinstance(log_level, int): + name_by_level = { + level: name for name, level in logging.getLevelNamesMapping().items() + } + return name_by_level.get(log_level, str(log_level)) + return log_level + + class LoggingManager: """ Used to manage logging across multiple processes, one listener thread collects logs from all worker sub-processes via a queue. """ - def __init__(self): - self.log_queue: Queue = None - self.listener: QueueListener = None - self.handlers = [] + def __init__(self) -> None: + self.log_queue: Queue[logging.LogRecord] | None = None + self.listener: QueueListener | None = None + self.handlers: list[logging.Handler] = [] - def start_listener(self, log_filepath: str, log_level: str | int): + def start_listener(self, log_filepath: str, log_level: str | int) -> None: """ For use in main process to setup queue and start listening for logs from sub processes. @@ -75,16 +86,16 @@ def start_listener(self, log_filepath: str, log_level: str | int): ) # '*' unpacks the array for us (unpacking operator) self.listener.start() - logger.info( - f"Centralized logging listener started with level {logging.getLevelName(log_level)}" - ) + log_level_name = _get_level_name(log_level) + + logger.info(f"Centralized logging listener started with level {log_level_name}") - def stop_listener(self): + def stop_listener(self) -> None: if self.listener: self.listener.stop() logger.info("Main process logging listener stopped") - def update_log_level(self, log_level: str | int): + def update_log_level(self, log_level: str | int) -> None: """ Update log level in the handlers and logger objects in main process (only two handlers atm, terminal and file). This affects which level of logs (from any process) get written to the logfile. @@ -108,11 +119,11 @@ def update_log_level(self, log_level: str | int): if isinstance(logger_obj, logging.Logger) and name.startswith("badger"): logger_obj.setLevel(log_level) - logger.info( - f"Log level updated to {logging.getLevelName(log_level)} for badger namespace." - ) + log_level_name = _get_level_name(log_level) - def update_logfile_path(self, new_logfile_path: str): + logger.info(f"Log level updated to {log_level_name} for badger namespace.") + + def update_logfile_path(self, new_logfile_path: str) -> None: """ Update logfile path, so logs get written to a different file. @@ -121,11 +132,19 @@ def update_logfile_path(self, new_logfile_path: str): the QueueListener while passing in the new handlers. """ # Get the existing handler - file_handler = None + file_handler: logging.FileHandler | None = None for h in self.handlers: if isinstance(h, logging.FileHandler): file_handler = h break + if file_handler is None: + logger.error("No file handler found, cannot update logfile path") + return + + if file_handler is None: + raise ValueError( + "No file handler found in handlers list. Cannot update logfile path." + ) # We want to keep using same formatter and loglevel formatter = file_handler.formatter @@ -146,6 +165,9 @@ def update_logfile_path(self, new_logfile_path: str): self.handlers.append(new_file_handler) # Restart QueueListener, using new handlers + if self.log_queue is None: + logger.error("No log queue found, cannot restart QueueListener") + return self.listener = QueueListener( self.log_queue, *self.handlers, respect_handler_level=True ) @@ -153,7 +175,7 @@ def update_logfile_path(self, new_logfile_path: str): logger.info(f"Logfile path updated to {new_logfile_path}") - def get_logfile_name(self): + def get_logfile_name(self) -> str: """ Get name of the logfile for today, which is in form of: "log__.log" @@ -162,11 +184,11 @@ def get_logfile_name(self): log_filename = f"log_{today.year:04d}_{today.month:02d}_{today.day:02d}.log" return log_filename - def get_queue(self) -> Queue: + def get_queue(self) -> Queue[logging.LogRecord] | None: """Get the logging queue for use by subprocesses.""" return self.log_queue - def create_log_dir(self, log_dir: str): + def create_log_dir(self, log_dir: str) -> None: """ Create the logs directory if it doesn't exist. @@ -174,7 +196,7 @@ def create_log_dir(self, log_dir: str): log_dir: directory to potentially create. """ # If not set, empty, or invalid, use default (/logs dir in root of repo) - if log_dir is None: + if log_dir == "": log_dir = "logs" # Expand user home directory if needed @@ -196,9 +218,9 @@ def create_log_dir(self, log_dir: str): def configure_process_logging( - log_queue: Queue | None = None, + log_queue: Queue[logging.LogRecord] | None = None, logger_name: str = "badger", - log_level: str = "DEBUG", + log_level: str | int = "DEBUG", process_name: str | None = None, ) -> None: """ @@ -209,16 +231,14 @@ def configure_process_logging( args: logger_name (str): Name of the logger (default: "badger") log_level (int or str): Logging level to set - process_name (str): Custom process name to output in logs + process_name (str | None): Custom process name to output in logs """ if isinstance(log_level, str): - log_level = getattr(logging, log_level.upper(), logging.DEBUG) + log_level = getattr(logging, log_level.upper(), "DEBUG") # Set custom process name if provided if process_name: - import multiprocessing as mp - - mp.current_process().name = process_name + current_process().name = process_name # Get top level logger for this namespace logger = logging.getLogger(logger_name) @@ -239,12 +259,12 @@ def configure_process_logging( if isinstance(logger_obj, logging.Logger) and name.startswith(logger_name): logger_obj.setLevel(log_level) - logger.info( - f"process logger configured with level {logging.getLevelName(log_level)}" - ) + log_level_name = _get_level_name(log_level) + + logger.info(f"process logger configured with level {log_level_name}") -def setup_logging(args): +def setup_logging(args: Namespace) -> None: """ Init Badger's multiprocess logging system """ @@ -274,7 +294,7 @@ def setup_logging(args): # Ensure listener shuts down at exit atexit.register( - lambda: logging_manager.listener and logging_manager.listener.stop() + lambda: logging_manager.listener.stop() if logging_manager.listener else None ) diff --git a/src/badger/logbook.py b/src/badger/logbook.py index d57a4a94..06a66ae0 100644 --- a/src/badger/logbook.py +++ b/src/badger/logbook.py @@ -6,8 +6,11 @@ import os from datetime import UTC, datetime +from PyQt5.QtWidgets import QWidget + from badger.archive import BADGER_ARCHIVE_ROOT from badger.errors import BadgerConfigError, BadgerLogbookError +from badger.routine import Routine from badger.settings import init_settings logger = logging.getLogger(__name__) @@ -22,7 +25,7 @@ logger.info(f"Badger logbook root {BADGER_LOGBOOK_ROOT} created") -def send_to_logbook(routine, widget=None): +def send_to_logbook(routine: Routine, widget: QWidget | None = None) -> None: from re import sub from xml.etree import ElementTree @@ -38,6 +41,8 @@ def send_to_logbook(routine, widget=None): obj_opt = obj_opt[0] data = routine.data + if data is None: + raise BadgerLogbookError("No data to log!") obj_start = data[obj_name].iloc[0] duration = data["timestamp"].iloc[-1] - data["timestamp"].iloc[0] n_point = len(data["timestamp"]) @@ -59,7 +64,7 @@ def send_to_logbook(routine, widget=None): timestr = curr_time.strftime("%Y-%m-%dT%H%M%S") else: timestr = curr_time.strftime("%Y-%m-%dT%H:%M:%S") - log_entry = ElementTree.Element(None) + log_entry = ElementTree.Element("logentry") severity = ElementTree.SubElement(log_entry, "severity") location = ElementTree.SubElement(log_entry, "location") keywords = ElementTree.SubElement(log_entry, "keywords") @@ -99,22 +104,24 @@ def send_to_logbook(routine, widget=None): screenshot(widget, f"{fileName}.png") -def screenshot(widget, filename): +def screenshot(widget: QWidget | None, filename: str) -> None: """ Takes a screenshot of the whole gui window, saves png and ps images to file """ if widget is None: raise BadgerLogbookError("No widget to take screenshot on!") - from PIL import Image + from PIL import Image, ImageFile pic = widget.grab() pic.save(filename) img = Image.open(filename) + image_to_save: Image.Image | ImageFile.ImageFile = img if img.mode in ("RGBA", "LA"): # https://pillow.readthedocs.io/en/stable/handbook/image-file-formats.html?highlight=eps#eps # logger.warning(f'Current figure mode "{img.mode}" cannot be directly saved to .ps and will be converted to "RGB" mode') - img = img.convert("RGB") + converted_img = img.convert("RGB") + image_to_save = converted_img # img = img.scaled(400, 600) name = os.path.splitext(filename)[0] - img.save(f"{name}.ps") + image_to_save.save(f"{name}.ps") diff --git a/src/badger/logger/__init__.py b/src/badger/logger/__init__.py index 659544cc..275c3760 100644 --- a/src/badger/logger/__init__.py +++ b/src/badger/logger/__init__.py @@ -4,34 +4,31 @@ import json import os +from collections.abc import Callable -from badger.logger.event import Events +from badger.logger.event import Events, Solution from badger.logger.observer import _Tracker from badger.logger.util import Colours -def _get_default_logger(verbose): - return ScreenLogger(verbose=verbose) - - class ScreenLogger(_Tracker): _default_cell_size = 10 _default_precision = 4 - def __init__(self, verbose=2): + def __init__(self, verbose: int = 2) -> None: self._verbose = verbose - self._header_length = None + self._header_length: int | None = None super().__init__() @property - def verbose(self): + def verbose(self) -> int: return self._verbose @verbose.setter - def verbose(self, v): + def verbose(self, v: int) -> None: self._verbose = v - def _format_number(self, x): + def _format_number(self, x: float) -> str: if isinstance(x, int): s = "{x:< {s}}".format( x=x, @@ -51,57 +48,71 @@ def _format_number(self, x): return s[: self._default_cell_size - 3] + "..." return s - def _format_key(self, key): + def _format_key(self, key: str) -> str: s = "{key:^{s}}".format(key=key, s=self._default_cell_size) if len(s) > self._default_cell_size: return s[: self._default_cell_size - 3] + "..." return s - def _step(self, solution, colour=Colours.black): + def _step( + self, + solution: Solution, + colour: Callable[[str], str] = Colours.black, + ) -> str: # solution: (x: 1d array, y: 1d array, c: 1d array, s: 1d array, is_optimal: bool, # vars: str list, obses: str list, cons: str list, stas: str list) - cells = [] + cells: list[str] = [] cells.append(self._format_number(self._iterations + 1)) - for o in solution[1]: + for o in solution.objectives or []: cells.append(self._format_number(o)) - for c in solution[2]: + for c in solution.constraints or []: cells.append(self._format_number(c)) - for v in solution[0]: + for v in solution.variables or []: cells.append(self._format_number(v)) - for s in solution[3]: + for s in solution.states or []: cells.append(self._format_number(s)) # TODO: deal with the str case return "| " + " | ".join(map(colour, cells)) + " |" - def _header(self, solution): + def _header( + self, + solution: Solution, + ) -> str: cells = [] cells.append(self._format_key("iter")) - for obs in solution[6]: - cells.append(self._format_key(obs)) + for obj in solution.objective_names: + cells.append(self._format_key(obj)) - for con in solution[7]: + for con in solution.constraint_names: cells.append(self._format_key(con)) - for var in solution[5]: + for var in solution.variable_names: cells.append(self._format_key(var)) - for sta in solution[8]: - cells.append(self._format_key(sta)) + for obs in solution.observable_names: + cells.append(self._format_key(obs)) - line = "| " + " | ".join(cells) + " |" + line: str = "| " + " | ".join(cells) + " |" self._header_length = len(line) return line + "\n" + ("-" * self._header_length) - def _is_new_max(self, solution): - return solution[4] - - def update(self, event, solution): + def _is_new_max( + self, + solution: Solution, + ) -> bool: + return solution.is_optimal + + def update( + self, + event: Events, + solution: Solution, + ) -> None: if event == Events.OPTIMIZATION_START: line = self._header(solution) + "\n" elif event == Events.OPTIMIZATION_STEP: @@ -112,15 +123,19 @@ def update(self, event, solution): colour = Colours.purple if is_new_max else Colours.black line = self._step(solution, colour=colour) + "\n" elif event == Events.OPTIMIZATION_END: - line = "=" * self._header_length + "\n" + line = "=" * (self._header_length or 0) + "\n" if self._verbose: print(line, end="") self._update_tracker(event, solution) +def _get_default_logger(verbose: int) -> ScreenLogger: + return ScreenLogger(verbose=verbose) + + class JSONLogger(_Tracker): - def __init__(self, path, reset=True): + def __init__(self, path: str, reset: bool = True) -> None: self._path = path if path[-5:] == ".json" else path + ".json" if reset: try: @@ -129,14 +144,14 @@ def __init__(self, path, reset=True): pass super().__init__() - def update(self, event, solution): + def update(self, event: Events, solution: Solution) -> None: if event == Events.OPTIMIZATION_STEP: - data = { - "x": solution[0], - "y": solution[1], - "c": solution[2], - "s": solution[3], - "is_optimal": solution[4], + data: dict[str, object] = { + "x": solution.variables, + "y": solution.objectives, + "c": solution.constraints, + "s": solution.states, + "is_optimal": solution.is_optimal, } now, time_elapsed, time_delta = self._time_metrics() diff --git a/src/badger/logger/event.py b/src/badger/logger/event.py index e6b12100..752a26ce 100644 --- a/src/badger/logger/event.py +++ b/src/badger/logger/event.py @@ -1,8 +1,23 @@ """Event constants (start, step, end) that the optimization loop fires to notify loggers and observers of progress.""" +from enum import StrEnum +from typing import NamedTuple -class Events: + +class Solution(NamedTuple): + variables: list[float] | None + objectives: list[float] | None + constraints: list[float] | None + states: list[float] | None + is_optimal: bool + variable_names: list[str] + objective_names: list[str] + constraint_names: list[str] + observable_names: list[str] + + +class Events(StrEnum): OPTIMIZATION_START = "optimization:start" OPTIMIZATION_STEP = "optimization:step" OPTIMIZATION_END = "optimization:end" diff --git a/src/badger/logger/observer.py b/src/badger/logger/observer.py index 72c39c50..4b266cd7 100644 --- a/src/badger/logger/observer.py +++ b/src/badger/logger/observer.py @@ -3,26 +3,26 @@ from datetime import UTC, datetime -from badger.logger.event import Events +from badger.logger.event import Events, Solution class Observer: - def update(self, event, solution): + def update(self, event: Events, solution: Solution) -> None: raise NotImplementedError class _Tracker: - def __init__(self): + def __init__(self) -> None: self._iterations = 0 - self._start_time = None - self._previous_time = None + self._start_time: datetime | None = None + self._previous_time: datetime | None = None - def _update_tracker(self, event, solution): + def _update_tracker(self, event: Events, solution: Solution) -> None: if event == Events.OPTIMIZATION_STEP: self._iterations += 1 - def _time_metrics(self): + def _time_metrics(self) -> tuple[str, float, float]: now = datetime.now(tz=UTC) if self._start_time is None: self._start_time = now diff --git a/src/badger/logger/util.py b/src/badger/logger/util.py index 400621ba..cfe61792 100644 --- a/src/badger/logger/util.py +++ b/src/badger/logger/util.py @@ -17,55 +17,55 @@ class Colours: YELLOW = "\033[93m" @classmethod - def _wrap_colour(cls, s, colour): + def _wrap_colour(cls, s: str, colour: str) -> str: return colour + s + cls.END @classmethod - def black(cls, s): + def black(cls, s: str) -> str: """Wrap text in black.""" return cls._wrap_colour(s, cls.END) @classmethod - def blue(cls, s): + def blue(cls, s: str) -> str: """Wrap text in blue.""" return cls._wrap_colour(s, cls.BLUE) @classmethod - def bold(cls, s): + def bold(cls, s: str) -> str: """Wrap text in bold.""" return cls._wrap_colour(s, cls.BOLD) @classmethod - def cyan(cls, s): + def cyan(cls, s: str) -> str: """Wrap text in cyan.""" return cls._wrap_colour(s, cls.CYAN) @classmethod - def darkcyan(cls, s): + def darkcyan(cls, s: str) -> str: """Wrap text in darkcyan.""" return cls._wrap_colour(s, cls.DARKCYAN) @classmethod - def green(cls, s): + def green(cls, s: str) -> str: """Wrap text in green.""" return cls._wrap_colour(s, cls.GREEN) @classmethod - def purple(cls, s): + def purple(cls, s: str) -> str: """Wrap text in purple.""" return cls._wrap_colour(s, cls.PURPLE) @classmethod - def red(cls, s): + def red(cls, s: str) -> str: """Wrap text in red.""" return cls._wrap_colour(s, cls.RED) @classmethod - def underline(cls, s): + def underline(cls, s: str) -> str: """Wrap text in underline.""" return cls._wrap_colour(s, cls.UNDERLINE) @classmethod - def yellow(cls, s): + def yellow(cls, s: str) -> str: """Wrap text in yellow.""" return cls._wrap_colour(s, cls.YELLOW) diff --git a/src/badger/routine.py b/src/badger/routine.py index 843eb1c7..02c0a3c7 100644 --- a/src/badger/routine.py +++ b/src/badger/routine.py @@ -14,15 +14,11 @@ import json import logging from copy import deepcopy -from typing import Any +from typing import Any, cast import numpy as np import pandas as pd - -# Import xopt.generators at startup so they don't need to be imported -# each time a Routine is created -import xopt.generators.bayesian -import xopt.generators.sequential # noqa: F401 +from gest_api.vocs import VOCS from pandas import DataFrame from pydantic import ( ConfigDict, @@ -32,12 +28,12 @@ field_validator, model_validator, ) -from xopt import VOCS, Evaluator, Xopt +from xopt import Evaluator, Xopt from xopt.generators import get_generator from xopt.generators.sequential import SequentialGenerator from xopt.vocs import get_local_region -from badger.environment import BaseEnvironment, instantiate_env +from badger.environment import BaseEnvironment, Environment, instantiate_env from badger.factory import get_env from badger.utils import curr_ts @@ -50,15 +46,17 @@ class Routine(Xopt): name: str description: str | None = Field(None) environment: SerializeAsAny[BaseEnvironment] + # Built from `environment` in validate_model, so callers never pass it. + evaluator: SerializeAsAny[Evaluator] | None = Field(None) initial_points: DataFrame | None = Field(None) critical_constraint_names: list[str] | None = Field([]) - tags: list | None = Field(None) + tags: dict[str, Any] | None = Field(None) script: str | None = Field(None) # Store relative to current params relative_to_current: bool | None = Field(False) - vrange_limit_options: dict | None = Field(None) - vrange_hard_limit: dict | None = Field({}) # override hard limits - initial_point_actions: list | None = Field(None) + vrange_limit_options: dict[str, dict[str, Any]] | None = Field(None) + vrange_hard_limit: dict[str, list[float]] | None = Field({}) # override hard limits + initial_point_actions: list[dict[str, Any]] | None = Field(None) additional_variables: list[str] | None = Field([]) formulas: dict[str, dict[str, Any]] | None = Field({}) constraint_formulas: dict[str, dict[str, Any]] | None = Field({}) @@ -71,7 +69,7 @@ class Routine(Xopt): @model_validator(mode="before") @classmethod - def validate_model(cls, data: Any): + def validate_model(cls, data: Any) -> Any: logger.info("Validating Routine model from input data.") if isinstance(data, dict): logger.debug(f"Routine data dict received: {list(data.keys())}") @@ -113,20 +111,20 @@ def validate_model(cls, data: Any): if "data" in data and isinstance(data["data"], dict): logger.debug("Validating and converting data to DataFrame.") try: - data["data"] = pd.DataFrame(data["data"]) + df = pd.DataFrame(data["data"]) except IndexError: - data["data"] = pd.DataFrame(data["data"], index=[0]) - - data["data"].index = data["data"].index.astype(int) - data["data"].sort_index(inplace=True) + df = pd.DataFrame(data["data"], index=[0]) + df.index = df.index.astype(int) + df.sort_index(inplace=True) + data["data"] = df # Add data one row at a time to avoid generator issues if isinstance(data["generator"], SequentialGenerator): logger.debug("Setting data for SequentialGenerator.") - data["generator"].set_data(data["data"]) + data["generator"].set_data(df) else: logger.debug("Adding data to generator.") - data["generator"].add_data(data["data"]) + data["generator"].add_data(df) # instantiate env if isinstance(data["environment"], dict): @@ -142,11 +140,13 @@ def validate_model(cls, data: Any): data["environment"] = instantiate_env(env_class, configs_env) # create evaluator - env = data["environment"] + env: Environment = data["environment"] - def evaluate_point(point: dict): + def evaluate_point( + point: dict[str, float], + ) -> dict[str, float | list[float]]: logger.debug(f"Evaluating point: {point}") - point = pd.Series(point).explode().to_dict() + point = cast(dict[str, float], pd.Series(point).explode().to_dict()) env.set_variables(point) obs = env.get_observables(data["generator"].vocs.output_names) ts = curr_ts() @@ -160,7 +160,7 @@ def evaluate_point(point: dict): return data @field_validator("initial_points", mode="before") - def validate_data(cls, v, info: ValidationInfo): + def validate_data(cls, v: Any, info: ValidationInfo) -> Any: logger.debug("Validating initial_points field.") if isinstance(v, dict): try: @@ -170,7 +170,7 @@ def validate_data(cls, v, info: ValidationInfo): return v @property - def sorted_data(self): + def sorted_data(self) -> pd.DataFrame | None: logger.debug("Sorting routine data.") data_copy = deepcopy(self.data) if data_copy is not None: @@ -178,7 +178,7 @@ def sorted_data(self): data_copy.sort_index(inplace=True) return data_copy - def json(self, **kwargs) -> str: + def json(self, **kwargs: Any) -> str: logger.info("Serializing Routine to JSON.") """Handle custom serialization of environment""" @@ -200,25 +200,30 @@ def json(self, **kwargs) -> str: dict_result["environment"] = {"name": self.environment.name} | dict_result[ "environment" ] - try: - dict_result["environment"]["interface"] = { - "name": self.environment.interface.name - } | dict_result["environment"]["interface"] - except KeyError: - pass - except AttributeError: - pass + # Only Environment (not BaseEnvironment) carries an interface. + if ( + isinstance(self.environment, Environment) + and self.environment.interface is not None + ): + try: + dict_result["environment"]["interface"] = { + "name": self.environment.interface.name + } | dict_result["environment"]["interface"] + except KeyError: + pass return json.dumps(dict_result) -def calculate_variable_bounds(limit_options, vocs, env): +def calculate_variable_bounds( + limit_options: dict[str, dict[str, Any]], vocs: VOCS, env: BaseEnvironment +) -> dict[str, list[float]]: logger.info("Calculating variable bounds.") vnames = vocs.variable_names var_curr = env.get_variables(vnames) var_range = env.get_bounds(vnames) - variables_updated = {} + variables_updated: dict[str, list[float]] = {} for name in vnames: try: limit_option = limit_options[name] @@ -250,10 +255,12 @@ def calculate_variable_bounds(limit_options, vocs, env): return variables_updated -def calculate_initial_points(init_actions, vocs, env): +def calculate_initial_points( + init_actions: list[dict[str, Any]], vocs: VOCS, env: BaseEnvironment +) -> dict[str, list[float]]: logger.info("Calculating initial points.") vnames = vocs.variable_names - init_points = {k: [] for k in vnames} + init_points: dict[str, list[float]] = {k: [] for k in vnames} for action in init_actions: logger.debug(f"Processing initial point action: {action}") diff --git a/src/badger/settings.py b/src/badger/settings.py index 7eab7aa5..c26df299 100644 --- a/src/badger/settings.py +++ b/src/badger/settings.py @@ -13,7 +13,7 @@ import platform import shutil from importlib import resources -from typing import Any +from typing import Any, Self, cast import yaml from pydantic import BaseModel, Field, ValidationError @@ -152,25 +152,28 @@ class BadgerConfig(BaseModel): class ConfigSingleton: - _instance = None + _instance: "ConfigSingleton | None" = None + user_flag: bool + _config: BadgerConfig + config_path: str | None - def __new__(cls, config_path: str | None = None, user_flag: bool = False): + def __new__(cls, config_path: str | None = None, user_flag: bool = False) -> Self: if cls._instance is None: cls._instance = super().__new__(cls) cls._instance.user_flag = user_flag cls._instance._config = cls.load_or_create_config(config_path) cls._instance.config_path = config_path - return cls._instance + return cast(Self, cls._instance) @classmethod - def load_or_create_config(cls, config_path: str) -> BadgerConfig: + def load_or_create_config(cls, config_path: str | None) -> BadgerConfig: """ Loads the config file from a given yaml file if it exists, otherwise creates an instance of BadgerConfig with default settings. Parameters ---------- - config_path: str + config_path: str | None Path to the user config file. Returns @@ -179,7 +182,7 @@ def load_or_create_config(cls, config_path: str) -> BadgerConfig: An instance of BadgerConfig populated with the data from the config file, or with default settings if the file does not exist. """ - if os.path.exists(config_path): + if config_path is not None and os.path.exists(config_path): with open(config_path, "r") as config_file: config_data = yaml.safe_load(config_file) @@ -211,7 +214,7 @@ def load_or_create_config(cls, config_path: str) -> BadgerConfig: print(f"Error validating config file: {e}") raise else: - if cls._instance.user_flag: + if cls._instance and cls._instance.user_flag: err_msg = f"Error loading config {config_path}: invalid path." raise BadgerLoadConfigError(err_msg) @@ -236,8 +239,9 @@ def update_and_save_config(self, updates: dict[str, Any]) -> None: self._update_config_by_dot_key(config_data, dot_key, value) # Save updated config to file - with open(self.config_path, "w") as file: - yaml.dump(config_data, file, default_flow_style=False) + if self.config_path is not None: + with open(self.config_path, "w") as file: + yaml.dump(config_data, file, default_flow_style=False) self._config = BadgerConfig(**config_data) @@ -260,7 +264,7 @@ def list_settings(self) -> dict[str, Any]: Returns ------- - result: Dict + result: dict[str, Any] A dictionary containing the settings. Keys in the dict are fields of the settings, the value for each key is the current value for that setting. """ @@ -271,7 +275,7 @@ def list_path_settings(self) -> dict[str, Any]: Returns ------- - result: Dict + result: dict[str, Any] A dictionary containing the path-related settings. Keys in the dict are fields of the settings, the value for each key is the current value for that setting. @@ -408,7 +412,7 @@ def write_value(self, key: str, value: Any) -> None: The value that is being saved. """ keys = key.split(".") - updates = {} + updates: dict[str, Any] = {} sub_dict = updates for k in keys[:-1]: @@ -535,7 +539,7 @@ def apply_pytorch_multiprocess_tensor_sharing_setting( return try: - torch.multiprocessing.set_sharing_strategy(strategy) + torch.multiprocessing.set_sharing_strategy(strategy) # type: ignore[no-untyped-call] except Exception: logger.exception( "Can't set pytorch multiprocess tensor-sharing strategy to '%s'", strategy @@ -545,7 +549,7 @@ def apply_pytorch_multiprocess_tensor_sharing_setting( logger.info("Set pytorch multiprocess tensor-sharing strategy to '%s'", strategy) -def mock_settings(): +def mock_settings() -> None: """A method for setting up mock settings""" config_singleton = init_settings() app_data_dir = get_datadir() / "Badger" @@ -555,7 +559,7 @@ def mock_settings(): plugins_dir = str(app_data_dir / "plugins") os.makedirs(plugins_dir, exist_ok=True) config_singleton.write_value("BADGER_PLUGIN_ROOT", plugins_dir) - built_in_plugins_dir = resources.files(__package__) / "built_in_plugins" + built_in_plugins_dir = str(resources.files(__package__) / "built_in_plugins") shutil.copytree(built_in_plugins_dir, plugins_dir, dirs_exist_ok=True) logbook_dir = str(app_data_dir / "logbook") @@ -599,9 +603,14 @@ def get_user_config_folder() -> str: If the operating system is not supported. """ system = platform.system() + config_folder: str | None if system == "Windows": config_folder = os.getenv("APPDATA") or os.getenv("LOCALAPPDATA") + if config_folder is None: + raise OSError( + "Neither APPDATA nor LOCALAPPDATA environment variable is set on this Windows system." + ) elif system == "Darwin": config_folder = os.path.expanduser("~/Library/Application Support/Badger") elif system == "Linux": diff --git a/src/badger/stats.py b/src/badger/stats.py index 6f4d4e66..45a67964 100644 --- a/src/badger/stats.py +++ b/src/badger/stats.py @@ -1,48 +1,51 @@ """Reducer functions (mean, median, percentile, etc.) that collapse array-valued observations into a single scalar for the optimizer.""" +from typing import Any + import numpy as np +from numpy.typing import NDArray -def none(data): +def none(data: NDArray[Any]) -> NDArray[Any]: return data -def median(data): +def median(data: NDArray[Any]) -> Any: return np.median(data) -def std_deviation(data): +def std_deviation(data: NDArray[Any]) -> Any: return np.std(data) -def median_deviation(data): +def median_deviation(data: NDArray[Any]) -> Any: median = np.median(data) return np.median(np.abs(data - median)) -def max(data): +def max(data: NDArray[Any]) -> Any: return np.max(data) -def min(data): +def min(data: NDArray[Any]) -> Any: return np.min(data) -def percent_80(data): +def percent_80(data: NDArray[Any]) -> Any: return np.percentile(data, 80) -def percent_20(data): +def percent_20(data: NDArray[Any]) -> Any: return np.percentile(data, 20) -def avg_mean(data): +def avg_mean(data: NDArray[Any]) -> Any: percentile = np.percentile(data, 50) return np.mean(data[data > percentile]) -def mean(data): +def mean(data: NDArray[Any]) -> Any: return np.mean(data) diff --git a/src/badger/tests/conftest.py b/src/badger/tests/conftest.py index 1fc25c09..2f5041f7 100644 --- a/src/badger/tests/conftest.py +++ b/src/badger/tests/conftest.py @@ -1,12 +1,68 @@ import os import shutil +from collections.abc import Generator import pytest +from _pytest.fixtures import FixtureRequest from PyQt5.QtWidgets import QDialog +from pytest_mock import MockerFixture + +# Stash key holding the user's real plugin root so it can be restored at session end. +_OLD_PLUGIN_ROOT: pytest.StashKey[str | None] = pytest.StashKey() + + +def pytest_configure(config: pytest.Config) -> None: + # badger.factory captures BADGER_PLUGIN_ROOT at import time, which happens during + # collection. Point it at the mock plugins now, before any test module is imported. + from badger.settings import init_settings + + settings = init_settings() + try: + config.stash[_OLD_PLUGIN_ROOT] = settings.read_value("BADGER_PLUGIN_ROOT") + except KeyError: + config.stash[_OLD_PLUGIN_ROOT] = None + + mock_plugin_root = os.path.join(os.path.dirname(__file__), "mock", "plugins") + settings.write_value("BADGER_PLUGIN_ROOT", mock_plugin_root) + + +def pytest_unconfigure(config: pytest.Config) -> None: + old_plugin_root = config.stash.get(_OLD_PLUGIN_ROOT, None) + if old_plugin_root is not None: + from badger.settings import init_settings + + init_settings().write_value("BADGER_PLUGIN_ROOT", old_plugin_root) @pytest.fixture(autouse=True) -def suppress_popups(mocker): +def _ensure_qapp(qapp: object) -> None: + # Guarantee pytest-qt's QApplication exists before any widget is constructed, + # so tests that build Qt widgets without requesting qtbot don't abort when run + # in isolation (e.g. VS Code's Testing panel). + return None + + +@pytest.fixture(autouse=True) +def _reap_subprocesses() -> Generator[None, None, None]: + # Safety net: terminate any optimization subprocess a test leaves running. + # A leaked child inherits the pytest stdout pipe and keeps it open, so the + # test runner never sees EOF and hangs "indefinitely" (e.g. VS Code Testing). + import multiprocessing + + yield + + children = multiprocessing.active_children() + for child in children: + child.terminate() + for child in children: + child.join(timeout=2) + if child.is_alive(): + child.kill() + child.join(timeout=1) + + +@pytest.fixture(autouse=True) +def suppress_popups(mocker: MockerFixture) -> None: mocker.patch( "badger.gui.windows.expandable_message_box.ExpandableMessageBox.exec_", return_value=None, @@ -19,14 +75,14 @@ def suppress_popups(mocker): @pytest.fixture(scope="module", autouse=True) def config_test_settings( - mock_plugin_root, - mock_template_root, - mock_logbook_root, - mock_archive_root, - mock_log_directory, - mock_logging_level, - mock_temp_directory, -): + mock_template_root: str, + mock_logbook_root: str, + mock_archive_root: str, + mock_log_directory: str, + mock_logging_level: str, + mock_temp_directory: str, +) -> Generator[None, None, None]: + from badger.settings import init_settings config_singleton = init_settings() @@ -34,8 +90,8 @@ def config_test_settings( # Store the old values # If user's config is missing any values (for example if have older config missing newer added config options), # 'read_value' throws a KeyError which we can ignore + # BADGER_PLUGIN_ROOT is handled in pytest_configure (it must be set before collection). try: - old_root = config_singleton.read_value("BADGER_PLUGIN_ROOT") old_template = config_singleton.read_value("BADGER_TEMPLATE_ROOT") old_logbook = config_singleton.read_value("BADGER_LOGBOOK_ROOT") old_archived = config_singleton.read_value("BADGER_ARCHIVE_ROOT") @@ -46,7 +102,6 @@ def config_test_settings( pass # Assign values for test - config_singleton.write_value("BADGER_PLUGIN_ROOT", mock_plugin_root) config_singleton.write_value("BADGER_TEMPLATE_ROOT", mock_template_root) config_singleton.write_value("BADGER_LOGBOOK_ROOT", mock_logbook_root) config_singleton.write_value("BADGER_ARCHIVE_ROOT", mock_archive_root) @@ -57,7 +112,6 @@ def config_test_settings( # Restoring the original settings try: - config_singleton.write_value("BADGER_PLUGIN_ROOT", old_root) config_singleton.write_value("BADGER_TEMPLATE_ROOT", old_template) config_singleton.write_value("BADGER_LOGBOOK_ROOT", old_logbook) config_singleton.write_value("BADGER_ARCHIVE_ROOT", old_archived) @@ -71,12 +125,12 @@ def config_test_settings( @pytest.fixture(scope="module", autouse=True) def clean_up( - mock_template_root, - mock_logbook_root, - mock_archive_root, - mock_log_directory, - mock_temp_directory, -): + mock_template_root: str, + mock_logbook_root: str, + mock_archive_root: str, + mock_log_directory: str, + mock_temp_directory: str, +) -> Generator[None, None, None]: # Clean before tests shutil.rmtree(mock_template_root, True) # ignore errors shutil.rmtree(mock_logbook_root, True) @@ -95,45 +149,45 @@ def clean_up( @pytest.fixture(scope="module") -def mock_root(request): - return os.path.join(request.fspath.dirname, "mock") +def mock_root(request: FixtureRequest) -> str: + return os.path.join(request.path.parent, "mock") @pytest.fixture(scope="module") -def mock_plugin_root(mock_root): +def mock_plugin_root(mock_root: str) -> str: return os.path.join(mock_root, "plugins") @pytest.fixture(scope="module") -def mock_template_root(mock_root): +def mock_template_root(mock_root: str) -> str: return os.path.join(mock_root, "templates") @pytest.fixture(scope="module") -def mock_logbook_root(mock_root): +def mock_logbook_root(mock_root: str) -> str: return os.path.join(mock_root, "logbook") @pytest.fixture(scope="module") -def mock_archive_root(mock_root): +def mock_archive_root(mock_root: str) -> str: return os.path.join(mock_root, "archived") @pytest.fixture(scope="module") -def mock_log_directory(mock_root): +def mock_log_directory(mock_root: str) -> str: return os.path.join(mock_root, "logs") @pytest.fixture(scope="module") -def mock_temp_directory(mock_root): +def mock_temp_directory(mock_root: str) -> str: return os.path.join(mock_root, "temp") @pytest.fixture(scope="module") -def mock_logging_level(mock_root): +def mock_logging_level(mock_root: str) -> str: return "WARNING" @pytest.fixture(scope="module") -def mock_config_root(mock_root): +def mock_config_root(mock_root: str) -> str: return os.path.join(mock_root, "configs") diff --git a/src/badger/tests/mock/plugins/environments/multiobjective_test/__init__.py b/src/badger/tests/mock/plugins/environments/multiobjective_test/__init__.py index 2b074eb6..d648081c 100644 --- a/src/badger/tests/mock/plugins/environments/multiobjective_test/__init__.py +++ b/src/badger/tests/mock/plugins/environments/multiobjective_test/__init__.py @@ -13,7 +13,7 @@ class Environment(environment.Environment): flag: int = 0 - def set_variables(self, variable_inputs: dict[str, float]): + def set_variables(self, variable_inputs: dict[str, float]) -> None: if not self.interface: raise BadgerNoInterfaceError diff --git a/src/badger/tests/mock/plugins/environments/test/__init__.py b/src/badger/tests/mock/plugins/environments/test/__init__.py index 9d76c265..2db19d4f 100644 --- a/src/badger/tests/mock/plugins/environments/test/__init__.py +++ b/src/badger/tests/mock/plugins/environments/test/__init__.py @@ -15,7 +15,7 @@ class Environment(environment.Environment): flag: int = 0 delay: float = 0.0 - def set_variables(self, variable_inputs: dict[str, float]): + def set_variables(self, variable_inputs: dict[str, float]) -> None: if not self.interface: raise BadgerNoInterfaceError @@ -28,7 +28,7 @@ def set_variables(self, variable_inputs: dict[str, float]): self.interface.set_value("c", float((x**2).sum().numpy())) time.sleep(self.delay) - def get_bounds(self, variable_names): + def get_bounds(self, variable_names: list[str]) -> dict[str, list[float]]: """ Returns the bounds of new variables (not already included in env). """ diff --git a/src/badger/tests/mock/plugins/interfaces/test/__init__.py b/src/badger/tests/mock/plugins/interfaces/test/__init__.py index 8bae3784..7be34ffa 100644 --- a/src/badger/tests/mock/plugins/interfaces/test/__init__.py +++ b/src/badger/tests/mock/plugins/interfaces/test/__init__.py @@ -6,15 +6,15 @@ class Interface(interface.Interface): flag: int = 0 # Private variables - _states: dict + _states: dict[str, float] - def __init__(self, **data): + def __init__(self, **data: dict[str, float]) -> None: super().__init__(**data) self._states = {} @interface.log - def get_values(self, channel_names): + def get_values(self, channel_names: list[str]) -> dict[str, float]: channel_outputs = {} for channel in channel_names: @@ -28,6 +28,6 @@ def get_values(self, channel_names): return channel_outputs @interface.log - def set_values(self, channel_inputs): + def set_values(self, channel_inputs: dict[str, float]) -> None: for channel, value in channel_inputs.items(): self._states[channel] = value diff --git a/src/badger/tests/multiprocess_logging.py b/src/badger/tests/multiprocess_logging.py index a6afea19..c23a4f24 100644 --- a/src/badger/tests/multiprocess_logging.py +++ b/src/badger/tests/multiprocess_logging.py @@ -1,6 +1,7 @@ import logging import time -from multiprocessing import Process +from multiprocessing import Process, Queue +from pathlib import Path from badger.log import configure_process_logging, get_logging_manager @@ -19,7 +20,7 @@ """ -def simple_subprocess(log_queue): +def simple_subprocess(log_queue: Queue[logging.LogRecord]) -> None: configure_process_logging(log_queue=log_queue, log_level="DEBUG") sub_logger = logging.getLogger("badger.subprocess") @@ -29,7 +30,7 @@ def simple_subprocess(log_queue): time.sleep(0.5) -def test_multiprocessed_logging(tmp_path): +def test_multiprocessed_logging(tmp_path: Path) -> None: # deleted after each run log_filepath_1 = tmp_path / "test_multiprocess_1.log" log_filepath_2 = tmp_path / "test_multiprocess_2.log" @@ -116,7 +117,6 @@ def test_multiprocessed_logging(tmp_path): if __name__ == "__main__": import tempfile - from pathlib import Path with tempfile.TemporaryDirectory() as tmpdir: test_multiprocessed_logging(Path(tmpdir)) diff --git a/src/badger/tests/test_cli_basic.py b/src/badger/tests/test_cli_basic.py index c61d20e1..090ad9fa 100644 --- a/src/badger/tests/test_cli_basic.py +++ b/src/badger/tests/test_cli_basic.py @@ -2,15 +2,15 @@ from importlib import metadata -def capture(command): +def capture(command: list[str]) -> tuple[str, str, int]: proc = subprocess.Popen(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE) - out, err = proc.communicate() - out = out.decode("utf-8") - err = err.decode("utf-8") + _out, _err = proc.communicate() + out = _out.decode("utf-8") + err = _err.decode("utf-8") return out, err, proc.returncode -def test_cli_main(): +def test_cli_main() -> None: command = ["badger"] out, _, exitcode = capture(command) @@ -32,7 +32,7 @@ def test_cli_main(): assert outlines[1] == f"version: {version}" -def test_list_algo(): +def test_list_algo() -> None: from badger.factory import ALGO_EXCLUDED command = ["badger", "generator"] @@ -45,20 +45,3 @@ def test_list_algo(): for algo in ["expected_improvement", "neldermead", "rcds"]: if algo not in ALGO_EXCLUDED: assert f"- {algo}" in outlines - - -# def test_cli_run(mock_config_root): -# command = ['badger', 'run', '-a', 'upper_confidence_bound', '-ap', -# '{max_evaluations: 10}', '-e', 'silly', '-c', -# os.path.join(mock_config_root, 'test.yaml'), '-y'] -# out, err, exitcode = capture(command) - -# assert exitcode == 0 - -# # Check output lines -# outlines = out.splitlines() -# assert len(outlines) == 15 - -# # Check table header -# assert outlines[1] == '| iter | l2 ' \ -# + '| q1 | q2 |' diff --git a/src/badger/tests/test_core.py b/src/badger/tests/test_core.py index bc2127e4..b89deba3 100644 --- a/src/badger/tests/test_core.py +++ b/src/badger/tests/test_core.py @@ -12,7 +12,7 @@ class TestCore: """ @pytest.fixture(autouse=True, scope="function") - def test_core_setup(self, *args, **kwargs) -> None: + def test_core_setup(self, *args: object, **kwargs: object) -> None: super().__init__(*args, **kwargs) self.count = 0 self.candidates = None diff --git a/src/badger/tests/test_core_subprocess.py b/src/badger/tests/test_core_subprocess.py index 5e20cbf6..db53cc78 100644 --- a/src/badger/tests/test_core_subprocess.py +++ b/src/badger/tests/test_core_subprocess.py @@ -1,9 +1,14 @@ import multiprocessing import time +from collections.abc import Generator import pandas as pd import pytest +from badger.gui.components.process_manager import ProcessManager +from badger.gui.components.routine_runner import ArgumentQueueType +from badger.types import TerminationConditionConfig + class TestCore: """ @@ -11,7 +16,7 @@ class TestCore: """ @pytest.fixture - def process_manager(self): + def process_manager(self) -> Generator[ProcessManager, None, None]: from badger.gui.components.create_process import CreateProcess from badger.gui.components.process_manager import ProcessManager @@ -22,10 +27,10 @@ def process_manager(self): yield process_manager - process_manager.close_proccesses() + process_manager.close_processes() @pytest.fixture(scope="session") - def init_multiprocessing(self): + def init_multiprocessing(self) -> None: # Use 'spawn' on Windows, 'fork' on Unix-like systems method = ( "spawn" @@ -35,7 +40,7 @@ def init_multiprocessing(self): multiprocessing.set_start_method(method, force=True) @pytest.fixture(autouse=True, scope="function") - def test_core_setup(self, *args, **kwargs) -> None: + def test_core_setup(self, *args: object, **kwargs: object) -> None: super().__init__(*args, **kwargs) self.count = 0 self.candidates = None @@ -56,9 +61,8 @@ def test_core_setup(self, *args, **kwargs) -> None: self.points_eval_target = pd.DataFrame(data_eval_target) - def test_run_routine_subprocess( - self, process_manager, init_multiprocessing - ) -> None: + @pytest.mark.usefixtures("init_multiprocessing") + def test_run_routine_subprocess(self, process_manager: ProcessManager) -> None: """ A unit test to ensure the core functionality of run_routine_xopt is functioning as intended. @@ -72,27 +76,32 @@ def test_run_routine_subprocess( self.routine = create_routine() time.sleep(1) tmp_filename = save_tmp_run(self.routine) - self.termination_condition = { - "tc_idx": 0, - "max_eval": 3, - } + self.termination_condition = TerminationConditionConfig( + tc_idx=0, + max_eval=3, + ) process_with_args = process_manager.remove_from_queue() - pause_event = process_with_args["pause_event"] - args_queue = process_with_args["args_queue"] - wait_event = process_with_args["wait_event"] - routine_process = process_with_args["process"] - evaluate_queue = process_with_args["evaluate_queue"] - - arg_dict = { - "routine_id": self.routine.id, - "routine_filename": tmp_filename, - "routine_name": self.routine.name, - "variable_ranges": self.routine.vocs.variables, - "initial_points": self.routine.initial_points, - "evaluate": True, - "termination_condition": self.termination_condition, - "start_time": time.time(), - } + assert process_with_args is not None + pause_event = process_with_args.pause_event + args_queue = process_with_args.args_queue + wait_event = process_with_args.wait_event + routine_process = process_with_args.process + evaluate_queue = process_with_args.evaluate_queue + + arg_dict = ArgumentQueueType( + routine_id=self.routine.id, + routine_filename=tmp_filename, + routine_name=self.routine.name, + variable_ranges=self.routine.vocs.variables, + initial_points=self.routine.initial_points, + evaluate=True, + archive=True, + termination_condition=self.termination_condition, + start_time=time.time(), + testing=True, + run_data=True, + init_points=self.routine.initial_points, + ) args_queue.put(arg_dict) wait_event.set() @@ -142,59 +151,3 @@ def test_evaluate_points(self) -> None: .astype(float) .equals(self.points_eval_target.astype(float)) ) - - def test_run_turbo(self, process_manager, init_multiprocessing) -> None: - """ - A unit test to ensure TuRBO can run in Badger. - """ - return - - from badger.db import save_routine - from badger.tests.utils import create_routine_turbo - - self.count = 0 - self.num_of_points = 3 - self.routine = create_routine_turbo() - time.sleep(1) - assert self.routine.generator.turbo_controller.best_value is None - save_routine(self.routine) - self.termination_condition = { - "tc_idx": 0, - "max_eval": 3, - } - process_with_args = process_manager.remove_from_queue() - pause_event = process_with_args["pause_event"] - args_queue = process_with_args["args_queue"] - wait_event = process_with_args["wait_event"] - routine_process = process_with_args["process"] - evaluate_queue = process_with_args["evaluate_queue"] - - arg_dict = { - "routine_id": self.routine.id, - "evaluate": True, - "termination_condition": self.termination_condition, - "start_time": time.time(), - } - - args_queue.put(arg_dict) - wait_event.set() - pause_event.set() - - time.sleep(0.20) - routine_process.terminate() - time.sleep(1) - - while evaluate_queue[1].poll(): - self.results = evaluate_queue[1].recv() - - # assert len(self.candidates_list) == self.count - 1 - - # assert len(self.results) == self.num_of_points - - assert self.states is None - - """ - path = "./test.yaml" - assert os.path.exists(path) is True - os.remove("./test.yaml") - """ diff --git a/src/badger/tests/test_create_process.py b/src/badger/tests/test_create_process.py index 02fcfddf..e2118d0e 100644 --- a/src/badger/tests/test_create_process.py +++ b/src/badger/tests/test_create_process.py @@ -1,21 +1,29 @@ import multiprocessing as mp import sys +from dataclasses import fields +from typing import TYPE_CHECKING from unittest.mock import MagicMock, patch import pytest from PyQt5.QtWidgets import QApplication +from pytestqt.qtbot import QtBot + +if TYPE_CHECKING: + from badger.gui.components.create_process import CreateProcess app = QApplication(sys.argv) @pytest.fixture -def process_creator(): +def process_creator() -> "CreateProcess": from badger.gui.components.create_process import CreateProcess return CreateProcess() -def test_create_subprocess_emits_signals(qtbot, process_creator): +def test_create_subprocess_emits_signals( + qtbot: QtBot, process_creator: "CreateProcess" +) -> None: with ( patch("badger.gui.components.create_process.Process") as mock_process, patch("badger.gui.components.create_process.run_routine_subprocess"), @@ -34,9 +42,9 @@ def test_create_subprocess_emits_signals(qtbot, process_creator): assert blocker_finished.signal_triggered mock_process.assert_called_once() - # Verify that the emitted object contains the expected keys + # Verify that the emitted object exposes the expected attributes emitted_args = blocker_subprocess_prepared.args[0] - assert set(emitted_args.keys()) == { + assert {f.name for f in fields(emitted_args)} == { "process", "args_queue", "stop_event", @@ -47,10 +55,10 @@ def test_create_subprocess_emits_signals(qtbot, process_creator): "dialog_action_queue", } - assert isinstance(emitted_args["args_queue"], mp.queues.Queue) - assert isinstance(emitted_args["data_queue"], mp.queues.Queue) - assert isinstance(emitted_args["dialog_action_queue"], mp.queues.Queue) - assert isinstance(emitted_args["evaluate_queue"], tuple) - assert isinstance(emitted_args["wait_event"], mp.synchronize.Event) - assert isinstance(emitted_args["stop_event"], mp.synchronize.Event) - assert isinstance(emitted_args["pause_event"], mp.synchronize.Event) + assert isinstance(emitted_args.args_queue, mp.queues.Queue) + assert isinstance(emitted_args.data_queue, mp.queues.Queue) + assert isinstance(emitted_args.dialog_action_queue, mp.queues.Queue) + assert isinstance(emitted_args.evaluate_queue, tuple) + assert isinstance(emitted_args.wait_event, mp.synchronize.Event) + assert isinstance(emitted_args.stop_event, mp.synchronize.Event) + assert isinstance(emitted_args.pause_event, mp.synchronize.Event) diff --git a/src/badger/tests/test_env.py b/src/badger/tests/test_env.py index 6f783f4e..eb12559b 100644 --- a/src/badger/tests/test_env.py +++ b/src/badger/tests/test_env.py @@ -1,7 +1,7 @@ import pytest -def test_find_env(): +def test_find_env() -> None: from badger.factory import get_env, list_env assert len(list_env()) == 2 @@ -11,7 +11,7 @@ def test_find_env(): assert configs["version"] == "1.0" -def test_get_params(): +def test_get_params() -> None: from badger.factory import get_env _, configs = get_env("test") @@ -20,7 +20,7 @@ def test_get_params(): assert params == {"delay": 0.0, "flag": 0} -def test_set_params(): +def test_set_params() -> None: from badger.factory import get_env, get_intf Interface, _ = get_intf("test") @@ -39,7 +39,7 @@ def test_set_params(): assert env.flag == 0 -def test_list_variables(): +def test_list_variables() -> None: from badger.factory import get_env Environment, _ = get_env("test") @@ -51,7 +51,7 @@ def test_list_variables(): assert env.variable_names == list(variables.keys()) -def test_list_observables(): +def test_list_observables() -> None: from badger.factory import get_env Environment, _ = get_env("test") @@ -60,7 +60,7 @@ def test_list_observables(): assert env.observables == ["f", "c"] -def test_get_variables(): +def test_get_variables() -> None: from badger.errors import ( BadgerNoInterfaceError, ) @@ -92,7 +92,7 @@ def test_get_variables(): # assert "Getting them through interface is not allowed" in str(e.value) -def testset_variables(): +def testset_variables() -> None: from badger.errors import ( BadgerEnvVarError, BadgerNoInterfaceError, @@ -140,7 +140,7 @@ def testset_variables(): assert variable_outputs == {"x1": 1, "x2": -1} # values shouldn't change -def testget_observables(): +def testget_observables() -> None: from badger.errors import BadgerNoInterfaceError from badger.factory import get_env, get_intf diff --git a/src/badger/tests/test_environment.py b/src/badger/tests/test_environment.py index ffeeabc4..7931597f 100644 --- a/src/badger/tests/test_environment.py +++ b/src/badger/tests/test_environment.py @@ -7,6 +7,7 @@ from badger.environment import BaseEnvironment, Environment from badger.errors import BadgerEnvVarError, BadgerNoInterfaceError +from badger.factory import BadgerPluginConfig from badger.interface import Interface # TEST_VOCS_BASE replacement for testing @@ -19,12 +20,12 @@ class TestEnvironment: """Test suite for the Environment classes and decorators.""" - def test_base_environment_abstract(self): + def test_base_environment_abstract(self) -> None: """Test that BaseEnvironment cannot be instantiated directly.""" with pytest.raises(TypeError): BaseEnvironment() - def test_basic_env(self): + def test_basic_env(self) -> None: """Test basic environment functionality with custom implementation.""" class TestEnv(BaseEnvironment): @@ -39,10 +40,12 @@ class TestEnv(BaseEnvironment): def get_variables(self, variable_names: list[str]) -> dict[str, float]: return {name: 0.5 for name in variable_names} - def set_variables(self, variable_inputs: dict[str, float]): - pass + def set_variables(self, variable_inputs: dict[str, float]) -> None: + return None - def get_observables(self, observable_names: list[str]) -> dict[str, float]: + def get_observables( + self, observable_names: list[str] + ) -> dict[str, float | list[float]]: return {ele: 1.0 for ele in observable_names} env = TestEnv() @@ -56,7 +59,7 @@ def get_observables(self, observable_names: list[str]) -> dict[str, float]: result = json.loads(env.model_dump_json()) assert result["my_flag"] == 1 - def test_environment_with_interface(self): + def test_environment_with_interface(self) -> None: """Test Environment class with interface.""" # Create a mock interface mock_interface = Mock(spec=Interface) @@ -88,7 +91,7 @@ class TestEnv(Environment): env.reset_environment() mock_interface.reset_interface.assert_called_once() - def test_environment_no_interface_error(self): + def test_environment_no_interface_error(self) -> None: """Test that Environment raises error when no interface is provided.""" class TestEnv(Environment): @@ -107,7 +110,7 @@ class TestEnv(Environment): with pytest.raises(BadgerNoInterfaceError): env.get_observables(["f"]) - def test_bounds_validation(self): + def test_bounds_validation(self) -> None: """Test bounds validation decorator.""" class TestEnv(BaseEnvironment): @@ -135,7 +138,7 @@ def get_bounds(self, variable_names: list[str]) -> dict[str, float]: ): env.get_bounds(["x1"]) - def test_setpoint_validation(self): + def test_setpoint_validation(self) -> None: """Test setpoint validation decorator.""" class TestEnv(BaseEnvironment): @@ -168,7 +171,7 @@ def get_observables(self, observable_names: list[str]) -> dict[str, float]: ): env.set_variables({"x2": -1.0}) # Outside lower bound - def test_formula_processing(self): + def test_formula_processing(self) -> None: """Test formula processing decorator for observables.""" class TestEnv(BaseEnvironment): @@ -202,7 +205,7 @@ def get_observables(self, observable_names: list[str]) -> dict[str, float]: # Should not include the formula variables in output assert "f" not in result - def test_convenience_methods(self): + def test_convenience_methods(self) -> None: """Test convenience methods for single variable/observable operations.""" class TestEnv(BaseEnvironment): @@ -231,7 +234,7 @@ def get_observables(self, observable_names: list[str]) -> dict[str, float]: # Test get_observable assert env.get_observable("f") == 1.5 - def test_variable_names_property(self): + def test_variable_names_property(self) -> None: """Test variable_names property for Environment class.""" mock_interface = Mock(spec=Interface) @@ -247,7 +250,7 @@ class TestEnv(Environment): env = TestEnv(interface=mock_interface) assert set(env.variable_names) == {"x1", "x2", "x3"} - def test_get_system_states(self): + def test_get_system_states(self) -> None: """Test get_system_states default implementation.""" class TestEnv(BaseEnvironment): @@ -267,7 +270,7 @@ def get_observables(self, observable_names: list[str]) -> dict[str, float]: env = TestEnv() assert env.get_system_states() == {} - def test_search_not_implemented(self): + def test_search_not_implemented(self) -> None: """Test search method raises NotImplementedError by default.""" class TestEnv(BaseEnvironment): @@ -291,7 +294,7 @@ def get_observables(self, observable_names: list[str]) -> dict[str, float]: ): env.search("test", lambda x: x) - def test_env_in_routine(self): + def test_env_in_routine(self) -> None: """Test environment integration with routine.""" from badger.routine import Routine @@ -326,7 +329,7 @@ def get_observables(self, observable_names: list[str]) -> dict[str, float]: result = json.loads(routine.json()) assert result["environment"]["my_flag"] == 1 - def test_env_from_get_env(self): + def test_env_from_get_env(self) -> None: """Test environment from factory function.""" from badger.factory import get_env from badger.routine import Routine @@ -348,7 +351,7 @@ def test_env_from_get_env(self): result = json.loads(routine.json()) assert result["environment"]["flag"] == 1 - def test_complex_formula_processing(self): + def test_complex_formula_processing(self) -> None: """Test complex formula processing with numpy functions.""" class TestEnv(BaseEnvironment): @@ -392,7 +395,7 @@ def get_observables(self, observable_names: list[str]) -> dict[str, float]: result["sqrt(mean(`data`**2)) + percentile95(`signal`)"], expected ) - def test_instantiate_env_function(self): + def test_instantiate_env_function(self) -> None: """Test environment instantiation function.""" from badger.environment import instantiate_env from badger.interface import Interface @@ -411,7 +414,7 @@ class TestEnv(Environment): original_get_intf = getattr(badger.factory, "get_intf", None) - def mock_get_intf(name): + def mock_get_intf(name: str) -> tuple[Mock, BadgerPluginConfig]: mock_interface_class = Mock() mock_interface_class.return_value = Mock(spec=Interface) return mock_interface_class, {} diff --git a/src/badger/tests/test_factory.py b/src/badger/tests/test_factory.py index c9ebf368..869c601a 100644 --- a/src/badger/tests/test_factory.py +++ b/src/badger/tests/test_factory.py @@ -7,7 +7,7 @@ class TestFactory: - def test_generator_generation(self): + def test_generator_generation(self) -> None: from badger.factory import list_generators generator_names = list_generators() diff --git a/src/badger/tests/test_formulas.py b/src/badger/tests/test_formulas.py index ff1da129..dea3b5fa 100644 --- a/src/badger/tests/test_formulas.py +++ b/src/badger/tests/test_formulas.py @@ -1,3 +1,5 @@ +from typing import Any + import numpy as np import pytest @@ -13,29 +15,29 @@ class TestSafeVarName: """Test the safe_var_name function.""" - def test_alphanumeric_unchanged(self): + def test_alphanumeric_unchanged(self) -> None: """Test that alphanumeric names with underscores remain unchanged.""" assert safe_var_name("valid_name123") == "valid_name123" assert safe_var_name("ABC_def_456") == "ABC_def_456" assert safe_var_name("_private") == "_private" - def test_special_characters_replaced(self): + def test_special_characters_replaced(self) -> None: """Test that special characters are replaced with underscores.""" assert safe_var_name("test-name") == "test_name" assert safe_var_name("test.name") == "test_name" assert safe_var_name("test name") == "test_name" assert safe_var_name("test@name#") == "test_name_" - def test_multiple_special_chars(self): + def test_multiple_special_chars(self) -> None: """Test handling of multiple special characters.""" assert safe_var_name("test-@#$%name") == "test_____name" assert safe_var_name("!@#$%^&*()") == "__________" - def test_empty_string(self): + def test_empty_string(self) -> None: """Test handling of empty string.""" assert safe_var_name("") == "" - def test_numbers_only(self): + def test_numbers_only(self) -> None: """Test handling of numeric strings.""" assert safe_var_name("123") == "123" assert safe_var_name("456.789") == "456_789" @@ -44,13 +46,13 @@ def test_numbers_only(self): class TestExtractVariableKeys: """Test the extract_variable_keys function.""" - def test_backticks(self): + def test_backticks(self) -> None: """Test extraction of backtick-quoted variables.""" expr = "`var1` + `var2`" result = extract_variable_keys(expr) assert result == ["var1", "var2"] - def test_duplicate_variables(self): + def test_duplicate_variables(self) -> None: """Test that duplicate variables are handled correctly.""" expr = "`var1` + `var1` + `var1`" result = extract_variable_keys(expr) @@ -58,25 +60,27 @@ def test_duplicate_variables(self): "var1", "var1", "var1", - ] # findall returns list with duplicates def test_no_variables(self): + ] # findall returns list with duplicates + + def test_no_variables(self) -> None: """Test expression with no quoted variables.""" expr = "x + y + 5" result = extract_variable_keys(expr) assert result == [] - def test_empty_quotes(self): + def test_empty_quotes(self) -> None: """Test handling of empty quotes.""" expr = "``" result = extract_variable_keys(expr) assert result == [] # Empty backticks don't match the regex - def test_special_characters_in_quotes(self): + def test_special_characters_in_quotes(self) -> None: """Test variables with special characters.""" expr = "`var-1` + `var.2` + `var@3`" result = extract_variable_keys(expr) assert result == ["var-1", "var.2", "var@3"] - def test_nested_backticks_not_supported(self): + def test_nested_backticks_not_supported(self) -> None: """Test that nested backticks are not properly handled.""" expr = "`var`1` + `var2`" result = extract_variable_keys(expr) @@ -87,35 +91,35 @@ def test_nested_backticks_not_supported(self): class TestFindUsedNames: """Test the find_used_names function.""" - def test_simple_variables(self): + def test_simple_variables(self) -> None: """Test finding simple variable names.""" expr = "x + y" result = find_used_names(expr) assert result == {"x", "y"} - def test_function_calls(self): + def test_function_calls(self) -> None: """Test finding function names.""" expr = "sin(x) + cos(y)" result = find_used_names(expr) assert result == {"sin", "cos", "x", "y"} - def test_complex_expression(self): + def test_complex_expression(self) -> None: """Test complex expression with various names.""" expr = "sqrt(mean(x**2)) + percentile(data, 95)" result = find_used_names(expr) assert result == {"sqrt", "mean", "x", "percentile", "data"} - def test_invalid_syntax(self): + def test_invalid_syntax(self) -> None: """Test that SyntaxError is raised for invalid expressions.""" with pytest.raises(SyntaxError, match="Invalid syntax in expression"): find_used_names("x +++") # Use a clearly invalid syntax - def test_empty_expression(self): + def test_empty_expression(self) -> None: """Test empty expression.""" with pytest.raises(SyntaxError): find_used_names("") - def test_literals_only(self): + def test_literals_only(self) -> None: """Test expression with only literals.""" expr = "5 + 3.14" result = find_used_names(expr) @@ -125,7 +129,7 @@ def test_literals_only(self): class TestSuggestName: """Test the suggest_name function.""" - def test_close_match_found(self): + def test_close_match_found(self) -> None: """Test finding close matches for unknown names.""" unknown = ["sine", "cosine"] known = ["sin", "cos", "tan", "exp"] @@ -133,14 +137,14 @@ def test_close_match_found(self): assert "sine" in result assert result["sine"] == "sin" - def test_no_close_match(self): + def test_no_close_match(self) -> None: """Test when no close match is found.""" unknown = ["xyz"] known = ["sin", "cos", "tan"] result = suggest_name(unknown, known) assert result == {} - def test_exact_match_not_suggested(self): + def test_exact_match_not_suggested(self) -> None: """Test that exact matches are still suggested (difflib behavior).""" unknown = ["sin"] known = ["sin", "cos", "tan"] @@ -148,7 +152,7 @@ def test_exact_match_not_suggested(self): # difflib actually returns exact matches too assert result == {"sin": "sin"} - def test_multiple_suggestions(self): + def test_multiple_suggestions(self) -> None: """Test multiple suggestions.""" unknown = ["sine", "cosine", "xyz"] known = ["sin", "cos", "tan"] @@ -156,7 +160,7 @@ def test_multiple_suggestions(self): assert len(result) <= 2 # xyz should not match assert "sine" in result - def test_empty_inputs(self): + def test_empty_inputs(self) -> None: """Test empty inputs.""" assert suggest_name([], ["sin", "cos"]) == {} assert suggest_name(["sine"], []) == {} @@ -166,21 +170,21 @@ def test_empty_inputs(self): class TestInterpretExpression: """Test the interpret_expression function.""" - def test_simple_arithmetic(self): + def test_simple_arithmetic(self) -> None: """Test simple arithmetic with variables.""" expr = "`x` + `y`" variables = {"x": 5, "y": 3} result = interpret_expression(expr, variables) assert result == 8 - def test_numpy_functions(self): + def test_numpy_functions(self) -> None: """Test using numpy functions.""" expr = "sin(`x`) + cos(`y`)" variables = {"x": 0, "y": 0} result = interpret_expression(expr, variables) assert np.isclose(result, 1.0) # sin(0) + cos(0) = 0 + 1 = 1 - def test_percentile_function(self): + def test_percentile_function(self) -> None: """Test the percentile function transformation.""" expr = "percentile90(`data`)" variables = {"data": np.array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10])} @@ -188,7 +192,7 @@ def test_percentile_function(self): expected = np.percentile(variables["data"], 90) assert np.isclose(result, expected) - def test_rms_function(self): + def test_rms_function(self) -> None: """Test the RMS function transformation.""" expr = "rms(`data`)" variables = {"data": np.array([3, 4])} @@ -196,21 +200,21 @@ def test_rms_function(self): expected = np.sqrt(np.mean(variables["data"] ** 2)) assert np.isclose(result, expected) - def test_backtick_variables(self): + def test_backtick_variables(self) -> None: """Test with backtick-quoted variables.""" expr = "`x` + `y` + `z`" variables = {"x": 1, "y": 2, "z": 3} result = interpret_expression(expr, variables) assert result == 6 - def test_special_character_variables(self): + def test_special_character_variables(self) -> None: """Test variables with special characters.""" expr = "`var-1` + `var.2`" variables = {"var-1": 10, "var.2": 20} result = interpret_expression(expr, variables) assert result == 30 - def test_missing_variables_error(self): + def test_missing_variables_error(self) -> None: """Test error when variables are missing.""" expr = "`x` + `y`" variables = {"x": 5} # missing 'y' @@ -219,28 +223,28 @@ def test_missing_variables_error(self): ): interpret_expression(expr, variables) - def test_unknown_names_error(self): + def test_unknown_names_error(self) -> None: """Test error for unknown function/variable names.""" expr = "unknown_func(`x`)" variables = {"x": 5} with pytest.raises(NameError, match="Unknown names in expression"): interpret_expression(expr, variables) - def test_unknown_names_with_suggestions(self): + def test_unknown_names_with_suggestions(self) -> None: """Test that suggestions are provided for unknown names.""" expr = "sine(`x`)" # should suggest 'sin' variables = {"x": 0} with pytest.raises(NameError, match="Did you mean"): interpret_expression(expr, variables) - def test_invalid_expression_syntax(self): + def test_invalid_expression_syntax(self) -> None: """Test error for invalid expression syntax.""" expr = "`x` +" # incomplete expression variables = {"x": 5} with pytest.raises(SyntaxError, match="Invalid syntax in expression"): interpret_expression(expr, variables) - def test_complex_expression(self): + def test_complex_expression(self) -> None: """Test a complex expression with multiple operations.""" expr = "sqrt(mean(`data`**2)) + percentile95(`values`)" variables = { @@ -256,14 +260,14 @@ def test_complex_expression(self): assert np.isclose(result, expected) - def test_array_operations(self): + def test_array_operations(self) -> None: """Test operations with numpy arrays.""" expr = "sum(`arr`) / len(`arr`)" variables = {"arr": np.array([1, 2, 3, 4, 5])} result = interpret_expression(expr, variables) assert np.isclose(result, 3.0) # mean of [1,2,3,4,5] - def test_nested_functions(self): + def test_nested_functions(self) -> None: """Test nested function calls.""" expr = "sqrt(abs(sin(`x`)))" variables = {"x": -np.pi / 2} @@ -271,7 +275,7 @@ def test_nested_functions(self): expected = np.sqrt(np.abs(np.sin(-np.pi / 2))) assert np.isclose(result, expected) - def test_multiple_percentile_calls(self): + def test_multiple_percentile_calls(self) -> None: """Test multiple percentile function calls.""" expr = "percentile25(`data`) + percentile75(`data`)" variables = {"data": np.array([1, 2, 3, 4, 5])} @@ -281,10 +285,10 @@ def test_multiple_percentile_calls(self): ) assert np.isclose(result, expected) - def test_empty_variable_name(self): + def test_empty_variable_name(self) -> None: """Test handling of empty variable names.""" expr = "42" # Simple constant expression instead of empty backticks - variables = {} + variables: dict[str, Any] = {} result = interpret_expression(expr, variables) assert result == 42 @@ -293,7 +297,7 @@ def test_empty_variable_name(self): class TestIntegration: """Integration tests combining multiple functions.""" - def test_full_workflow(self): + def test_full_workflow(self) -> None: """Test the full workflow of expression interpretation.""" # This tests how all functions work together expr = "`sensor_1` + `sensor-2` * sin(`angle`)" @@ -307,7 +311,7 @@ def test_full_workflow(self): expected = 10 + 5 * np.sin(np.pi / 6) # 10 + 5 * 0.5 = 12.5 assert np.isclose(result, expected) - def test_error_handling_workflow(self): + def test_error_handling_workflow(self) -> None: """Test error handling across the workflow.""" expr = "`missing_var` + unknown_func(`x`)" variables = {"x": 5} diff --git a/src/badger/tests/test_gui_basic.py b/src/badger/tests/test_gui_basic.py index bfd3711a..18955278 100644 --- a/src/badger/tests/test_gui_basic.py +++ b/src/badger/tests/test_gui_basic.py @@ -3,198 +3,218 @@ from unittest.mock import patch import pytest -from PyQt5.QtCore import QEventLoop, Qt, QTimer +from PyQt5.QtCore import QEventLoop, QTimer +from PyQt5.QtGui import QShowEvent from PyQt5.QtWidgets import QDialog +from pytestqt.qtbot import QtBot @pytest.fixture(scope="session") -def init_multiprocessing(): +def init_multiprocessing() -> None: multiprocessing.set_start_method("fork", force=True) -def test_gui_main(qtbot, init_multiprocessing): +@pytest.mark.usefixtures("init_multiprocessing") +def test_gui_main(qtbot: QtBot) -> None: from badger.gui.windows.main_window import BadgerMainWindow from badger.tests.utils import fix_path_issues fix_path_issues() window = BadgerMainWindow() + try: + while window.thread_list: + loop = QEventLoop() + QTimer.singleShot(1000, loop.quit) # 1000 ms pause + loop.exec_() + time.sleep(1) + print("test", print(window.thread_list)) + + qtbot.addWidget(window) - while window.thread_list: loop = QEventLoop() - QTimer.singleShot(1000, loop.quit) # 1000 ms pause + QTimer.singleShot(3000, loop.quit) # 1000 ms pause loop.exec_() - time.sleep(1) - print("test", print(window.thread_list)) - - qtbot.addWidget(window) - loop = QEventLoop() - QTimer.singleShot(3000, loop.quit) # 1000 ms pause - loop.exec_() + # TODO: Test if generator tab has been filled + # assert window.stacks.currentWidget().tabs.currentIndex() == 1 + finally: + window.process_manager.close_processes() - # TODO: Test if generator tab has been filled - # assert window.stacks.currentWidget().tabs.currentIndex() == 1 - window.process_manager.close_proccesses() - - -def test_close_main(qtbot, init_multiprocessing): +@pytest.mark.usefixtures("init_multiprocessing") +def test_close_main(qtbot: QtBot) -> None: from badger.archive import save_tmp_run from badger.gui.windows.main_window import BadgerMainWindow from badger.tests.utils import create_routine, fix_path_issues + from badger.types import TerminationConditionConfig fix_path_issues() window = BadgerMainWindow() + try: + qtbot.addWidget(window) - qtbot.addWidget(window) + loop = QEventLoop() + QTimer.singleShot(5000, loop.quit) # 1000 ms pause + loop.exec_() - loop = QEventLoop() - QTimer.singleShot(5000, loop.quit) # 1000 ms pause - loop.exec_() + routine = create_routine() + home_page = window.home_page + home_page.current_routine = routine + tmp_filename = save_tmp_run(routine) + home_page.run_monitor.testing = True + home_page.run_monitor.routine_filename = tmp_filename + home_page.run_monitor.termination_condition = TerminationConditionConfig( + tc_idx=0, + max_eval=3, + ) + home_page.go_run(-1) + home_page.run_monitor.start(True) - routine = create_routine() - home_page = window.home_page - home_page.current_routine = routine - tmp_filename = save_tmp_run(routine) - home_page.run_monitor.testing = True - home_page.run_monitor.routine_filename = tmp_filename - home_page.run_monitor.termination_condition = { - "tc_idx": 0, - "max_eval": 3, - } - home_page.go_run(-1) - home_page.run_monitor.start(True) + loop = QEventLoop() + QTimer.singleShot(1000, loop.quit) # 1000 ms pause + loop.exec_() - loop = QEventLoop() - QTimer.singleShot(1000, loop.quit) # 1000 ms pause - loop.exec_() + window.close() # this action should release the env - window.close() # this action should release the env - # So we expect an AttributeError here - with pytest.raises(AttributeError): - _ = home_page.run_monitor.routine.environment + assert home_page.run_monitor.routine is not None - window.process_manager.close_proccesses() + # So we expect an AttributeError here + with pytest.raises(AttributeError): + _ = home_page.run_monitor.routine.environment + finally: + window.process_manager.close_processes() -def test_traceback_during_run(qtbot, init_multiprocessing): +@pytest.mark.usefixtures("init_multiprocessing") +def test_traceback_during_run(qtbot: QtBot, monkeypatch: pytest.MonkeyPatch) -> None: with patch("badger.core.run_routine") as run_routine_mock: run_routine_mock.side_effect = Exception("Test exception") from badger.gui.windows.main_window import BadgerMainWindow from badger.gui.windows.message_dialog import BadgerScrollableMessageBox from badger.tests.utils import create_routine, fix_path_issues + from badger.types import TerminationConditionConfig fix_path_issues() window = BadgerMainWindow() - - loop = QEventLoop() - QTimer.singleShot(1000, loop.quit) # 1000 ms pause - loop.exec_() - - qtbot.addWidget(window) - - routine = create_routine() - home_page = window.home_page - home_page.current_routine = routine - home_page.run_monitor.testing = True - home_page.run_monitor.termination_condition = { - "tc_idx": 0, - "max_eval": 3, - } - home_page.go_run(-1) - - # Function to replace the original showEvent - def patched_showEvent(original_showEvent): - def inner(ins, event): - original_showEvent(ins, event) # Call the original showEvent + try: + loop = QEventLoop() + QTimer.singleShot(1000, loop.quit) # 1000 ms pause + loop.exec_() + + qtbot.addWidget(window) + + routine = create_routine() + home_page = window.home_page + home_page.current_routine = routine + home_page.run_monitor.testing = True + home_page.run_monitor.termination_condition = TerminationConditionConfig( + tc_idx=0, + max_eval=3, + ) + home_page.go_run(-1) + + # Patch showEvent to auto-accept the traceback dialog when it appears. + original_show_event = BadgerScrollableMessageBox.showEvent + + def patched_show_event( + ins: BadgerScrollableMessageBox, event: QShowEvent | None + ) -> None: + original_show_event(ins, event) # Call the original showEvent assert ins # make sure the dialog is created assert ins.detailedTextWidget.toPlainText() # make sure it's not empty QTimer.singleShot(100, ins.accept) # Close the dialog after 100 ms - return inner - - BadgerScrollableMessageBox.showEvent = patched_showEvent( - BadgerScrollableMessageBox.showEvent - ) - - home_page.run_monitor.start(True) - # Wait until the run is done - while home_page.run_monitor.running: - qtbot.wait(100) + monkeypatch.setattr( + BadgerScrollableMessageBox, "showEvent", patched_show_event + ) - window.process_manager.close_proccesses() + home_page.run_monitor.start(True) + # Wait until the run is done + while home_page.run_monitor.running: + qtbot.wait(100) + finally: + window.process_manager.close_processes() -def test_measurement_retry_dialog_in_app_flow(qtbot, init_multiprocessing): +def test_measurement_retry_dialog_in_app_flow(qtbot: QtBot) -> None: from badger.gui.windows.main_window import BadgerMainWindow from badger.tests.utils import create_routine, fix_path_issues + from badger.types import ( + MeasurementActionMessage, + MeasurementErrorMessage, + TerminationConditionConfig, + ) fix_path_issues() window = BadgerMainWindow() - qtbot.addWidget(window) - - loop = QEventLoop() - QTimer.singleShot(1000, loop.quit) - loop.exec_() - - routine = create_routine() - home_page = window.home_page - home_page.current_routine = routine - home_page.go_run(-1) - home_page.run_monitor.init_routine_runner() - runner = home_page.run_monitor.routine_runner - - # Grab queue and process from the process manager - process_with_args = window.process_manager.remove_from_queue() - assert process_with_args is not None - runner.data_and_error_queue = process_with_args["data_queue"] - runner.dialog_action_queue = process_with_args["dialog_action_queue"] - runner.evaluate_queue = process_with_args["evaluate_queue"] - runner.routine_process = process_with_args["process"] - - with patch( - "badger.gui.windows.measurement_retry_dialog.BadgerMeasurementRetryDialog.exec_", - return_value=QDialog.Accepted, - ) as exec_mock: - runner.data_and_error_queue.put( - { - "type": "measurement_error", - "title": "Injected measurement error", - "traceback": "traceback details", - } - ) + process_with_args = None + try: + qtbot.addWidget(window) - # Wait up to 1 second in 10 for queue to receive the error. - for _ in range(10): - if not runner.data_and_error_queue.empty(): - break - qtbot.wait(100) + loop = QEventLoop() + QTimer.singleShot(1000, loop.quit) + loop.exec_() + + routine = create_routine() + home_page = window.home_page + home_page.current_routine = routine + home_page.run_monitor.termination_condition = TerminationConditionConfig( + tc_idx=0, + max_eval=3, + ) + home_page.go_run(-1) + home_page.run_monitor.init_routine_runner() + runner = home_page.run_monitor.routine_runner + assert runner is not None + + # Grab queue and process from the process manager + process_with_args = window.process_manager.remove_from_queue() + assert process_with_args is not None + runner.data_and_error_queue = process_with_args.data_queue + runner.dialog_action_queue = process_with_args.dialog_action_queue + runner.evaluate_queue = process_with_args.evaluate_queue + runner.routine_process = process_with_args.process + + with patch( + "badger.gui.windows.measurement_retry_dialog.BadgerMeasurementRetryDialog.exec_", + return_value=QDialog.Accepted, + ) as exec_mock: + runner.data_and_error_queue.put( + MeasurementErrorMessage( + title="Injected measurement error", + traceback="traceback details", + ) + ) - # Explicitly check if any errors were thrown - runner.check_queue() + # Wait up to 1 second in 10 for queue to receive the error. + for _ in range(10): + if not runner.data_and_error_queue.empty(): + break + qtbot.wait(100) - exec_mock.assert_called_once() - response = runner.dialog_action_queue.get(timeout=1) - assert response == { - "type": "measurement_action", - "action": "retry", - } + # Explicitly check if any errors were thrown + runner.check_queue() - window.process_manager.close_proccesses() - process_with_args["process"].terminate() - process_with_args["process"].join() + exec_mock.assert_called_once() + response = runner.dialog_action_queue.get(timeout=1) + assert response == MeasurementActionMessage(action="retry") + finally: + window.process_manager.close_processes() + if process_with_args is not None: + process_with_args.process.terminate() + process_with_args.process.join() # TODO: Check the use_low_noise_prior parameter in the routine # once it's running -- currently use_low_noise_prior is not exposed in the GUI # so need to check the routine object held by the monitor/runner -def test_default_low_noise_prior_in_bo(qtbot, init_multiprocessing): +def test_default_low_noise_prior_in_bo(qtbot: QtBot) -> None: import yaml from xopt.generators import all_generator_names @@ -204,61 +224,62 @@ def test_default_low_noise_prior_in_bo(qtbot, init_multiprocessing): fix_path_issues() window = BadgerMainWindow() + try: + loop = QEventLoop() + QTimer.singleShot(1000, loop.quit) # 1000 ms pause + loop.exec_() - loop = QEventLoop() - QTimer.singleShot(1000, loop.quit) # 1000 ms pause - loop.exec_() - - qtbot.addWidget(window) - - # Create and save a routine - - editor = window.home_page.routine_editor - cb_generator = editor.generator_box.cb - algos = [cb_generator.itemText(i) for i in range(cb_generator.count())] - for algo in algos: - if algo in all_generator_names["bo"]: - qtbot.keyClicks(editor.generator_box.cb, algo) - params = editor.generator_box.edit.get_parameters_yaml() - params_dict = yaml.safe_load(params) + qtbot.addWidget(window) - if "gp_constructor" in params_dict: - # use_low_noise_prior may not be exposed in the GUI for every - # generator; default to False so a hidden key doesn't error. - assert not params_dict["gp_constructor"].get( - "use_low_noise_prior", False - ) - else: # that part of params is hidden so we need to dig deeper - pass + # Create and save a routine + + editor = window.home_page.routine_editor + cb_generator = editor.generator_box.cb + algos = [cb_generator.itemText(i) for i in range(cb_generator.count())] + for algo in algos: + if algo in all_generator_names["bo"]: + qtbot.keyClicks(editor.generator_box.cb, algo) + params = editor.generator_box.edit.get_parameters_yaml() + params_dict = yaml.safe_load(params) + + if "gp_constructor" in params_dict: + # use_low_noise_prior may not be exposed in the GUI for every + # generator; default to False so a hidden key doesn't error. + assert not params_dict["gp_constructor"].get( + "use_low_noise_prior", False + ) + else: # that part of params is hidden so we need to dig deeper + pass + finally: + window.process_manager.close_processes() + + +# TODO: Fix default turbo in bo test +def test_default_turbo_in_bo(qtbot: QtBot) -> None: + return - window.process_manager.close_proccesses() + # import yaml + # from xopt.generators import all_generator_names + # from badger.gui.windows.main_window import BadgerMainWindow + # from badger.tests.utils import fix_db_path_issue -def test_default_turbo_in_bo(qtbot): - return + # fix_db_path_issue() - import yaml - from xopt.generators import all_generator_names + # window = BadgerMainWindow() + # qtbot.addWidget(window) - from badger.gui.windows.main_window import BadgerMainWindow - from badger.tests.utils import fix_db_path_issue + # # Create and save a routine + # qtbot.mouseClick(window.home_page.btn_new, Qt.MouseButton.LeftButton) + # assert window.home_page.tabs.currentIndex() == 1 # jump to the editor - fix_db_path_issue() + # editor = window.home_page.routine_editor + # cb_generator = editor.generator_box.cb + # algos = [cb_generator.itemText(i) for i in range(cb_generator.count())] + # for algo in algos: + # if algo in all_generator_names["bo"]: + # qtbot.keyClicks(editor.routine_page.generator_box.cb, algo) + # params = editor.routine_page.generator_box.edit.get_parameters_yaml() + # params_dict = yaml.safe_load(params) - window = BadgerMainWindow() - qtbot.addWidget(window) - - # Create and save a routine - qtbot.mouseClick(window.home_page.btn_new, Qt.MouseButton.LeftButton) - assert window.home_page.tabs.currentIndex() == 1 # jump to the editor - - editor = window.home_page.routine_editor - cb_generator = editor.generator_box.cb - algos = [cb_generator.itemText(i) for i in range(cb_generator.count())] - for algo in algos: - if algo in all_generator_names["bo"]: - qtbot.keyClicks(editor.routine_page.generator_box.cb, algo) - params = editor.routine_page.generator_box.edit.get_parameters_yaml() - params_dict = yaml.safe_load(params) - - assert params_dict["turbo_controller"] == "optimize" + # assert params_dict["turbo_controller"] == "optimize" diff --git a/src/badger/tests/test_home_page.py b/src/badger/tests/test_home_page.py index b5bd6192..c096ad0c 100644 --- a/src/badger/tests/test_home_page.py +++ b/src/badger/tests/test_home_page.py @@ -2,10 +2,11 @@ import pytest from PyQt5.QtCore import QEventLoop, QTimer +from pytestqt.qtbot import QtBot @pytest.fixture(scope="session") -def init_multiprocessing(): +def init_multiprocessing() -> None: # Use 'spawn' on Windows, 'fork' on Unix-like systems method = ( "spawn" @@ -15,7 +16,8 @@ def init_multiprocessing(): multiprocessing.set_start_method(method, force=True) -def test_home_page_run_routine(qtbot, init_multiprocessing): +@pytest.mark.usefixtures("init_multiprocessing") +def test_home_page_run_routine(qtbot: QtBot) -> None: from badger.archive import save_tmp_run from badger.gui.windows.main_window import BadgerMainWindow from badger.tests.utils import ( @@ -23,48 +25,49 @@ def test_home_page_run_routine(qtbot, init_multiprocessing): create_routine, fix_path_issues, ) + from badger.types import TerminationConditionConfig fix_path_issues() main_page = BadgerMainWindow() + try: + loop = QEventLoop() + QTimer.singleShot(1000, loop.quit) # 1000 ms pause + loop.exec_() - loop = QEventLoop() - QTimer.singleShot(1000, loop.quit) # 1000 ms pause - loop.exec_() + home_page = main_page.home_page + # test running routines w high level interface + routines = [ + create_routine(), + create_multiobjective_routine(), + # create_routine_turbo(), + ] - home_page = main_page.home_page - # test running routines w high level interface - routines = [ - create_routine(), - create_multiobjective_routine(), - # create_routine_turbo(), - ] + for ele in routines: + tmp_filename = save_tmp_run(ele) + home_page.current_routine = ele + home_page.run_monitor.testing = True + home_page.run_monitor.routine_filename = tmp_filename + home_page.run_monitor.termination_condition = TerminationConditionConfig( + tc_idx=0, + max_eval=3, + ) + home_page.go_run(-1) + # start run in a thread and wait some time for it to finish + home_page.run_monitor.start(True) - for ele in routines: - tmp_filename = save_tmp_run(ele) - home_page.current_routine = ele - home_page.run_monitor.testing = True - home_page.run_monitor.routine_filename = tmp_filename - home_page.run_monitor.termination_condition = { - "tc_idx": 0, - "max_eval": 3, - } - home_page.go_run(-1) - # start run in a thread and wait some time for it to finish - home_page.run_monitor.start(True) + with qtbot.waitSignal( + home_page.run_monitor.sig_progress, timeout=1000 + ) as blocker: + pass - with qtbot.waitSignal( - home_page.run_monitor.sig_progress, timeout=1000 - ) as blocker: - pass + assert blocker.signal_triggered - assert blocker.signal_triggered + # Wait until the run is done + while home_page.run_monitor.running: + qtbot.wait(100) - # Wait until the run is done - while home_page.run_monitor.running: - qtbot.wait(100) - - # assert we get the right result, ie. correct number of samples - assert len(home_page.run_monitor.routine.data) == 3 - - main_page.process_manager.close_proccesses() + # assert we get the right result, ie. correct number of samples + assert len(home_page.run_monitor.routine.data) == 3 + finally: + main_page.process_manager.close_processes() diff --git a/src/badger/tests/test_intf.py b/src/badger/tests/test_intf.py index 82d17966..8960e4fe 100644 --- a/src/badger/tests/test_intf.py +++ b/src/badger/tests/test_intf.py @@ -1,4 +1,4 @@ -def test_find_intf(): +def test_find_intf() -> None: from badger.factory import get_intf, list_intf assert len(list_intf()) == 1 @@ -8,7 +8,7 @@ def test_find_intf(): assert configs["version"] == "1.0" -def test_get_params(): +def test_get_params() -> None: from badger.factory import get_intf _, configs = get_intf("test") @@ -17,7 +17,7 @@ def test_get_params(): assert params == {"flag": 0} -def test_set_params(): +def test_set_params() -> None: from badger.factory import get_intf Interface, _ = get_intf("test") @@ -26,7 +26,7 @@ def test_set_params(): assert intf.flag == 1 -def test_get_values(): +def test_get_values() -> None: from badger.factory import get_intf Interface, _ = get_intf("test") @@ -40,7 +40,7 @@ def test_get_values(): assert value == 0.5 -def test_set_values(): +def test_set_values() -> None: from badger.factory import get_intf Interface, _ = get_intf("test") @@ -57,7 +57,7 @@ def test_set_values(): assert value == 5 -def test_recording(): +def test_recording() -> None: from badger.factory import get_intf Interface, _ = get_intf("test") diff --git a/src/badger/tests/test_lib_basic.py b/src/badger/tests/test_lib_basic.py index 5e59aecb..9f44388a 100644 --- a/src/badger/tests/test_lib_basic.py +++ b/src/badger/tests/test_lib_basic.py @@ -5,7 +5,7 @@ # assert configs['name'] == 'upper_confidence_bound' -def test_env_api(): +def test_env_api() -> None: from badger.factory import get_env, list_env assert len(list_env()) == 2 @@ -14,7 +14,7 @@ def test_env_api(): assert configs["name"] == "test" -def test_intf_api(): +def test_intf_api() -> None: from badger.factory import get_intf, list_intf assert len(list_intf()) == 1 diff --git a/src/badger/tests/test_measurement_retry_dialog.py b/src/badger/tests/test_measurement_retry_dialog.py index 868399a1..6360ce5f 100644 --- a/src/badger/tests/test_measurement_retry_dialog.py +++ b/src/badger/tests/test_measurement_retry_dialog.py @@ -1,15 +1,16 @@ import pytest from PyQt5.QtCore import QTimer from PyQt5.QtWidgets import QDialog +from pytestqt.qtbot import QtBot # conftest suppresses ExpandableMessageBox's exec_ call, but we need it here. @pytest.fixture(autouse=True) -def suppress_popups(): +def suppress_popups() -> None: pass -def test_measurement_retry_dialog_retry(qtbot): +def test_measurement_retry_dialog_retry(qtbot: QtBot) -> None: from badger.gui.windows.measurement_retry_dialog import ( BadgerMeasurementRetryDialog, ) @@ -22,7 +23,7 @@ def test_measurement_retry_dialog_retry(qtbot): assert dialog.exec_() == QDialog.Accepted -def test_measurement_retry_dialog_stop(qtbot): +def test_measurement_retry_dialog_stop(qtbot: QtBot) -> None: from badger.gui.windows.measurement_retry_dialog import ( BadgerMeasurementRetryDialog, ) diff --git a/src/badger/tests/test_process_manager.py b/src/badger/tests/test_process_manager.py index 6a4dd2b4..2babb8e3 100644 --- a/src/badger/tests/test_process_manager.py +++ b/src/badger/tests/test_process_manager.py @@ -1,15 +1,18 @@ import pytest +from pytestqt.qtbot import QtBot + +from badger.gui.components.process_manager import ProcessManager @pytest.fixture -def process_manager(): +def process_manager() -> ProcessManager: from badger.gui.components.process_manager import ProcessManager return ProcessManager() class TestProcessManager: - def test_add_to_queue(self, process_manager): + def test_add_to_queue(self, process_manager: ProcessManager) -> None: """ Test that a process can be added to the queue. """ @@ -18,7 +21,9 @@ def test_add_to_queue(self, process_manager): assert len(process_manager.processes_queue) == 1 assert process_manager.processes_queue[0] == process_with_args - def test_remove_from_queue(self, process_manager, qtbot): + def test_remove_from_queue( + self, process_manager: ProcessManager, qtbot: QtBot + ) -> None: """ Test that a process can be removed from the queue and the correct process is returned. """ diff --git a/src/badger/tests/test_routine.py b/src/badger/tests/test_routine.py index 2d5313af..ef411890 100644 --- a/src/badger/tests/test_routine.py +++ b/src/badger/tests/test_routine.py @@ -1,10 +1,11 @@ import os +from collections.abc import Generator import pytest class TestRoutine: - def test_routine_init(self): + def test_routine_init(self) -> None: from badger.tests.utils import create_routine routine = create_routine() @@ -12,7 +13,7 @@ def test_routine_init(self): assert len(routine.data) == 1 - def test_routine_serialization(self): + def test_routine_serialization(self) -> None: from badger.routine import Routine from badger.tests.utils import create_routine @@ -28,7 +29,7 @@ def test_routine_serialization(self): assert len(lroutine.data) == 2 - def test_routine_env_dump(self): + def test_routine_env_dump(self) -> None: from badger.environment import Environment from badger.routine import Routine from badger.tests.utils import create_routine @@ -44,7 +45,7 @@ def test_routine_env_dump(self): assert routine_re.environment.flag == 1 @pytest.fixture(scope="module", autouse=True) - def clean_up(self): + def clean_up(self) -> Generator[None, None, None]: yield files = ["test.yaml"] for f in files: diff --git a/src/badger/tests/test_routine_page.py b/src/badger/tests/test_routine_page.py index e286e763..1aa49cd5 100644 --- a/src/badger/tests/test_routine_page.py +++ b/src/badger/tests/test_routine_page.py @@ -13,7 +13,7 @@ from pytestqt.qtbot import QtBot -def test_routine_page_init(qtbot: QtBot): +def test_routine_page_init(qtbot: QtBot) -> None: from badger.gui.components.routine_page import BadgerRoutinePage window = BadgerRoutinePage() @@ -21,7 +21,7 @@ def test_routine_page_init(qtbot: QtBot): qtbot.addWidget(window) -def test_set_routine(qtbot: QtBot): +def test_set_routine(qtbot: QtBot) -> None: from badger.gui.components.routine_page import BadgerRoutinePage from badger.tests.utils import create_routine @@ -31,7 +31,7 @@ def test_set_routine(qtbot: QtBot): window.set_routine(routine) -def test_routine_generation(qtbot: QtBot): +def test_routine_generation(qtbot: QtBot) -> None: from badger.errors import BadgerRoutineError # test if a simple routine can be created @@ -79,7 +79,7 @@ def test_routine_generation(qtbot: QtBot): assert routine.xopt_version == get_xopt_version() -def test_add_additional_vars(qtbot: QtBot): +def test_add_additional_vars(qtbot: QtBot) -> None: from badger.gui.components.routine_page import BadgerRoutinePage window = BadgerRoutinePage() @@ -122,7 +122,7 @@ def test_add_additional_vars(qtbot: QtBot): assert window.env_box.var_table.rowCount() == n_rows + 1 -def test_initial_points(qtbot: QtBot): +def test_initial_points(qtbot: QtBot) -> None: # test to make sure initial points widget works properly from badger.gui.components.routine_page import BadgerRoutinePage @@ -151,7 +151,7 @@ def test_initial_points(qtbot: QtBot): ) -def test_ui_update(qtbot: QtBot): +def test_ui_update(qtbot: QtBot) -> None: # test to make sure initial points widget works properly from badger.gui.components.routine_page import BadgerRoutinePage from badger.tests.utils import create_routine @@ -183,7 +183,7 @@ def test_ui_update(qtbot: QtBot): assert actual_params == expected_params -def test_constraints(qtbot: QtBot): +def test_constraints(qtbot: QtBot) -> None: # test if a simple routine can be created from badger.gui.components.routine_page import BadgerRoutinePage @@ -206,7 +206,7 @@ def test_constraints(qtbot: QtBot): assert routine.critical_constraint_names == ["c"] -def test_observables(qtbot: QtBot): +def test_observables(qtbot: QtBot) -> None: # test if a simple routine can be created from badger.gui.components.routine_page import BadgerRoutinePage @@ -228,7 +228,7 @@ def test_observables(qtbot: QtBot): assert routine.vocs.observables == {"c": Observable(dtype=None)} -def test_add_random_points(qtbot: QtBot): +def test_add_random_points(qtbot: QtBot) -> None: # test to add random points to initial points table from badger.gui.components.routine_page import BadgerRoutinePage @@ -241,12 +241,20 @@ def test_add_random_points(qtbot: QtBot): qtbot.keyClicks(window.generator_box.cb, "random") qtbot.keyClicks(window.env_box.cb, "test") - window.env_box.var_table.cellWidget(0, 0).setChecked(True) - window.env_box.var_table.cellWidget(1, 0).setChecked(True) - window.env_box.var_table.cellWidget(2, 0).setChecked(True) - window.env_box.obj_table.cellWidget(0, 0).setChecked(True) - - def handle_dialog(): + var_table_0_0 = window.env_box.var_table.cellWidget(0, 0) + if var_table_0_0 is not None: + var_table_0_0.setChecked(True) + var_table_1_0 = window.env_box.var_table.cellWidget(1, 0) + if var_table_1_0 is not None: + var_table_1_0.setChecked(True) + var_table_2_0 = window.env_box.var_table.cellWidget(2, 0) + if var_table_2_0 is not None: + var_table_2_0.setChecked(True) + obj_table_0_0 = window.env_box.obj_table.cellWidget(0, 0) + if obj_table_0_0 is not None: + obj_table_0_0.setChecked(True) + + def handle_dialog() -> None: while window.rc_dialog is None: QApplication.processEvents() @@ -257,8 +265,9 @@ def handle_dialog(): QTimer.singleShot(0, handle_dialog) - qtbot.mouseClick(window.env_box.btn_add_rand, Qt.LeftButton) + qtbot.mouseClick(window.env_box.btn_add_rand, Qt.MouseButton.LeftButton) routine = window._compose_routine() + assert routine.initial_points is not None assert routine.initial_points.shape[0] == 5 # curr is 0.5, frac is 0.05, range is [-1, 1] @@ -269,28 +278,28 @@ def handle_dialog(): # TODO: Test if the EI, Simplex, and RCDS params show o the params editor # are the simplified versions -def test_simplified_generator_params(qtbot: QtBot): +def test_simplified_generator_params(qtbot: QtBot) -> None: pass # TODO: First load an old routine w/ initail points, # then create a new routine and check if the initial points panel # is cleared -def test_initial_points_clear_when_create_routine(qtbot: QtBot): +def test_initial_points_clear_when_create_routine(qtbot: QtBot) -> None: pass # TODO: Test if env selector reacts to scroll events, it should not -def test_scroll_on_environment_selector(qtbot: QtBot): +def test_scroll_on_environment_selector(qtbot: QtBot) -> None: pass # TODO: Test if generator selector reacts to scroll events, it should not -def test_scroll_on_generator_selector(qtbot: QtBot): +def test_scroll_on_generator_selector(qtbot: QtBot) -> None: pass # TODO: Test relative to current behavior, including the auto calculated # bounds and initial points wrt the current variable values -def test_relative_to_current(qtbot: QtBot): +def test_relative_to_current(qtbot: QtBot) -> None: pass diff --git a/src/badger/tests/test_routine_runner.py b/src/badger/tests/test_routine_runner.py index f4b3fe0b..e162e149 100644 --- a/src/badger/tests/test_routine_runner.py +++ b/src/badger/tests/test_routine_runner.py @@ -1,15 +1,18 @@ import multiprocessing +from collections.abc import Generator from unittest.mock import Mock import pytest -from PyQt5.QtCore import QEventLoop, Qt, QTimer +from PyQt5.QtCore import QTimer from PyQt5.QtTest import QSignalSpy -from PyQt5.QtWidgets import QApplication + +from badger.gui.components.process_manager import ProcessManager +from badger.gui.components.routine_runner import BadgerRoutineSubprocess class TestRoutineRunner: - @pytest.fixture(scope="session") - def init_multiprocessing(self): + @pytest.fixture(scope="session", autouse=True) + def init_multiprocessing(self) -> None: # Use 'spawn' on Windows, 'fork' on Unix-like systems method = ( "spawn" @@ -19,9 +22,8 @@ def init_multiprocessing(self): multiprocessing.set_start_method(method, force=True) @pytest.fixture - def process_manager(self): + def process_manager(self) -> Generator[ProcessManager, None, None]: from badger.gui.components.create_process import CreateProcess - from badger.gui.components.process_manager import ProcessManager process_manager = ProcessManager() process_builder = CreateProcess() @@ -30,10 +32,10 @@ def process_manager(self): yield process_manager - process_manager.close_proccesses() + process_manager.close_processes() @pytest.fixture - def instance(self, process_manager, init_multiprocessing): + def instance(self, process_manager: ProcessManager) -> BadgerRoutineSubprocess: from badger.archive import save_tmp_run from badger.gui.components.routine_runner import ( BadgerRoutineSubprocess, @@ -46,7 +48,7 @@ def instance(self, process_manager, init_multiprocessing): instance.pause_event = multiprocessing.Event() return instance - def test_ctrl_routine(self, instance) -> None: + def test_ctrl_routine(self, instance: BadgerRoutineSubprocess) -> None: # Test setting the event instance.ctrl_routine(True) assert not instance.pause_event.is_set() @@ -55,24 +57,24 @@ def test_ctrl_routine(self, instance) -> None: instance.ctrl_routine(False) assert instance.pause_event.is_set() - def test_stop_routine(self, instance): + def test_stop_routine(self, instance: BadgerRoutineSubprocess) -> None: instance.run() instance.stop_routine() assert instance.stop_event.is_set() - def test_save_init_vars(self, instance): + def test_save_init_vars(self, instance: BadgerRoutineSubprocess) -> None: sig_env_ready_spy = QSignalSpy(instance.signals.env_ready) instance.save_init_vars() assert len(sig_env_ready_spy) == 1 - def test_after_evaluate(self, instance): + def test_after_evaluate(self, instance: BadgerRoutineSubprocess) -> None: sig_progress_spy = QSignalSpy(instance.signals.progress) instance.setup_timer() instance.after_evaluate(True) assert len(sig_progress_spy) == 1 instance.timer.stop() - def test_check_queue(self, instance): + def test_check_queue(self, instance: BadgerRoutineSubprocess) -> None: sig_finished_spy = QSignalSpy(instance.signals.finished) instance.run() instance.data_and_error_queue.empty = Mock(return_value=True) @@ -87,7 +89,7 @@ def test_check_queue(self, instance): # instance.stop_routine() # assert len(sig_progress_spy) > 0 - def test_setup_timer(self, instance): + def test_setup_timer(self, instance: BadgerRoutineSubprocess) -> None: instance.setup_timer() assert isinstance(instance.timer, QTimer) assert instance.timer.interval() == 100 @@ -95,79 +97,17 @@ def test_setup_timer(self, instance): instance.timer.stop() - def test_run(self, instance): + def test_run(self, instance: BadgerRoutineSubprocess) -> None: instance.run() instance.ctrl_routine(True) assert instance.routine_process.is_alive() assert instance.wait_event.is_set() instance.stop_routine() - def test_set_termination_condition(self, instance): - instance.set_termination_condition(True) - assert instance.termination_condition - - # TODO: check for signal emit message - - def test_turbo_with_routine_runner(self, qtbot, init_multiprocessing): - # TODO: make this test more stable - return - - from badger.gui.windows.main_window import BadgerMainWindow - from badger.gui.windows.message_dialog import ( - BadgerScrollableMessageBox, - ) - - window = BadgerMainWindow() - # qtbot.addWidget(window) - - loop = QEventLoop() - QTimer.singleShot(1000, loop.quit) # 1000 ms pause - loop.exec_() - - # Create and save a routine - editor = window.home_page.routine_editor - # Turn off relative to current - editor.env_box.relative_to_curr.setChecked(False) - # Config env and vocs - qtbot.keyClicks(editor.env_box.cb, "test") - editor.env_box.var_table.cellWidget(0, 0).setChecked(True) - editor.env_box.obj_table.cellWidget(0, 0).setChecked(True) - qtbot.mouseClick(editor.env_box.btn_add_curr, Qt.MouseButton.LeftButton) - - # Run the routine - monitor = window.home_page.run_monitor - monitor.testing = True + def test_set_termination_condition(self, instance: BadgerRoutineSubprocess) -> None: + from badger.types import TerminationConditionConfig - def handle_dialog(): - while monitor.tc_dialog is None: - QApplication.processEvents() - - # Set max evaluation to 2, then run the optimization - monitor.tc_dialog.sb_max_eval.setValue(2) - qtbot.mouseClick(monitor.tc_dialog.btn_run, Qt.MouseButton.LeftButton) - - QTimer.singleShot(0, handle_dialog) - - # Bypass the possible Badger error dialog - def patched_showEvent(original_showEvent): - def inner(ins, event): - original_showEvent(ins, event) # Call the original showEvent - QTimer.singleShot(100, ins.accept) # Close the dialog after 100 ms - - return inner - - BadgerScrollableMessageBox.showEvent = patched_showEvent( - BadgerScrollableMessageBox.showEvent + instance.set_termination_condition( + TerminationConditionConfig(tc_idx=0, max_eval=2) ) - - action_bar = window.home_page.run_action_bar - action_bar.run_until_action.trigger() - monitor.routine_runner.data_and_error_queue.empty = Mock(return_value=True) - - # Wait until the run is done - while monitor.running: - qtbot.wait(100) - - assert len(monitor.routine.data) == 2 - - window.process_manager.close_proccesses() + assert instance.termination_condition diff --git a/src/badger/tests/test_run_monitor.py b/src/badger/tests/test_run_monitor.py index fd220c6d..03257cd4 100644 --- a/src/badger/tests/test_run_monitor.py +++ b/src/badger/tests/test_run_monitor.py @@ -1,5 +1,6 @@ import multiprocessing import warnings +from collections.abc import Generator from unittest.mock import patch import numpy as np @@ -8,11 +9,19 @@ from PyQt5.QtGui import QMouseEvent from PyQt5.QtTest import QSignalSpy from PyQt5.QtWidgets import QApplication, QMessageBox +from pytest_mock import MockerFixture +from pytestqt.qtbot import QtBot + +from badger.gui.components.process_manager import ProcessManager +from badger.gui.components.run_monitor import BadgerOptMonitor +from badger.gui.pages.home_page import BadgerHomePage +from badger.gui.windows.main_window import BadgerMainWindow +from badger.types import TerminationConditionConfig class TestRunMonitor: @pytest.fixture(scope="session") - def init_multiprocessing(self): + def init_multiprocessing(self) -> None: # Use 'spawn' on Windows, 'fork' on Unix-like systems method = ( "spawn" @@ -22,9 +31,8 @@ def init_multiprocessing(self): multiprocessing.set_start_method(method, force=True) @pytest.fixture - def process_manager(self): + def process_manager(self) -> Generator[ProcessManager, None, None]: from badger.gui.components.create_process import CreateProcess - from badger.gui.components.process_manager import ProcessManager process_manager = ProcessManager() process_builder = CreateProcess() @@ -32,12 +40,13 @@ def process_manager(self): process_builder.create_subprocess() yield process_manager - process_manager.close_proccesses() + process_manager.close_processes() @pytest.fixture - def monitor(self, process_manager, init_multiprocessing): + def monitor( + self, process_manager: ProcessManager, init_multiprocessing + ) -> BadgerOptMonitor: from badger.archive import save_tmp_run - from badger.gui.components.run_monitor import BadgerOptMonitor from badger.tests.utils import create_routine routine = create_routine() @@ -50,9 +59,8 @@ def monitor(self, process_manager, init_multiprocessing): return monitor @pytest.fixture - def home_page(self, process_manager): + def home_page(self, process_manager: ProcessManager) -> BadgerHomePage: from badger.archive import save_tmp_run - from badger.gui.pages.home_page import BadgerHomePage from badger.tests.utils import create_routine routine = create_routine() @@ -68,9 +76,8 @@ def home_page(self, process_manager): return home @pytest.fixture - def home_page_critical(self, process_manager): + def home_page_critical(self, process_manager: ProcessManager) -> BadgerHomePage: from badger.archive import save_tmp_run - from badger.gui.pages.home_page import BadgerHomePage from badger.tests.utils import create_routine_critical routine = create_routine_critical() @@ -85,13 +92,13 @@ def home_page_critical(self, process_manager): return home - def add_data(self, monitor): + def add_data(self, monitor: BadgerOptMonitor) -> None: monitor.routine.random_evaluate(10) monitor.init_plots(monitor.routine) assert len(monitor.routine.data) == 10 - def test_run_monitor(self, process_manager): + def test_run_monitor(self, process_manager: ProcessManager) -> None: from badger.gui.components.run_monitor import BadgerOptMonitor from badger.tests.utils import create_routine @@ -123,18 +130,20 @@ def test_run_monitor(self, process_manager): # set up run monitor and test it monitor.init_routine_runner() - monitor.routine_runner.set_termination_condition({"tc_idx": 0, "max_eval": 2}) + monitor.routine_runner.set_termination_condition( + TerminationConditionConfig(tc_idx=0, max_eval=2) + ) spy = QSignalSpy(monitor.routine_runner.signals.progress) assert spy.isValid() monitor.start() - def test_routine_identity(self, home_page): + def test_routine_identity(self, home_page: BadgerHomePage) -> None: monitor = home_page.run_monitor monitor.init_routine_runner() assert monitor.routine_runner.routine == monitor.routine - def test_plotting(self, qtbot, monitor): + def test_plotting(self, qtbot: QtBot, monitor: BadgerOptMonitor) -> None: self.add_data(monitor) monitor.update_curves() @@ -150,7 +159,9 @@ def test_plotting(self, qtbot, monitor): monitor.x_plot_y_axis = 1 monitor.update_curves() - def test_click_graph(self, qtbot, monitor, mocker): + def test_click_graph( + self, qtbot: QtBot, monitor: BadgerOptMonitor, mocker: MockerFixture + ) -> None: self.add_data(monitor) sig_inspect_spy = QSignalSpy(monitor.sig_inspect) monitor.plot_x_axis = True @@ -165,7 +176,9 @@ def test_click_graph(self, qtbot, monitor, mocker): assert new_variable_value != orginal_value assert len(sig_inspect_spy) == 1 - def create_test_run_monitor(self, process_manager, add_data=True): + def create_test_run_monitor( + self, process_manager: ProcessManager, add_data: bool = True + ) -> BadgerOptMonitor: from badger.gui.components.run_monitor import BadgerOptMonitor from badger.tests.utils import create_routine @@ -182,7 +195,9 @@ def create_test_run_monitor(self, process_manager, add_data=True): return monitor - def test_x_axis_specification(self, qtbot, monitor, mocker): + def test_x_axis_specification( + self, qtbot: QtBot, monitor: BadgerOptMonitor, mocker: MockerFixture + ) -> None: # check iteration/time drop down menu self.add_data(monitor) @@ -242,11 +257,13 @@ def test_x_axis_specification(self, qtbot, monitor, mocker): monitor.cb_plot_x.setCurrentIndex(1) assert current_index == monitor.inspector_variable.value() - def test_y_axis_specification(self, qtbot, monitor): - monitor.termination_condition = { - "tc_idx": 0, - "max_eval": 10, - } + def test_y_axis_specification( + self, qtbot: QtBot, monitor: BadgerOptMonitor + ) -> None: + monitor.termination_condition = TerminationConditionConfig( + tc_idx=0, + max_eval=10, + ) monitor.start(True) # Wait until the run is done @@ -284,14 +301,14 @@ def test_y_axis_specification(self, qtbot, monitor): normalized_raw_value = monitor.curves_variable["x0"].getData()[1][index] assert normalized_raw_value == 0.75 - def test_pause_play(self, qtbot, home_page): + def test_pause_play(self, qtbot: QtBot, home_page: BadgerHomePage) -> None: monitor = home_page.run_monitor action_bar = home_page.run_action_bar - monitor.termination_condition = { - "tc_idx": 0, - "max_eval": 10, - } + monitor.termination_condition = TerminationConditionConfig( + tc_idx=0, + max_eval=10, + ) spy = QSignalSpy(monitor.sig_pause) monitor.start(True) @@ -308,7 +325,7 @@ def test_pause_play(self, qtbot, home_page): while monitor.running: qtbot.wait(100) - def test_jump_to_optimum(self, qtbot, home_page): + def test_jump_to_optimum(self, qtbot: QtBot, home_page: BadgerHomePage) -> None: monitor = home_page.run_monitor action_bar = home_page.run_action_bar @@ -330,9 +347,10 @@ def test_jump_to_optimum(self, qtbot, home_page): # Check if it is going to be the optimal solution assert max_value == optimal_value - def test_reset_environment(self, qtbot, init_multiprocessing): + def test_reset_environment( + self, qtbot: QtBot, init_multiprocessing, home_page: BadgerHomePage + ) -> None: from badger.archive import save_tmp_run - from badger.gui.windows.main_window import BadgerMainWindow from badger.tests.utils import ( create_routine, get_current_vars, @@ -358,10 +376,10 @@ def test_reset_environment(self, qtbot, init_multiprocessing): # check if reset button click signal is trigged and if state is same as original state after click init_vars = get_current_vars(routine) - monitor.termination_condition = { - "tc_idx": 0, - "max_eval": 10, - } + monitor.termination_condition = TerminationConditionConfig( + tc_idx=0, + max_eval=10, + ) home_page.go_run(-1) monitor.start(True) @@ -400,7 +418,7 @@ def test_reset_environment(self, qtbot, init_multiprocessing): # Check if the env has been reset curr_vars = get_current_vars(monitor.routine) assert np.all(curr_vars == init_vars) - window.process_manager.close_proccesses() + window.process_manager.close_processes() def test_dial_in_solution(self, qtbot, home_page): monitor = home_page.run_monitor diff --git a/src/badger/tests/test_settings.py b/src/badger/tests/test_settings.py index 04c2c66e..5c697e37 100644 --- a/src/badger/tests/test_settings.py +++ b/src/badger/tests/test_settings.py @@ -17,7 +17,7 @@ class TestBadgerConfig: config_file = "/mock/config.yaml" - def setup_method(self, method): + def setup_method(self) -> None: self.mock_badger_config = BadgerConfig( BADGER_PLUGIN_ROOT=Setting( display_name="plugin root", @@ -63,11 +63,13 @@ def setup_method(self, method): ), ) - def teardown_method(self, method): + def teardown_method(self) -> None: ConfigSingleton._instance = None self.mock_badger_config = None - def test_get_user_config_folder_windows(self, monkeypatch): + def test_get_user_config_folder_windows( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: monkeypatch.setattr(platform, "system", lambda: "Windows") monkeypatch.setenv("APPDATA", r"C:\Users\test\AppData\Roaming") @@ -75,7 +77,7 @@ def test_get_user_config_folder_windows(self, monkeypatch): assert config_folder == r"C:\Users\test\AppData\Roaming" - def test_get_user_config_folder_mac(self, monkeypatch): + def test_get_user_config_folder_mac(self, monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(platform, "system", lambda: "Darwin") monkeypatch.setattr( os.path, @@ -87,7 +89,9 @@ def test_get_user_config_folder_mac(self, monkeypatch): assert config_folder == "/Users/test/Library/Application Support/Badger" - def test_get_user_config_folder_linux(self, monkeypatch): + def test_get_user_config_folder_linux( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: monkeypatch.setattr(platform, "system", lambda: "Linux") monkeypatch.setattr(os.path, "expanduser", lambda x: "/home/test/.config") @@ -95,13 +99,15 @@ def test_get_user_config_folder_linux(self, monkeypatch): assert config_folder == "/home/test/.config" - def test_get_user_config_folder_unsupported_os(self, monkeypatch): + def test_get_user_config_folder_unsupported_os( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: monkeypatch.setattr(platform, "system", lambda: "UnsupportedOS") with pytest.raises(OSError): get_user_config_folder() - def test_init_settings(self): + def test_init_settings(self) -> None: mock_config_singleton = MagicMock(spec=ConfigSingleton) with ( patch( @@ -124,7 +130,7 @@ def test_init_settings(self): ) assert config_singleton == mock_config_singleton - def test_config_singleton_initialization(self): + def test_config_singleton_initialization(self) -> None: # Patch the open function and ensure it reads the correct mock configuration with ( patch("os.path.exists", return_value=True), @@ -146,7 +152,7 @@ def test_config_singleton_initialization(self): == "/mock/template/root" ) - def test_config_singleton_create_new_config_if_not_exists(self): + def test_config_singleton_create_new_config_if_not_exists(self) -> None: with ( patch("os.path.exists", return_value=False), patch("badger.settings.BadgerConfig", return_value=self.mock_badger_config), @@ -154,7 +160,7 @@ def test_config_singleton_create_new_config_if_not_exists(self): config_singleton = ConfigSingleton(self.config_file) assert isinstance(config_singleton.config, BadgerConfig) - def test_update_and_save_config(self): + def test_update_and_save_config(self) -> None: with ( patch("os.path.exists", return_value=True), patch( @@ -173,7 +179,7 @@ def test_update_and_save_config(self): config_singleton.config.BADGER_PLUGIN_ROOT.value == "/new/plugin/root" ) - def test_list_settings(self): + def test_list_settings(self) -> None: with ( patch("os.path.exists", return_value=True), patch( @@ -190,7 +196,7 @@ def test_list_settings(self): assert "BADGER_PLUGIN_ROOT" in settings assert settings["BADGER_PLUGIN_ROOT"]["value"] == "/mock/plugin/root" - def test_read_value(self): + def test_read_value(self) -> None: with ( patch("os.path.exists", return_value=True), patch( @@ -206,7 +212,7 @@ def test_read_value(self): value = config_singleton.read_value("BADGER_PLUGIN_ROOT") assert value == "/mock/plugin/root" - def test_read_display_name(self): + def test_read_display_name(self) -> None: with ( patch("os.path.exists", return_value=True), patch( @@ -222,7 +228,7 @@ def test_read_display_name(self): display_name = config_singleton.read_display_name("BADGER_PLUGIN_ROOT") assert display_name == "plugin root" - def test_read_description(self): + def test_read_description(self) -> None: with ( patch("os.path.exists", return_value=True), patch( @@ -238,7 +244,7 @@ def test_read_description(self): description = config_singleton.read_description("BADGER_PLUGIN_ROOT") assert description == "Mock plugin root" - def test_read_is_path(self): + def test_read_is_path(self) -> None: with ( patch("os.path.exists", return_value=True), patch( @@ -254,7 +260,7 @@ def test_read_is_path(self): is_path = config_singleton.read_is_path("BADGER_PLUGIN_ROOT") assert is_path - def test_write_value(self): + def test_write_value(self) -> None: with ( patch("os.path.exists", return_value=True), patch( @@ -272,7 +278,7 @@ def test_write_value(self): config_singleton.config.BADGER_PLUGIN_ROOT.value == "/new/plugin/root" ) - def test_reset_settings(self): + def test_reset_settings(self) -> None: with ( patch("os.path.exists", return_value=False), patch("badger.settings.BadgerConfig", return_value=self.mock_badger_config), diff --git a/src/badger/tests/utils.py b/src/badger/tests/utils.py index f57b890d..0b934778 100644 --- a/src/badger/tests/utils.py +++ b/src/badger/tests/utils.py @@ -1,15 +1,18 @@ import os +from typing import Any +import numpy as np import pandas as pd -from xopt import VOCS +from gest_api.vocs import VOCS from xopt.generators import RandomGenerator from xopt.generators.bayesian import UpperConfidenceBoundGenerator +from badger.routine import Routine -def create_routine(): - from badger.routine import Routine - test_routine = { +def create_routine() -> Routine: + + test_routine: dict[str, Any] = { "name": "routine-for-core-test", "generator": "random", "env": "test", @@ -35,10 +38,9 @@ def create_routine(): ) -def create_multiobjective_routine(): - from badger.routine import Routine +def create_multiobjective_routine() -> Routine: - test_routine = { + test_routine: dict[str, Any] = { "name": "routine-for-core-test", "generator": "random", "generator_params": {}, @@ -57,17 +59,15 @@ def create_multiobjective_routine(): return Routine( name="test", - vocs=vocs, generator=generator, environment={"name": "multiobjective_test"}, initial_points=pd.DataFrame(test_routine["init_points"]), ) -def create_routine_turbo(): - from badger.routine import Routine +def create_routine_turbo() -> Routine: - test_routine = { + test_routine: dict[str, Any] = { "name": "routine-for-turbo-test", "generator": "expected_improvement", "env": "test", @@ -98,17 +98,15 @@ def create_routine_turbo(): return Routine( name="test-turbo", - vocs=vocs, generator=generator, environment={"name": "test"}, initial_points=pd.DataFrame(test_routine["config"]["init_points"]), ) -def create_routine_critical(): - from badger.routine import Routine +def create_routine_critical() -> Routine: - test_routine = { + test_routine: dict[str, Any] = { "name": "routine-for-critical-test", "generator": "random", "env": "test", @@ -128,7 +126,6 @@ def create_routine_critical(): return Routine( name="test", - vocs=vocs, generator=generator, environment={"name": "test"}, initial_points=pd.DataFrame(test_routine["init_points"]), @@ -136,10 +133,9 @@ def create_routine_critical(): ) -def create_routine_constrained_ucb(): - from badger.routine import Routine +def create_routine_constrained_ucb() -> Routine: - test_routine = { + test_routine: dict[str, Any] = { "name": "routine-for-ucb-cons-test", "generator": "expected_improvement", "env": "test", @@ -174,19 +170,23 @@ def create_routine_constrained_ucb(): ) -def get_current_vars(routine): +def get_current_vars(routine: Routine) -> list[float]: var_names = routine.vocs.variable_names var_dict = routine.environment.get_variables(var_names) return list(var_dict.values()) -def get_vars_in_row(routine, idx=0): +def get_vars_in_row(routine: Routine, idx: int = 0) -> np.ndarray: var_names = routine.vocs.variable_names - return routine.data.iloc[idx][var_names].to_numpy() + if routine.data is None: + raise ValueError("Routine data is None. Unable to get variables in row.") + output: np.ndarray = routine.data.iloc[idx][var_names].to_numpy() + # output = routine.data.iat[idx][var_names].to_numpy() + return output -def fix_path_issues(): +def fix_path_issues() -> None: from badger.archive import BADGER_ARCHIVE_ROOT from badger.settings import init_settings diff --git a/src/badger/tests/x_test_db.py b/src/badger/tests/x_test_db.py index 98e3f948..1f321b48 100644 --- a/src/badger/tests/x_test_db.py +++ b/src/badger/tests/x_test_db.py @@ -1,20 +1,20 @@ class TestDB: - def test_save_routine(self): + def test_save_routine(self) -> None: from badger.db import remove_routine, save_routine - from badger.tests.utils import create_routine, fix_db_path_issue + from badger.tests.utils import create_routine, fix_path_issues - fix_db_path_issue() + fix_path_issues() routine = create_routine() save_routine(routine) remove_routine(routine.id) - def test_load_routine(self): + def test_load_routine(self) -> None: from badger.db import load_routine, remove_routine, save_routine - from badger.tests.utils import create_routine, fix_db_path_issue + from badger.tests.utils import create_routine, fix_path_issues - fix_db_path_issue() + fix_path_issues() routine = create_routine() save_routine(routine) diff --git a/src/badger/tests/x_test_routine_id.py b/src/badger/tests/x_test_routine_id.py index ba9bfa2e..6ec691b6 100644 --- a/src/badger/tests/x_test_routine_id.py +++ b/src/badger/tests/x_test_routine_id.py @@ -1,4 +1,7 @@ -def test_routines_same_name(): +from pytestqt.qtbot import QtBot + + +def test_routines_same_name() -> None: from badger.db import remove_routine, save_routine from badger.tests.utils import create_routine, fix_db_path_issue @@ -15,7 +18,7 @@ def test_routines_same_name(): remove_routine(routine2.id) -def test_modify_routine_no_runs(qtbot): +def test_modify_routine_no_runs(qtbot: QtBot) -> None: from badger.db import list_routine, load_routine, remove_routine from badger.gui.components.routine_page import BadgerRoutinePage @@ -47,7 +50,7 @@ def test_modify_routine_no_runs(qtbot): # TODO: write test for modifying name of routine with runs -def test_modify_routine_name(qtbot): +def test_modify_routine_name(qtbot: QtBot) -> None: from badger.db import list_routine, load_routine, remove_routine, save_run from badger.gui.components.routine_page import BadgerRoutinePage @@ -87,7 +90,7 @@ def test_modify_routine_name(qtbot): # TODO: write test for modifying algorithm of routine with runs -def test_modify_routine_algorithm(qtbot): +def test_modify_routine_algorithm(qtbot: QtBot) -> None: from badger.db import list_routine, load_routine, remove_routine, save_run from badger.gui.components.routine_page import BadgerRoutinePage diff --git a/src/badger/types.py b/src/badger/types.py new file mode 100644 index 00000000..9eeef765 --- /dev/null +++ b/src/badger/types.py @@ -0,0 +1,129 @@ +from __future__ import annotations + +from dataclasses import dataclass +from multiprocessing import Process, Queue +from multiprocessing.connection import Connection +from multiprocessing.synchronize import Event +from typing import Any + +import pandas as pd + +from badger.errors import ( + MEASUREMENT_ACTION_TYPE, + MEASUREMENT_ERROR_TYPE, + ROUTINE_ERROR_TYPE, + TERMINATION_ACTION_TYPE, + TERMINATION_REACHED_TYPE, +) + + +@dataclass +class MeasurementErrorMessage: + """Subprocess -> routine runner: an environment measurement raised. + + Prompts the user to retry or abort the run. + """ + + title: str + traceback: str + type: str = MEASUREMENT_ERROR_TYPE + + +@dataclass +class TerminationReachedMessage: + """Subprocess -> routine runner: a run-until condition was reached. + + Prompts the user to continue or end the run. ``tc_condition`` is a flat + dict with keys ``type``, ``config`` and ``state``. + """ + + tc_condition: dict[str, Any] + title: str = "Termination condition reached" + type: str = TERMINATION_REACHED_TYPE + + +@dataclass +class RoutineErrorMessage: + """Subprocess -> routine runner: a fatal error ended the run. + + Shown to the user in an error dialog. + """ + + title: str + traceback: str + type: str = ROUTINE_ERROR_TYPE + + +# Messages carried by the data/error queue from the subprocess to the runner. +DataQueueMessage = ( + MeasurementErrorMessage | TerminationReachedMessage | RoutineErrorMessage +) + + +@dataclass +class MeasurementActionMessage: + """Routine runner -> subprocess: user's response to a measurement error. + + ``action`` is ``MEASUREMENT_ACTION_RETRY`` or ``MEASUREMENT_ACTION_ABORT``. + """ + + action: str + type: str = MEASUREMENT_ACTION_TYPE + + +@dataclass +class TerminationActionMessage: + """Routine runner -> subprocess: user's response to a termination prompt. + + ``action`` is ``TERMINATION_ACTION_CONTINUE`` or ``TERMINATION_ACTION_END``. + """ + + action: str + type: str = TERMINATION_ACTION_TYPE + + +# Messages carried by the dialog action queue from the runner to the subprocess. +DialogActionMessage = MeasurementActionMessage | TerminationActionMessage + + +@dataclass +class TerminationConditionConfig: + """Run-until configuration set in the termination condition dialog. + + ``tc_idx`` selects which limit is active: 0 = max evaluations (``max_eval``), + 1 = max running time in seconds (``max_time``). ``ftol`` is reserved for a + convergence-tolerance mode that is not yet wired up. + """ + + tc_idx: int = 0 + max_eval: int = 0 + max_time: float = 0.0 + ftol: float = 0.0 + + +@dataclass +class ArgumentQueueType: + routine_id: str + routine_filename: str + routine_name: str + variable_ranges: dict[str, Any] + initial_points: pd.DataFrame | None + evaluate: bool + archive: bool + termination_condition: TerminationConditionConfig | None + start_time: float + testing: bool + run_data: bool + init_points: bool + + +@dataclass +class ProcessWithArgs: + process: Process + stop_event: Event + pause_event: Event + args_queue: Queue[ArgumentQueueType] + data_queue: Queue[DataQueueMessage] + evaluate_queue: tuple[Connection, Connection] + wait_event: Event + dialog_action_queue: Queue[DialogActionMessage] diff --git a/src/badger/utils.py b/src/badger/utils.py index e89ae964..9bea5ccb 100644 --- a/src/badger/utils.py +++ b/src/badger/utils.py @@ -11,14 +11,17 @@ from datetime import UTC, datetime from importlib import metadata from types import TracebackType -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, TypedDict import yaml +from pandas import DataFrame from PyQt5.QtWidgets import QLayout, QWidget from badger.errors import BadgerLoadConfigError if TYPE_CHECKING: + from xopt.generator import Generator + from badger.routine import Routine from decimal import ROUND_CEILING, ROUND_FLOOR, Decimal @@ -31,13 +34,15 @@ class BlockSignalsContext: widgets: Iterable[QWidget | QLayout] - def __init__(self, widgets: QWidget | QLayout | Iterable[QWidget | QLayout]): + def __init__( + self, widgets: QWidget | QLayout | Iterable[QWidget | QLayout] + ) -> None: if isinstance(widgets, Iterable): self.widgets = widgets else: self.widgets = [widgets] - def __enter__(self): + def __enter__(self) -> None: for widget in self.widgets: if widget.signalsBlocked(): logger.warning( @@ -50,7 +55,7 @@ def __exit__( exc_type: type[BaseException] | None, exc_value: BaseException | None, exc_traceback: TracebackType | None, - ): + ) -> None: for widget in self.widgets: if not widget.signalsBlocked(): logger.warning( @@ -62,34 +67,38 @@ def __exit__( # https://stackoverflow.com/a/39681672/4263605 # https://github.com/yaml/pyyaml/issues/234#issuecomment-765894586 class Dumper(yaml.Dumper): - def increase_indent(self, flow=False, indentless=False): + def increase_indent(self, flow: bool = False, indentless: bool = False) -> None: return super().increase_indent(flow, False) -def get_yaml_string(content): +def get_yaml_string(content: Any) -> str: if content is None: return "" - return yaml.dump(content, Dumper=Dumper, default_flow_style=False, sort_keys=False) + return str( + yaml.dump(content, Dumper=Dumper, default_flow_style=False, sort_keys=False) + ) -def yprint(content): +def yprint(content: Any) -> None: print(get_yaml_string(content), end="") -def norm(x, lb, ub): +def norm(x: float, lb: float, ub: float) -> float: return (x - lb) / (ub - lb) -def denorm(x, lb, ub): +def denorm(x: float, lb: float, ub: float) -> float: return (1 - x) * lb + x * ub -def config_list_to_dict(config_list): +def config_list_to_dict( + config_list: Iterable[dict[str, Any]] | None, +) -> dict[str, Any]: if not config_list: return {} - book = {} + book: dict[str, Any] = {} for config in config_list: for k, v in config.items(): book[k] = v @@ -97,8 +106,8 @@ def config_list_to_dict(config_list): return book -def load_config(fname): - configs = None +def load_config(fname: str | None) -> dict[str, Any] | None: + configs: dict[str, Any] | None = None if fname is None: return configs @@ -127,7 +136,9 @@ def load_config(fname): return configs -def merge_params(default_params, params): +def merge_params( + default_params: dict[str, Any] | None, params: dict[str, Any] | None +) -> dict[str, Any] | None: merged_params = None if params is None: @@ -140,9 +151,9 @@ def merge_params(default_params, params): return merged_params -def range_to_str(vranges): +def range_to_str(vranges: list[dict[str, tuple[float, float]]]) -> list[dict[str, str]]: # Transfer the range list to a string for better printing - vranges_str = [] + vranges_str: list[dict[str, str]] = [] for var_dict in vranges: var = next(iter(var_dict)) vrange = var_dict[var] @@ -188,19 +199,26 @@ def curr_ts_to_str(format: str = "lcls-log") -> str: def create_archive_run_filename(routine: "Routine", format: str = "lcls-fname") -> str: - data = routine.sorted_data - env_name = routine.environment.name - data_dict = data.to_dict("list") if hasattr(routine, "creation_ts"): suffix = routine.creation_ts else: # compatibility with old routines + data = routine.sorted_data + if data is None: + raise ValueError("Unable to get timestamp. Sorted Data is None") + data_dict = data.to_dict("list") ts_float = data_dict["timestamp"][0] # time of the first evaluated point suffix = ts_float_to_str(ts_float, format) + env_name = routine.environment.name fname = f"{env_name}-{suffix}.yaml" return fname def get_header(routine: "Routine") -> list[str]: + obj_names: list[str] = [] + var_names: list[str] = [] + con_names: list[str] = [] + sta_names: list[str] = [] + try: obj_names = routine.vocs.objective_names except AttributeError: @@ -279,8 +297,14 @@ def convert_str_to_value(s: str) -> str | int | float | bool: return s -def parse_rule(rule: dict | str) -> dict: - if type(rule) is str: +class Rule(TypedDict): + direction: str + filter: str + reducer: str + + +def parse_rule(rule: Rule | str) -> Rule: + if isinstance(rule, str): return { "direction": rule, "filter": "ignore_nan", @@ -301,14 +325,10 @@ def parse_rule(rule: dict | str) -> dict: except KeyError: reducer = "percentile_80" - return { - "direction": direction, - "filter": filter, - "reducer": reducer, - } + return Rule(direction=direction, filter=filter, reducer=reducer) -def get_value_or_none(book, key): +def get_value_or_none(book: dict[str, Any], key: str) -> Any: try: value = book[key] except KeyError: @@ -317,7 +337,7 @@ def get_value_or_none(book, key): return value -def dump_state(dump_file, generator, data): +def dump_state(dump_file: str | None, generator: "Generator", data: DataFrame) -> None: """dump data to file""" if dump_file is not None: output = state_to_dict(generator, data) @@ -326,9 +346,17 @@ def dump_state(dump_file, generator, data): logger.debug(f"Dumped state to YAML file: {dump_file}") -def state_to_dict(generator, data, include_data=True): +class StateDict(TypedDict, total=False): + generator: dict[str, Any] + vocs: dict[str, Any] + data: dict[str, Any] | None + + +def state_to_dict( + generator: "Generator", data: DataFrame, include_data: bool = True +) -> StateDict: # dump data to dict with config metadata - output = { + output: StateDict = { "generator": { "name": type(generator).name, type(generator).name: json.loads(generator.model_dump_json()), @@ -351,7 +379,7 @@ def strtobool(val: str) -> bool: try: val = val.lower() except AttributeError: - return val + return bool(val) if val in ("y", "yes", "t", "true", "on", "1"): return True @@ -380,13 +408,14 @@ def get_datadir() -> pathlib.Path: return home / ".local/share" elif sys.platform == "darwin": return home / "Library/Application Support" + return home -def get_badger_version(): +def get_badger_version() -> str: return metadata.version("badger-opt") -def get_xopt_version(): +def get_xopt_version() -> str: return metadata.version("xopt") @@ -440,4 +469,4 @@ def _round_bounds_inward( lower = _round_to_sigfig(float(bounds[0]), sigfigs, mode="ceil") upper = _round_to_sigfig(float(bounds[1]), sigfigs, mode="floor") - return [lower, upper] + return lower, upper