From 21f8ac9279d85b84883c87da0826bf66e9cd0059 Mon Sep 17 00:00:00 2001 From: Mitchell Victoriano Date: Mon, 9 Feb 2026 14:48:25 -0800 Subject: [PATCH 01/15] Adding mypy requirement to pre-commit --- .pre-commit-config.yaml | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 8d919046..5bb4af03 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -22,3 +22,8 @@ repos: - id: ruff-check args: [--fix] - id: ruff-format + - repo: https://github.com/pre-commit/mirrors-mypy + rev: v1.19.1 + hooks: + - id: mypy + args: [--strict, --ignore-missing-imports] From 8e337ff468f8ac4b8e6b3eb0a2e78fd702b65183 Mon Sep 17 00:00:00 2001 From: Mitchell Victoriano Date: Mon, 9 Feb 2026 14:51:15 -0800 Subject: [PATCH 02/15] Updated pre-commit hooks --- .pre-commit-config.yaml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 5bb4af03..dfc683ab 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -2,7 +2,7 @@ # See https://pre-commit.com/hooks.html for more hooks repos: - repo: https://github.com/pre-commit/pre-commit-hooks.git - rev: v5.0.0 + rev: v6.0.0 hooks: - id: no-commit-to-branch - id: trailing-whitespace @@ -17,7 +17,7 @@ repos: exclude: "^(python/src/badger/_version.py)$" - repo: https://github.com/astral-sh/ruff-pre-commit - rev: v0.12.2 + rev: v0.15.0 hooks: - id: ruff-check args: [--fix] From d1f4da6cb5919e5faabe81198ae35fe438a36957 Mon Sep 17 00:00:00 2001 From: Mitchell Victoriano Date: Tue, 17 Feb 2026 20:51:15 -0800 Subject: [PATCH 03/15] Adding more type hints --- .pre-commit-config.yaml | 10 +-- src/badger/actions/__init__.py | 3 +- src/badger/actions/doctor.py | 4 +- src/badger/actions/run.py | 7 +- src/badger/gui/__init__.py | 8 +- src/badger/gui/components/pydantic_editor.py | 59 ++++++------ src/badger/log.py | 70 ++++++++------- src/badger/settings.py | 7 +- src/badger/stats.py | 27 +++--- src/badger/tests/utils.py | 41 +++++---- src/badger/tests/x-test_routine_id.py | 11 ++- src/badger/utils.py | 95 ++++++++++++-------- 12 files changed, 199 insertions(+), 143 deletions(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index dfc683ab..140f1e17 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -22,8 +22,8 @@ repos: - id: ruff-check args: [--fix] - id: ruff-format - - repo: https://github.com/pre-commit/mirrors-mypy - rev: v1.19.1 - hooks: - - id: mypy - args: [--strict, --ignore-missing-imports] + # - repo: https://github.com/pre-commit/mirrors-mypy + # rev: v1.19.1 + # hooks: + # - id: mypy + # args: [--strict, --ignore-missing-imports] diff --git a/src/badger/actions/__init__.py b/src/badger/actions/__init__.py index 42d97ed9..5cb3e184 100644 --- a/src/badger/actions/__init__.py +++ b/src/badger/actions/__init__.py @@ -1,9 +1,10 @@ from importlib import metadata +from argparse import Namespace from badger.actions.doctor import check_n_config_paths from badger.utils import yprint -def show_info(args): +def show_info(args: Namespace): config_path = None if args.config_filepath: diff --git a/src/badger/actions/doctor.py b/src/badger/actions/doctor.py index 2c5c79b1..d70b7429 100644 --- a/src/badger/actions/doctor.py +++ b/src/badger/actions/doctor.py @@ -31,7 +31,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): if config_filepath is not None: config = init_settings(config_filepath) else: @@ -39,7 +39,7 @@ 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 = [] + issue_list: list[str] = [] for pname, pvalue in config._config.model_dump(by_alias=True).items(): if config.read_value(pname) is None: diff --git a/src/badger/actions/run.py b/src/badger/actions/run.py index c80bdc61..d5c7bfef 100644 --- a/src/badger/actions/run.py +++ b/src/badger/actions/run.py @@ -16,7 +16,12 @@ 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: int = 0, + flush_prompt: bool = False, ): try: from badger.archive import archive_run diff --git a/src/badger/gui/__init__.py b/src/badger/gui/__init__.py index 6257e084..c4142660 100644 --- a/src/badger/gui/__init__.py +++ b/src/badger/gui/__init__.py @@ -13,7 +13,7 @@ import traceback from badger.errors import BadgerError from types import TracebackType -from typing import Type, NoReturn +from typing import Any, Type, NoReturn import logging logger = logging.getLogger(__name__) @@ -32,10 +32,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): if TIMER["time"] is None: print("Press Ctrl/Cmd + C again within 1s to quit Badger") TIMER["time"] = time.time() @@ -81,7 +81,7 @@ def error_handler( raise BadgerError(error_title, error_msg) -def launch_gui(config_path=None): +def launch_gui(config_path: str | None = None): sys.excepthook = error_handler app = QApplication(sys.argv) diff --git a/src/badger/gui/components/pydantic_editor.py b/src/badger/gui/components/pydantic_editor.py index 044d70b2..4fbce157 100644 --- a/src/badger/gui/components/pydantic_editor.py +++ b/src/badger/gui/components/pydantic_editor.py @@ -84,7 +84,7 @@ def convert_to_type(value: Any, type: Callable[[Any], T]) -> T: def _set_value_for_basic_widget( widget: Any, value: str | float | int | bool | None, -): +) -> None: if isinstance(widget, QLabel) or isinstance(widget, QLineEdit): widget.setText("null" if value is None else str(value)) elif isinstance(widget, QDoubleSpinBox): @@ -308,7 +308,7 @@ def resolve_qt( return widget -def handle_changed(editor_info: tuple["BadgerPydanticEditor", QTreeWidgetItem]): +def handle_changed(editor_info: tuple["BadgerPydanticEditor", QTreeWidgetItem]) -> None: tree_widget, _ = editor_info tree_widget.validate() @@ -414,7 +414,9 @@ def _qt_widgets_to_values_recurse( 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) @@ -458,13 +460,13 @@ def __init__(self, editor: "BadgerListEditor", parent: QWidget | None = None): remove_button.clicked.connect(self.remove) layout.addWidget(remove_button, alignment=Qt.AlignmentFlag.AlignRight) - def parameter1(self): + def parameter1(self) -> QWidget | None: return self.parameter_value - def parameter2(self): + def parameter2(self) -> QWidget | None: return self.parameter_value2 - def remove(self): + def remove(self) -> None: self.setParent(None) self.deleteLater() self.editor.listChanged.emit() @@ -478,7 +480,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 @@ -504,16 +506,16 @@ def __init__( layout.addLayout(button_layout) - def handle_button_click(self): + def handle_button_click(self) -> None: self.add_widget() - def add_widget(self): + def add_widget(self) -> BadgerListItem: widget = BadgerListItem(self) self.list_layout.addWidget(widget) self.listChanged.emit() return widget - def get_parameters_yaml(self): + def get_parameters_yaml(self) -> str: child_values: list[str | None] = [ _qt_widget_to_yaml_value(child.parameter_value) for child in self.list_container.children() @@ -556,6 +558,8 @@ def get_parameters_dict(self): 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)} @@ -570,7 +574,7 @@ def __init__( parent: QTreeWidget | None = None, value_col: int = 1, update_callback: Callable[["BadgerPydanticEditor"], None] | None = None, - ): + ) -> None: QTreeWidget.__init__(self, parent) if value_col < 1: value_col = 1 @@ -593,7 +597,7 @@ def _set_params_recurse( fields: dict[str, FieldInfo], defaults: dict[str, Any] | None, hidden: bool, - ): + ) -> None: for field_name, field_info in fields.items(): child = QTreeWidgetItem( [field_name if i == 0 else "" for i in range(0, self.value_col + 1)] @@ -645,7 +649,7 @@ def initialize_combo_widget( self, widget: QComboBox, selections: Sequence[type[BaseModel] | None], - ): + ) -> None: # Clear out existing children widget.clear() for selection in selections: @@ -654,14 +658,14 @@ def initialize_combo_widget( else: widget.addItem(selection.model_fields["name"].default, selection) - def set_params_from_class(self, pydantic_class: type[Any]): + def set_params_from_class(self, pydantic_class: type[Any]) -> None: self.clear() self.model_class = pydantic_class self._set_params_recurse(None, self.model_class.model_fields, None, False) self.set_params_post_setup({}) self.validate() - def set_params_from_dict(self, params: dict[str, Any]): + def set_params_from_dict(self, params: dict[str, Any]) -> None: self.clear() field_definitions: dict[str, tuple[type, FieldInfo]] = { k: (type(v), Field()) for k, v in params.items() @@ -684,7 +688,7 @@ def set_params_from_generator( defaults: dict[str, Any], vocs: VOCS | None = None, validate: bool = True, - ): + ) -> None: logger.debug(f"vocs: {vocs}") logger.debug(f"defaults: {defaults}") self.vocs = vocs or VOCS() @@ -722,7 +726,7 @@ def set_params_from_generator( if validate: self.validate() - def set_params_post_setup(self, defaults: dict[str, Any]): + def set_params_post_setup(self, defaults: dict[str, Any]) -> None: if self.model_class is None: raise ValueError("Model class is not set.") @@ -733,7 +737,7 @@ def set_params_post_setup(self, defaults: dict[str, Any]): if self.model_class.model_fields.get("numerical_optimizer") is not None: self.initialize_special_field(defaults, "numerical_optimizer") - def initialize_special_field(self, defaults: dict[str, Any], field: str): + def initialize_special_field(self, defaults: dict[str, Any], field: str) -> None: widget_items = self.findItems(field, Qt.MatchFlag.MatchExactly) if len(widget_items) == 0: @@ -781,7 +785,9 @@ def initialize_special_field(self, defaults: dict[str, Any], field: str): lambda: self.on_radio_changed(special_item, "turbo_controller") ) - def get_all_compatible_classes(self, field_name: str): + def get_all_compatible_classes( + self, field_name: str + ) -> Sequence[type[BaseModel] | None]: if self.model_class is None: raise ValueError("Model class is not set.") @@ -825,7 +831,7 @@ def on_radio_changed( self, tree_widget_item: QTreeWidgetItem, field_name: str, - ): + ) -> None: # Clear out existing children for cc in tree_widget_item.takeChildren(): del cc @@ -857,7 +863,7 @@ def update_params_from_generator_class( name: str, field_name: str, defaults: dict[str, Any], - ): + ) -> None: try: pydantic_class = self.get_compatible_class(name, field_name) except ValueError: @@ -920,7 +926,7 @@ def exclude_condition(k: str) -> bool: return filtered_class_fields, removed_class_fields @staticmethod - def get_defaults_from_type(pydantic_class: type[Any]): + def get_defaults_from_type(pydantic_class: type[Any]) -> dict[str, Any]: if not issubclass(pydantic_class, BaseModel): raise ValueError("Provided class is not a Pydantic model") defaults: dict[str, Any] = {} @@ -972,7 +978,7 @@ def find_widget_at_path( return None - def remove_style(self, item: QTreeWidgetItem | None): + def remove_style(self, item: QTreeWidgetItem | None) -> None: # Have to reset border styling in case some errors were fixed if item is None: return @@ -984,7 +990,7 @@ def remove_style(self, item: QTreeWidgetItem | None): for j in range(item.childCount()): self.remove_style(item.child(j)) - def update_vocs(self, vocs: VOCS): + def update_vocs(self, vocs: VOCS) -> None: logger.debug(f"Updating VOCS in BadgerPydanticEditor: {vocs}") self.vocs = vocs @@ -996,7 +1002,7 @@ def update_vocs(self, vocs: VOCS): self.set_params_from_generator(self.generator_name, defaults, self.vocs) self.validate() - def update_after_validate(self, defaults: dict[str, Any]): + def update_after_validate(self, defaults: dict[str, Any]) -> None: model_class = self.model_class if model_class is None: return @@ -1029,7 +1035,7 @@ def update_after_validate(self, defaults: dict[str, Any]): if self.update_callback is not None: self.update_callback(self) - def validate(self): + def validate(self) -> bool: if self.model_class is None: return False @@ -1056,6 +1062,7 @@ def validate(self): return True except KeyError as e: logger.error(e) + return False except ValidationError as e: logger.error(e) diff --git a/src/badger/log.py b/src/badger/log.py index a33dea79..0c82dbda 100644 --- a/src/badger/log.py +++ b/src/badger/log.py @@ -1,10 +1,12 @@ +from __future__ import annotations +from argparse import Namespace import os import datetime import logging import atexit from logging.handlers import QueueHandler, QueueListener -from multiprocessing import Queue +from multiprocessing import Queue, current_process from badger.settings import get_user_config_folder from badger.settings import init_settings @@ -28,12 +30,12 @@ class LoggingManager: 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[str] = Queue() + 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. @@ -49,7 +51,7 @@ def start_listener(self, log_filepath: str, log_level: str | int): # Queue for sending all the logs to self.log_queue = Queue() - self.handlers = [] + self.handlers: list[logging.Handler] = [] # File handler file_handler = logging.FileHandler(log_filepath, mode="a") file_formatter = logging.Formatter( @@ -74,11 +76,11 @@ 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 = logging.getLevelName(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") @@ -107,11 +109,12 @@ 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." - ) + # deprecated, getLevelNamesMapping()[log_level] is the replacement + log_level_name = logging.getLevelName(log_level) + + logger.info(f"Log level updated to {log_level_name} for badger namespace.") - def update_logfile_path(self, new_logfile_path: str): + def update_logfile_path(self, new_logfile_path: str) -> None: """ Update logfile path, so logs get written to a different file. @@ -120,12 +123,17 @@ 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: + 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 log_level = file_handler.level @@ -152,7 +160,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" @@ -161,11 +169,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[str]: """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. @@ -173,7 +181,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 @@ -195,11 +203,11 @@ def create_log_dir(self, log_dir: str): def configure_process_logging( - log_queue: Queue = None, + log_queue: Queue[str] | None = None, logger_name: str = "badger", - log_level: str = "DEBUG", - process_name: str = None, -): + log_level: str | int = "DEBUG", + process_name: str | None = None, +) -> None: """ Configure logging in a process to send logs to the shared queue. This should be ran in all subprocesses b4 logging, and also in the main process where @@ -208,16 +216,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) # 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) @@ -238,12 +244,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 = logging.getLevelName(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 """ diff --git a/src/badger/settings.py b/src/badger/settings.py index 7438f9d8..3980f1d6 100644 --- a/src/badger/settings.py +++ b/src/badger/settings.py @@ -124,9 +124,12 @@ 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, user_flag: bool = False): + def __new__(cls, config_path: str, user_flag: bool = False): if cls._instance is None: cls._instance = super(ConfigSingleton, cls).__new__(cls) cls._instance.user_flag = user_flag diff --git a/src/badger/stats.py b/src/badger/stats.py index 310a5f22..af13c034 100644 --- a/src/badger/stats.py +++ b/src/badger/stats.py @@ -1,45 +1,48 @@ +from typing import Any + import numpy as np +from numpy.typing import NDArray -def none(data): +def none(data: NDArray[np.floating[Any]]) -> NDArray[np.floating[Any]]: return data -def median(data): +def median(data: NDArray[np.floating[Any]]) -> np.floating[Any]: return np.median(data) -def std_deviation(data): +def std_deviation(data: NDArray[np.floating[Any]]) -> np.floating[Any]: return np.std(data) -def median_deviation(data): - median = np.median(data) +def median_deviation(data: NDArray[np.floating[Any]]) -> np.floating[Any]: + med = np.median(data) - return np.median(np.abs(data - median)) + return np.median(np.abs(data - med)) -def max(data): +def max(data: NDArray[np.floating[Any]]) -> np.floating[Any]: return np.max(data) -def min(data): +def min(data: NDArray[np.floating[Any]]) -> np.floating[Any]: return np.min(data) -def percent_80(data): +def percent_80(data: NDArray[np.floating[Any]]) -> np.floating[Any]: return np.percentile(data, 80) -def percent_20(data): +def percent_20(data: NDArray[np.floating[Any]]) -> np.floating[Any]: return np.percentile(data, 20) -def avg_mean(data): +def avg_mean(data: NDArray[np.floating[Any]]) -> np.floating[Any]: percentile = np.percentile(data, 50) return np.mean(data[data > percentile]) -def mean(data): +def mean(data: NDArray[np.floating[Any]]) -> np.floating[Any]: return np.mean(data) diff --git a/src/badger/tests/utils.py b/src/badger/tests/utils.py index 737d9012..4ec9d6e0 100644 --- a/src/badger/tests/utils.py +++ b/src/badger/tests/utils.py @@ -1,14 +1,17 @@ import os +from typing import Any +import numpy as np import pandas as pd from xopt import VOCS from xopt.generators.bayesian import UpperConfidenceBoundGenerator from xopt.generators import RandomGenerator +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": {}, @@ -64,10 +66,9 @@ def create_multiobjective_routine(): ) -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", @@ -105,10 +106,9 @@ def create_routine_turbo(): ) -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", @@ -136,10 +136,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", @@ -175,19 +174,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 = 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.settings import init_settings from badger.archive import BADGER_ARCHIVE_ROOT diff --git a/src/badger/tests/x-test_routine_id.py b/src/badger/tests/x-test_routine_id.py index 943fd711..4190fd2a 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 save_routine, remove_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.gui.components.routine_page import BadgerRoutinePage from badger.db import list_routine, remove_routine, load_routine @@ -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.gui.components.routine_page import BadgerRoutinePage from badger.db import list_routine, remove_routine, load_routine, save_run @@ -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.gui.components.routine_page import BadgerRoutinePage from badger.db import list_routine, remove_routine, load_routine, save_run diff --git a/src/badger/utils.py b/src/badger/utils.py index cb58342c..b4f7e025 100644 --- a/src/badger/utils.py +++ b/src/badger/utils.py @@ -6,13 +6,18 @@ import pathlib from datetime import datetime from types import TracebackType -from typing import Iterable, Optional +from typing import TYPE_CHECKING, Any, Iterable, Optional, TypedDict +from pandas import DataFrame import yaml from badger.errors import BadgerLoadConfigError from PyQt5.QtWidgets import QWidget, QLayout +if TYPE_CHECKING: + from xopt.generators import Generator + from badger.routine import Routine + logger = logging.getLogger(__name__) @@ -50,34 +55,36 @@ 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): return super(Dumper, self).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) -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: Optional[Iterable[dict[str, Any]]], +) -> 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 @@ -85,7 +92,7 @@ def config_list_to_dict(config_list): return book -def load_config(fname): +def load_config(fname: Optional[str]) -> Optional[dict[str, Any]]: configs = None if fname is None: @@ -115,7 +122,9 @@ def load_config(fname): return configs -def merge_params(default_params, params): +def merge_params( + default_params: Optional[dict[str, Any]], params: Optional[dict[str, Any]] +) -> Optional[dict[str, Any]]: merged_params = None if params is None: @@ -128,9 +137,9 @@ def merge_params(default_params, params): return merged_params -def range_to_str(vranges): +def range_to_str(vranges: Iterable[dict[str, Any]]) -> 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] @@ -140,7 +149,7 @@ def range_to_str(vranges): return vranges_str -def ts_to_str(ts, format="lcls-log"): +def ts_to_str(ts: datetime, format: str = "lcls-log") -> str: if format == "lcls-log": return ts.strftime("%d-%b-%Y %H:%M:%S") elif format == "lcls-log-full": @@ -151,7 +160,7 @@ def ts_to_str(ts, format="lcls-log"): return ts.isoformat() -def str_to_ts(timestr, format="lcls-log"): +def str_to_ts(timestr: str, format: str = "lcls-log") -> datetime: if format == "lcls-log": return datetime.strptime(timestr, "%d-%b-%Y %H:%M:%S") elif format == "lcls-log-full": @@ -162,30 +171,30 @@ def str_to_ts(timestr, format="lcls-log"): return datetime.fromisoformat(timestr) -def ts_float_to_str(ts_float, format="lcls-log"): +def ts_float_to_str(ts_float: float, format: str = "lcls-log") -> str: ts = datetime.fromtimestamp(ts_float) return ts_to_str(ts, format) -def curr_ts(): +def curr_ts() -> datetime: return datetime.now() -def curr_ts_to_str(format="lcls-log"): +def curr_ts_to_str(format: str = "lcls-log") -> str: return ts_to_str(datetime.now(), format) -def create_archive_run_filename(routine, format: str = "lcls-fname") -> 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") + data_dict = data.to_dict("list") # type: ignore ts_float = data_dict["timestamp"][0] # time of the first evaluated point suffix = ts_float_to_str(ts_float, format) fname = f"{env_name}-{suffix}.yaml" return fname -def get_header(routine): +def get_header(routine: "Routine") -> list[str]: try: obj_names = routine.vocs.objective_names except Exception: @@ -206,8 +215,10 @@ def get_header(routine): return obj_names + con_names + var_names + sta_names -def run_names_to_dict(run_names): - runs = {} +def run_names_to_dict( + run_names: Iterable[str], +) -> dict[str, dict[str, dict[str, list[str]]]]: + runs: dict[str, dict[str, dict[str, list[str]]]] = {} for name in run_names: name = os.path.basename( name @@ -219,19 +230,19 @@ def run_names_to_dict(run_names): try: year_dict = runs[year] - except Exception: + except KeyError: runs[year] = {} year_dict = runs[year] key_month = f"{year}-{month}" try: month_dict = year_dict[key_month] - except Exception: + except KeyError: year_dict[key_month] = {} month_dict = year_dict[key_month] key_day = f"{year}-{month}-{day}" try: day_list = month_dict[key_day] - except Exception: + except KeyError: month_dict[key_day] = [] day_list = month_dict[key_day] day_list.append(name) @@ -239,7 +250,7 @@ def run_names_to_dict(run_names): return runs -def convert_str_to_value(str): +def convert_str_to_value(str: str) -> Any: try: return int(str) except ValueError: @@ -258,8 +269,14 @@ def convert_str_to_value(str): return str -def parse_rule(rule): - 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", @@ -287,7 +304,7 @@ def parse_rule(rule): } -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: @@ -296,7 +313,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): """dump data to file""" if dump_file is not None: output = state_to_dict(generator, data) @@ -305,9 +322,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: Optional[dict[str, Any]] + + +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()), @@ -315,13 +340,13 @@ def state_to_dict(generator, data, include_data=True): "vocs": json.loads(generator.vocs.model_dump_json()), } if include_data: - output["data"] = json.loads(data.to_json()) + output["data"] = json.loads(data.to_json()) # type: ignore return output # https://stackoverflow.com/a/18472142 -def strtobool(val): +def strtobool(val: str) -> bool: """Convert a string representation of truth to true (1) or false (0). True values are 'y', 'yes', 't', 'true', 'on', and '1'; false values are 'n', 'no', 'f', 'false', 'off', and '0'. Raises ValueError if @@ -330,7 +355,7 @@ def strtobool(val): try: val = val.lower() except AttributeError: - return val + return bool(val) if val in ("y", "yes", "t", "true", "on", "1"): return True From 8bb728ab0df9a4c5c3a00c999e2f635612a41bd3 Mon Sep 17 00:00:00 2001 From: Mitchell Victoriano Date: Tue, 9 Jun 2026 19:16:18 -0700 Subject: [PATCH 04/15] WIP: Type hinting improvements --- src/badger/__main__.py | 2 +- src/badger/actions/__init__.py | 2 +- src/badger/actions/config.py | 18 ++++++++++++----- src/badger/actions/doctor.py | 13 ++++++++++--- src/badger/actions/env.py | 3 ++- src/badger/actions/generator.py | 3 ++- src/badger/actions/routine.py | 3 ++- src/badger/db.py | 34 +++++++++++++++++++++------------ src/badger/extension.py | 24 +++++++++++++---------- src/badger/factory.py | 34 +++++++++++++++++++++++---------- src/badger/tests/x-test_db.py | 2 +- 11 files changed, 92 insertions(+), 46 deletions(-) diff --git a/src/badger/__main__.py b/src/badger/__main__.py index 614fcfb4..449b9e53 100644 --- a/src/badger/__main__.py +++ b/src/badger/__main__.py @@ -15,7 +15,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 5cb3e184..3c48e0c0 100644 --- a/src/badger/actions/__init__.py +++ b/src/badger/actions/__init__.py @@ -4,7 +4,7 @@ from badger.utils import yprint -def show_info(args: Namespace): +def show_info(args: Namespace) -> None: config_path = None if args.config_filepath: diff --git a/src/badger/actions/config.py b/src/badger/actions/config.py index ca3afb38..1068075d 100644 --- a/src/badger/actions/config.py +++ b/src/badger/actions/config.py @@ -1,3 +1,5 @@ +from argparse import Namespace + from badger.settings import init_settings import os import logging @@ -7,13 +9,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: @@ -25,12 +27,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) @@ -96,7 +99,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) @@ -128,6 +131,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 d70b7429..46ed0938 100644 --- a/src/badger/actions/doctor.py +++ b/src/badger/actions/doctor.py @@ -1,8 +1,10 @@ +from argparse import Namespace + from badger.settings import init_settings, mock_settings from badger.actions.config import _config_path_var -def self_check(args): +def self_check(args: Namespace) -> None: config = init_settings() # Reset Badger @@ -31,7 +33,7 @@ def self_check(args): print("Badger is healthy!") -def check_n_config_paths(config_filepath: str | None = None): +def check_n_config_paths(config_filepath: str | None = None) -> bool: if config_filepath is not None: config = init_settings(config_filepath) else: @@ -41,7 +43,12 @@ def check_n_config_paths(config_filepath: str | None = None): all_bad = True # if all config paths are empty, we'll suggest initialization issue_list: list[str] = [] - for pname, pvalue in config._config.model_dump(by_alias=True).items(): + 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 4d8d7f4c..3ac46eed 100644 --- a/src/badger/actions/env.py +++ b/src/badger/actions/env.py @@ -1,3 +1,4 @@ +from argparse import Namespace import logging from badger.utils import range_to_str, yprint @@ -5,7 +6,7 @@ logger = logging.getLogger(__name__) -def show_env(args): +def show_env(args: Namespace) -> None: try: from badger.factory import list_env, get_env except Exception as e: diff --git a/src/badger/actions/generator.py b/src/badger/actions/generator.py index 4d42877d..e0e41719 100644 --- a/src/badger/actions/generator.py +++ b/src/badger/actions/generator.py @@ -1,3 +1,4 @@ +from argparse import Namespace import logging from badger.utils import yprint @@ -5,7 +6,7 @@ logger = logging.getLogger(__name__) -def show_generator(args): +def show_generator(args: Namespace) -> None: try: from badger.factory import list_generators, get_generator except Exception as e: diff --git a/src/badger/actions/routine.py b/src/badger/actions/routine.py index b1476551..5c545220 100644 --- a/src/badger/actions/routine.py +++ b/src/badger/actions/routine.py @@ -1,3 +1,4 @@ +from argparse import Namespace import logging import pandas as pd @@ -8,7 +9,7 @@ logger = logging.getLogger(__name__) -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/db.py b/src/badger/db.py index 32ad4b2f..43223bd5 100644 --- a/src/badger/db.py +++ b/src/badger/db.py @@ -1,4 +1,5 @@ import os +from typing import Any, Callable, ParamSpec, TypeVar, cast import warnings from datetime import datetime import logging @@ -16,26 +17,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) @@ -53,12 +63,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) @@ -107,7 +117,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) @@ -126,7 +136,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) @@ -147,7 +157,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) @@ -172,7 +182,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() @@ -207,7 +217,7 @@ def load_routine(id: str): @ensure_routines_db_exists -def list_routine(keyword="", tags={}): +def list_routine(keyword: str = "", tags: dict[str, str] = {}) -> list[Routine]: db_routine = os.path.join(BADGER_DB_ROOT, "routines.db") con = sqlite3.connect(db_routine) cur = con.cursor() @@ -300,7 +310,7 @@ def list_routine(keyword="", tags={}): @ensure_runs_db_exists -def save_run(run): +def save_run(run: dict[str, Any]) -> int: db_run = os.path.join(BADGER_DB_ROOT, "runs.db") con = sqlite3.connect(db_run, timeout=30.0) diff --git a/src/badger/extension.py b/src/badger/extension.py index 1286c4f6..ca6ceb2b 100644 --- a/src/badger/extension.py +++ b/src/badger/extension.py @@ -1,29 +1,33 @@ from abc import ABC, abstractmethod -from typing import List +from typing import Any, Callable 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 + def list_generator(self) -> list[str]: + 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 1edfdd40..b4d661f3 100644 --- a/src/badger/factory.py +++ b/src/badger/factory.py @@ -1,4 +1,4 @@ -from typing import Any, TypedDict, cast, TYPE_CHECKING +from typing import Any, Literal, TypedDict, cast, TYPE_CHECKING from badger.settings import init_settings from badger.utils import get_value_or_none from badger.errors import ( @@ -13,6 +13,7 @@ import os import importlib import yaml +from xopt.generator import Generator from xopt.generators import generators, get_generator_defaults import logging @@ -34,7 +35,7 @@ ] -class BadgerPluginConfig(TypedDict): +class BadgerPluginConfig(TypedDict, total=False): name: str description: str version: str @@ -62,10 +63,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: @@ -73,7 +85,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") @@ -89,14 +101,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] | type[BadgerInterface] | type[Generator] | None, BadgerPluginConfig | None]": assert ptype in [ "generator", "interface", @@ -183,7 +195,7 @@ def load_plugin( return plugin -def load_docs(root, pname, ptype): +def load_docs(root: str, pname: str, ptype: str) -> str: assert ptype in [ "generator", "interface", @@ -215,7 +227,9 @@ def load_docs(root, pname, ptype): ) -def get_plug(root: str, name: str, ptype: str): +def get_plug( + root: str, name: str, ptype: Literal["environment", "interface", "generator"] +) -> "tuple[type[BadgerEnvironment] | type[BadgerInterface] | type[Generator] | None, BadgerPluginConfig | None]": try: plug = BADGER_FACTORY[ptype][name] if plug is None: # lazy loading diff --git a/src/badger/tests/x-test_db.py b/src/badger/tests/x-test_db.py index 25ccfd23..d314f07c 100644 --- a/src/badger/tests/x-test_db.py +++ b/src/badger/tests/x-test_db.py @@ -1,5 +1,5 @@ class TestDB: - def test_save_routine(self): + def test_save_routine(self) -> None: from badger.db import save_routine, remove_routine from badger.tests.utils import create_routine, fix_db_path_issue From afd5ebc161c0d3dd7307b6af99156525b4e4da86 Mon Sep 17 00:00:00 2001 From: Mitchell Victoriano Date: Tue, 9 Jun 2026 20:54:36 -0700 Subject: [PATCH 05/15] First pass at addressing the various mypy type errors using claude --- .pre-commit-config.yaml | 9 +- mypy_deferred_errors.md | 327 ++++++++++++++++++++++++++++++++ pyproject.toml | 4 +- src/badger/__main__.py | 2 +- src/badger/actions/__init__.py | 3 +- src/badger/actions/config.py | 7 +- src/badger/actions/doctor.py | 6 +- src/badger/actions/env.py | 5 +- src/badger/actions/generator.py | 5 +- src/badger/actions/install.py | 3 +- src/badger/actions/intf.py | 5 +- src/badger/actions/routine.py | 3 +- src/badger/actions/run.py | 24 ++- src/badger/actions/uninstall.py | 3 +- src/badger/core.py | 28 ++- src/badger/errors.py | 12 +- src/badger/extension.py | 10 +- src/badger/factory.py | 38 ++-- src/badger/formula.py | 20 +- src/badger/interface.py | 30 +-- src/badger/logger/observer.py | 14 +- src/badger/logger/util.py | 22 +-- src/badger/stats.py | 23 ++- src/badger/utils.py | 72 +++---- 24 files changed, 526 insertions(+), 149 deletions(-) create mode 100644 mypy_deferred_errors.md diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 8d919046..697dc9b0 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -2,7 +2,7 @@ # See https://pre-commit.com/hooks.html for more hooks repos: - repo: https://github.com/pre-commit/pre-commit-hooks.git - rev: v5.0.0 + rev: v6.0.0 hooks: - id: no-commit-to-branch - id: trailing-whitespace @@ -17,8 +17,13 @@ repos: exclude: "^(python/src/badger/_version.py)$" - repo: https://github.com/astral-sh/ruff-pre-commit - rev: v0.12.2 + rev: v0.15.16 hooks: - id: ruff-check args: [--fix] - id: ruff-format + # - repo: https://github.com/pre-commit/mirrors-mypy + # rev: v2.1.0 + # hooks: + # - id: mypy + # args: [--strict, --ignore-missing-imports] diff --git a/mypy_deferred_errors.md b/mypy_deferred_errors.md new file mode 100644 index 00000000..37df6b52 --- /dev/null +++ b/mypy_deferred_errors.md @@ -0,0 +1,327 @@ +# Mypy Deferred Errors + +This document captures the full state of `mypy --strict --ignore-missing-imports` on the badger repository as of 2026-06-09. The first run produced **1506 errors in 116 files**. After this PR's mechanical fixes (annotation-only changes, no runtime behavior changes) the count is **1473 errors in 98 files**. + +The remaining errors fall into two groups: + +1. **Mechanical work not yet done** — straightforward signature additions across GUI components and tests. Listed in §1 with file-by-file scopes for follow-up PRs. +2. **Real issues requiring business-logic review** — not mechanical fixes. Listed in §2 with concrete file:line references and what each one means. + +--- + +## §1. Remaining mechanical work (deferred to follow-up PRs) + +These are pure annotation additions. No runtime behavior changes required. Work can proceed file-by-file in independent PRs. + +### 1a. Test files — add `-> None` to test functions, type fixture parameters as `Any` + +| File | Approx. errors | +| --- | --- | +| `src/badger/tests/test_environment.py` | 40 | +| `src/badger/tests/test_formulas.py` | 39 | +| `src/badger/tests/test_routine_page.py` | 31 | +| `src/badger/tests/test_run_monitor.py` | 27 | +| `src/badger/tests/utils.py` | 17 | +| `src/badger/tests/test_settings.py` | 17 | +| `src/badger/tests/test_routine_runner.py` | 12 | +| `src/badger/tests/test_gui_basic.py` | 11 | +| `src/badger/tests/conftest.py` | 11 | +| `src/badger/tests/test_env.py` | 9 | +| `src/badger/tests/test_routine.py` | 4 | +| `src/badger/tests/test_intf.py` | 7 | +| `src/badger/tests/test_core_subprocess.py` | 8 | +| `src/badger/tests/test_cli_basic.py` | 7 | +| `src/badger/tests/test_core.py` | 6 | +| `src/badger/tests/test_lib_basic.py` | 2 | +| `src/badger/tests/test_factory.py` | 1 | +| `src/badger/tests/test_home_page.py` | 2 | +| `src/badger/tests/test_create_process.py` | 2 | +| `src/badger/tests/test_routine_id.py` (`x-test_routine_id.py`) | 4 | +| `src/badger/tests/x-test_db.py` | 2 | +| `src/badger/tests/test_settings.py` (continued) | — | +| `src/badger/tests/test_measurement_retry_dialog.py` | 2 | +| `src/badger/tests/test_process_manager.py` | 3 | +| `src/badger/tests/multiprocess_logging.py` | 3 | +| `src/badger/tests/mock/plugins/...` | ~10 | + +**Pattern**: +- `def test_xxx(...)` → `def test_xxx(...) -> None` +- `def fixture_xxx(mocker, qtbot, request)` → `def fixture_xxx(mocker: Any, qtbot: Any, request: Any) -> ...` +- `points_eval_list = []` → `points_eval_list: list[Any] = []` +- `out, err = "", ""` after `out: str` annotation in `test_cli_basic.py:8-9`. + +### 1b. GUI windows / dialogs — add signatures to `__init__`, `init_ui`, `config_logic`, slot handlers + +| File | Approx. errors | +| --- | --- | +| `src/badger/gui/windows/load_data_from_run_dialog.py` | 17 | +| `src/badger/gui/windows/settings_dialog.py` | 18 | +| `src/badger/gui/windows/ind_lim_vrange_dialog.py` | 14 | +| `src/badger/gui/windows/lim_vrange_dialog.py` | 11 | +| `src/badger/gui/windows/add_random_dialog.py` | 11 | +| `src/badger/gui/windows/main_window.py` | 12 | +| `src/badger/gui/windows/docs_window.py` | 11 | +| `src/badger/gui/windows/expandable_message_box.py` | 7 | +| `src/badger/gui/windows/edit_script_dialog.py` | 9 | +| `src/badger/gui/windows/var_dialog.py` | 7 | +| `src/badger/gui/windows/terminition_condition_dialog.py` | 11 | +| `src/badger/gui/windows/measurement_retry_dialog.py` | 1 | +| `src/badger/gui/windows/message_dialog.py` | 4 | +| `src/badger/gui/windows/review_dialog.py` | 4 | + +**Pattern**: +- `def __init__(self, ...):` → `def __init__(self, ...) -> None:` +- `def init_ui(self):` → `def init_ui(self) -> None:` +- `def config_logic(self):` → `def config_logic(self) -> None:` +- All event-handler slots → `-> None`. + +### 1c. GUI components/pages — same mechanical patterns + +| File | Approx. errors | +| --- | --- | +| `src/badger/gui/pages/home_page.py` | 40 (mostly logic — see §2) | +| `src/badger/gui/components/run_monitor.py` | 121 (mostly logic — see §2) | +| `src/badger/gui/components/routine_page.py` | 119 (logic-heavy — see §2) | +| `src/badger/gui/components/pydantic_editor.py` | 88 (some logic — see §2) | +| `src/badger/gui/components/var_table.py` | 60 (logic-heavy — see §2) | +| `src/badger/gui/components/data_table.py` | 26 (mostly mechanical) | +| `src/badger/gui/components/navigators.py` | 35 | +| `src/badger/gui/components/archive_search.py` | 26 | +| `src/badger/gui/components/action_bar.py` | 24 | +| `src/badger/gui/components/env_cbox.py` | 22 | +| `src/badger/gui/components/routine_runner.py` | 20 (mostly logic — see §2) | +| `src/badger/gui/components/editable_table.py` | 19 | +| `src/badger/gui/components/collapsible_box.py` | 17 | +| `src/badger/gui/components/bo_visualizer/bo_widget.py` | 17 | +| `src/badger/gui/components/plot_event_handlers.py` | 16 | +| `src/badger/gui/components/bo_visualizer/ui_components.py` | 16 | +| `src/badger/gui/components/syntax.py` | 15 | +| `src/badger/gui/components/routine_item.py` | 15 | +| `src/badger/gui/components/data_panel.py` | 14 (some logic — see §2) | +| `src/badger/gui/components/routine_editor.py` | 12 | +| `src/badger/gui/components/eliding_label.py` | 12 | +| `src/badger/gui/components/pf_viewer/pf_widget.py` | 40 | +| `src/badger/gui/components/process_manager.py` | 4 | +| `src/badger/gui/components/extension_utilities.py` | 2 | +| `src/badger/gui/components/extensions_palette.py` | 11 | +| `src/badger/gui/components/create_process.py` | 4 | +| `src/badger/gui/components/state_item.py` | 3 | +| `src/badger/gui/components/status_bar.py` | 6 | +| `src/badger/gui/components/search_bar.py` | 1 | +| `src/badger/gui/components/labeled_lineedit.py` | 1 | +| `src/badger/gui/components/robust_spinbox.py` | 1 | +| `src/badger/gui/components/reorderable_table.py` | 4 | +| `src/badger/gui/components/obs_table.py` | 1 | +| `src/badger/gui/components/obj_table.py` | 1 | +| `src/badger/gui/components/con_table.py` | — | +| `src/badger/gui/components/filter_cbox.py` | 3 | +| `src/badger/gui/components/generator_cbox.py` | 5 | +| `src/badger/gui/components/constraint_item.py` | 3 | +| `src/badger/gui/components/bo_visualizer/plotting_area.py` | 4 | +| `src/badger/gui/components/analysis_widget.py` | 1 | +| `src/badger/gui/components/analysis_extensions.py` | 1 | +| `src/badger/gui/utils.py` | 9 | +| `src/badger/gui/__init__.py` | 4 (one is logic — see §2) | + +### 1d. Plugin scaffolding + +| File | Errors | +| --- | --- | +| `src/badger/built_in_plugins/interfaces/default/__init__.py` | 5 | +| `src/badger/built_in_plugins/environments/sphere_2d/__init__.py` | 4 | +| `src/badger/tests/mock/plugins/interfaces/test/__init__.py` | 5 | +| `src/badger/tests/mock/plugins/environments/test/__init__.py` | 3 | +| `src/badger/tests/mock/plugins/environments/multiobjective_test/__init__.py` | 2 | + +### 1e. Other core/library modules with remaining mechanical work + +| File | Errors | +| --- | --- | +| `src/badger/db.py` | 27 (some logic — see §2.G) | +| `src/badger/log.py` | 23 | +| `src/badger/settings.py` | 18 (mostly logic — see §2.J) | +| `src/badger/routine.py` | 17 (some logic — see §2.E) | +| `src/badger/environment.py` | 29 (some logic — see §2.F) | +| `src/badger/core_subprocess.py` | 29 (mostly logic — see §2.H) | +| `src/badger/archive.py` | 9 | +| `src/badger/logger/__init__.py` | 29 | +| `src/badger/logbook.py` | 4 (one logic — see §2.K) | + +--- + +## §2. Errors requiring business-logic review (NOT addressed in this PR) + +These are real type-correctness issues — fixing them changes runtime behavior or requires a design decision. Each entry: file, line(s), what mypy says, why it's deferred. + +### A. Truthy-function checks in `core.py` + +`src/badger/core.py:108,132,136,165,169` — code like: +```python +if states_callback and (states is not None): +if evaluate_callback: +if dump_file_callback: +``` +Mypy: `[truthy-function]`. A function object is always truthy — these checks always pass. The intent appears to be `is not None`, but the parameters are typed `Callable[..., Any]` (non-Optional) except `dump_file_callback`. **Decision needed**: does each call site pass `None` for these callbacks, or always a callable? If never `None`, drop the checks; if sometimes `None`, change type to `Optional` and check `is not None`. + +### B. Unguarded `Optional` attribute access — `run_monitor.py` + +`src/badger/gui/components/run_monitor.py` — `self.routine` is initialized to `None` and accessed without None checks throughout. Specific lines: +- 109 (`vocs` property) +- 467, 469, 471, 473, 475 (`start()`) +- 503, 528, 538, 547 (data accessors) +- 602, 605, 617 (constraint plotting) +- 654, 655 (run_filename access) +- 688, 693 (routine + environment access) +- 795, 901, 909 (more environment / vocs) +- 956, 968, 969, 970 (vrange_hard_limit) +- 1053, 1107, 1109, 1115–1118, 1149–1152, 1161–1165 (more cascading) + +**Decision needed**: either +- Restructure so `self.routine` becomes non-Optional after a setter is called (and all access goes through that setter), or +- Add `if self.routine is None: return` guards at the top of every method that touches it. + +Also: `run_monitor.py:281–350` has `Cannot determine type of "plot_con"/"plot_obs"` errors — class attributes referenced before being assigned in `__init__`. Same issue as below in `routine_page.py`. These are forward-reference bugs in the class init path. + +### C. Unguarded `Optional` attribute access — `data_panel.py` + +`src/badger/gui/components/data_panel.py:139,170,226,267,272,347,382` — similar pattern: `self.run_monitor` or similar set to `None` then accessed. Plus: +- Line 181: `get_data_from_dialog` declared `-> None` but caller uses return value (`func-returns-value`). +- Line 267, 382: missing dict generic args. + +### D. Unguarded `Optional` access — `home_page.py`, `var_table.py` + +`src/badger/gui/pages/home_page.py:306,546,547,548,596,990` — `Item "None" of "Any | None" has no attribute "data"` etc. +`src/badger/gui/components/var_table.py:540,541,637,645` — `"None" has no attribute "get_variable"/"get_bounds"/"get_info"` and `Item "None" of "Any | None" has no attribute "text"`. + +Same pattern: attribute initialized to `None`, accessed unconditionally. + +### E. `routine.py:101,102` — wrong type annotation on `data` field + +```python +@validate_data +def evaluate_data(self, data: dict) -> dict: + ... + new_index = max(self.data.index, default=-1) + 1 + ... + self.data.sort_index(...) +``` +But `self.data` is annotated as `dict[Any, Any]`. The code expects a pandas DataFrame. **Decision needed**: the field type is wrong; should be `DataFrame | None`. Fixing reveals downstream callers expect dict-style access. + +### F. `environment.py` — multiple typing issues + +- **Line 223**: `def add_observation(self, callable=None) -> bool:` — uses bare `callable` (the builtin) as a type. Replace with `Callable[..., Any] | None`. +- **Line 285**: `__call__` declared `-> float` but actually returns `float | list[float]`. Need to widen return type or split into two methods. +- **Lines 298–324**: `Returning Any from function declared to return "dict"` — internal return-type drift. Several functions have explicit `-> dict[Any, Any]` but actually return `Any`. +- Lines 18, 36, 42, 82–83, 154, 188, 258, 304, 316, 327, 331 — missing annotations (mechanical). + +### G. `db.py:109,128,148,149,174,339` — untyped `@db.commit`-style decorators + +`Untyped decorator makes function "save_routine" untyped`. The SQLAlchemy / decorator library lacks type stubs. **Decision needed**: either accept `# type: ignore[misc]` per decorated function, or wait for upstream stubs. + +Plus `db.py:33,38,56,61,79,92,210,303,358,375,388,400,432` — straightforward missing annotations (mechanical). + +### H. `core_subprocess.py` — multiprocessing primitives used as types + +```python +event: multiprocessing.Event # Event is a factory, not a type +queue: multiprocessing.Queue # missing generic param +pipe: multiprocessing.Pipe # Pipe returns tuple, not a type +``` +Lines 37–39, 132–139, 179, 281, 285, 287, 289, 291, 329–349, 359–375. + +**Fix**: import the actual types from `multiprocessing.synchronize` (`Event`), `multiprocessing.connection` (`Connection`), `queue.Queue`. Non-trivial — affects many call sites and signatures across `log.py`, `routine_runner.py`, `create_process.py`. + +### I. `routine_page.py` — multiple structural issues + +`src/badger/gui/components/routine_page.py`: +- Lines 102, 104, 106: function declared `-> str` but returns tuples. Pick one. +- Lines 477, 824: `Incompatible types in assignment (expression has type "list[dict]", variable has type "dict")` — variable reused with mismatched types. +- Lines 512, 545, 584, 861, 897, 940, 1050, 1098, 1378: `Cannot determine type of "configs" / "archive_search"` — class attributes referenced before assignment in init flow. Init method order is fragile. +- Lines 551, 903: `Unpacking a string is disallowed` — code does `a, b, c = some_string` where the runtime value happens to be a tuple sometimes. Genuine bug or unclear intent. +- Lines 1176, 1280–1282: `List item N has incompatible type "float"/"bool"; expected "str"` — heterogeneous list passed where homogeneous expected. +- Lines 1245, 1273, 1308, 1394: `Item "None" of "Any | None"` not iterable / indexable. +- Lines 1903, 1907, 1911: more None-attribute access. + +### J. `settings.py` — `ConfigSingleton` pattern not modeled + +Lines 137, 140, 142, 194, 202, 212, 219, 247, 391, 406, 481, 492, 541, 1683, 1684, 1685, 1686, 1687, 1688, 1689, 1690, 1691, 1694: +- `"ConfigSingleton" has no attribute "user_flag"` etc. — singleton attributes set dynamically. +- `Cannot determine type of "_config"`. +- `Argument 1 to "copytree" has incompatible type "Traversable"; expected "str | PathLike[str]"`. + +**Decision needed**: refactor `ConfigSingleton` with explicit `__slots__` / class-level annotations, or sprinkle `cast()`/`# type: ignore`. + +### K. `logbook.py:60` — `Element` type-var bound + +`Value of type variable "_Tag" of "Element" cannot be "None"` — passing `None` where `xml.etree.ElementTree` expects a string tag. Likely a real bug if `name` is ever None at runtime. + +### L. `gui/__init__.py:85` — `sys.excepthook` signature + +```python +sys.excepthook = my_hook +``` +where `my_hook`'s third parameter is `TracebackType` (non-Optional) but `sys.excepthook` expects `Optional[TracebackType]`. Need to add a None guard inside the hook. + +### M. `pydantic_editor.py` — list/tuple variance and BaseModel typing + +Lines 150, 154, 171, 189, 254 in `src/badger/gui/components/pydantic_editor.py`: +- `Incompatible types in assignment (expression has type "list[Any]", variable has type "tuple[Any, ...]")` — type narrowing breaks because variable is reused. +- `Argument "main" to "BadgerResolvedType" has incompatible type "Any | type[Any] | | None"; expected "type[Any] | None"` — pydantic generic resolution returns more types than annotated. +- `Incompatible types in assignment (expression has type "BadgerResolvedType | None", variable has type "BadgerResolvedType")`. + +These need careful review of the pydantic-resolution machinery. + +### N. `pydantic_editor.py:1110` — redundant cast + +Already addressed in code; the `cast` was redundant but mypy says so as `[redundant-cast]`. Just delete the call. + +### O. Untyped pydantic decorators in `routine.py` + +Lines 52, 143: `@validate_model`, `@validate_data` make decorated methods untyped. Cascade is the same as §G. Wait for upstream pydantic typing fix or add `# type: ignore[misc]`. + +### P. Stale `# type: ignore` comments (mechanical) + +Files with `[unused-ignore]` errors — just delete the comments: +- `src/badger/gui/components/pydantic_editor.py`: lines 668, 674, 931 +- `src/badger/gui/components/plot_event_handlers.py`: lines 57, 76, 81, 318, 319 +- `src/badger/gui/components/pf_viewer/pf_widget.py`: lines 76, 77 +- `src/badger/gui/components/bo_visualizer/bo_widget.py`: lines 51, 52, 120, 129, 241 +- `src/badger/gui/windows/settings_dialog.py`: line 181 + +### Q. `editable_table.py:263` — unbound `ParamSpec` + +`ParamSpec "P" is unbound`. The class uses `ParamSpec` outside a generic context. Needs refactor to either remove the `ParamSpec` or make the surrounding class generic. + +### R. `var_table.py:263,333` — duplicate attribute assignment (`[no-redef]`) + +Class attributes `selected` and `addtl_vars` assigned in two places. Needs single-source declaration in `__init__`. + +### S. Test file logic mismatches + +- `src/badger/tests/test_environment.py:126` — dict literal type mismatch (`Dict entry 0 has incompatible type "str": "list[int]"; expected "str": "float"`). The test is feeding wrong-shape data. +- `src/badger/tests/test_core.py:64` — `Incompatible types in assignment (expression has type "dict[Any, Any]", variable has type "None")` — variable initialized as None then assigned dict; should be `dict | None`. +- `src/badger/tests/x-test_db.py` — file prefixed `x-`, not collected by pytest. Likely dead code; could be deleted. +- `src/badger/tests/utils.py:25,53,94,124,162,164` — `Argument after ** must be a mapping, not "Collection[str]"` — repeatedly unpacks a Collection as kwargs. Actual bug if executed; the file may be unused or only partially exercised. + +--- + +## §3. Suggested follow-up plan + +1. **Mechanical PR for tests** (§1a) — adds ~200 `-> None` annotations. Low risk, high signal; gates future strictness on test code. +2. **Mechanical PR for GUI windows/dialogs** (§1b) — ~120 annotations. +3. **Mechanical PR for GUI components leaf files** (the smaller files in §1c). +4. **Logic PR: callbacks & Optional core types** (§A, §B, §C, §D) — change `self.routine` etc. to proper `Optional` and add guards. +5. **Logic PR: multiprocessing types** (§H) — `multiprocessing.synchronize.Event` etc. across `core_subprocess.py`, `log.py`, `routine_runner.py`. +6. **Logic PR: settings.py + ConfigSingleton** (§J) — restructure or cast. +7. **Logic PR: routine_page.py, run_monitor.py, var_table.py large files** (§B, §I, §R) — biggest, most fragile; do last. + +Each can be a separate PR; none of them block each other except where noted. + +--- + +## §4. Counts at end of this PR + +- **Errors before this PR**: 1506 +- **Errors after this PR**: 1473 (down 33) +- **Errors after PyQt5-stubs install**: jumped to 1621 first (Qt now visible to mypy), then dropped to 1473 after fixing core/library + stubs benefit. Net effect positive vs. baseline despite the stub-revealed new errors. +- **Files fully clean now**: `stats.py`, `errors.py`, `extension.py`, `formula.py`, `interface.py`, `logger/util.py`, `logger/observer.py`, `actions/__init__.py`, `actions/doctor.py`, `actions/run.py`, `actions/routine.py`, `actions/intf.py`, `actions/generator.py`, `actions/env.py`, `actions/config.py`, `actions/install.py`, `actions/uninstall.py`, `__main__.py`, plus partial in `core.py`, `factory.py`, `utils.py`. diff --git a/pyproject.toml b/pyproject.toml index 0b1ad630..b6461e43 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -40,7 +40,9 @@ dev = [ "pytest", "pytest-cov", "pytest-qt", - "pytest-mock" + "pytest-mock", + "PyQt5-stubs", + "mypy", ] [project.urls] diff --git a/src/badger/__main__.py b/src/badger/__main__.py index 614fcfb4..449b9e53 100644 --- a/src/badger/__main__.py +++ b/src/badger/__main__.py @@ -15,7 +15,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 32a5be1d..fe7e88de 100644 --- a/src/badger/actions/__init__.py +++ b/src/badger/actions/__init__.py @@ -1,11 +1,12 @@ import os from importlib import metadata +from typing import Any from badger.actions.doctor import check_n_config_paths from badger.utils import yprint from badger.settings import init_settings, get_user_config_folder -def show_info(args): +def show_info(args: Any) -> None: config_path = None if args.config_filepath: diff --git a/src/badger/actions/config.py b/src/badger/actions/config.py index ca3afb38..0c92e899 100644 --- a/src/badger/actions/config.py +++ b/src/badger/actions/config.py @@ -1,13 +1,14 @@ from badger.settings import init_settings import os import logging +from typing import Any from badger.utils import yprint, convert_str_to_value logger = logging.getLogger(__name__) -def config_settings(args): +def config_settings(args: Any) -> Any: config_singleton = init_settings() key = args.key @@ -30,7 +31,7 @@ def config_settings(args): logger.error(f"{key} is not a valid Badger config key!") -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) @@ -96,7 +97,7 @@ def _config_path_var(var_name): return False -def _config_core_var(var_name): +def _config_core_var(var_name: str) -> None: config_singleton = init_settings() display_name = config_singleton.read_display_name(var_name) diff --git a/src/badger/actions/doctor.py b/src/badger/actions/doctor.py index 2c5c79b1..2ecb440a 100644 --- a/src/badger/actions/doctor.py +++ b/src/badger/actions/doctor.py @@ -1,8 +1,10 @@ +from typing import Any + from badger.settings import init_settings, mock_settings from badger.actions.config import _config_path_var -def self_check(args): +def self_check(args: Any) -> None: config = init_settings() # Reset Badger @@ -31,7 +33,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: diff --git a/src/badger/actions/env.py b/src/badger/actions/env.py index 4d8d7f4c..6aa7918b 100644 --- a/src/badger/actions/env.py +++ b/src/badger/actions/env.py @@ -1,11 +1,12 @@ import logging +from typing import Any from badger.utils import range_to_str, yprint logger = logging.getLogger(__name__) -def show_env(args): +def show_env(args: Any) -> None: try: from badger.factory import list_env, get_env except Exception as e: @@ -22,7 +23,7 @@ def show_env(args): logger.error(e) try: # The exception could carry the configs information - configs = e.configs + configs = e.configs # type: ignore[attr-defined] except: return diff --git a/src/badger/actions/generator.py b/src/badger/actions/generator.py index 4d42877d..d2c32927 100644 --- a/src/badger/actions/generator.py +++ b/src/badger/actions/generator.py @@ -1,11 +1,12 @@ import logging +from typing import Any from badger.utils import yprint logger = logging.getLogger(__name__) -def show_generator(args): +def show_generator(args: Any) -> None: try: from badger.factory import list_generators, get_generator except Exception as e: @@ -23,7 +24,7 @@ def show_generator(args): logger.error(e) try: # The exception could carry the configs information - configs = e.configs + configs = e.configs # type: ignore[attr-defined] yprint(configs) except: pass diff --git a/src/badger/actions/install.py b/src/badger/actions/install.py index 233b46e6..6aedf2ce 100644 --- a/src/badger/actions/install.py +++ b/src/badger/actions/install.py @@ -4,6 +4,7 @@ import shutil import os from os.path import exists +from typing import Any import yaml from tqdm.auto import tqdm @@ -13,7 +14,7 @@ logger = logging.getLogger(__name__) -def plugin_install(args): +def plugin_install(args: Any) -> 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 72704b35..e82e2bf1 100644 --- a/src/badger/actions/intf.py +++ b/src/badger/actions/intf.py @@ -1,11 +1,12 @@ import logging +from typing import Any from badger.utils import yprint logger = logging.getLogger(__name__) -def show_intf(args): +def show_intf(args: Any) -> None: try: from badger.factory import list_intf, get_intf except Exception as e: @@ -23,7 +24,7 @@ def show_intf(args): logger.error(e) try: # The exception could carry the configs information - configs = e.configs + configs = e.configs # type: ignore[attr-defined] yprint(configs) except: pass diff --git a/src/badger/actions/routine.py b/src/badger/actions/routine.py index b1476551..e2b67468 100644 --- a/src/badger/actions/routine.py +++ b/src/badger/actions/routine.py @@ -1,4 +1,5 @@ import logging +from typing import Any import pandas as pd import yaml @@ -8,7 +9,7 @@ logger = logging.getLogger(__name__) -def show_routine(args): +def show_routine(args: Any) -> 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 c80bdc61..573cdeec 100644 --- a/src/badger/actions/run.py +++ b/src/badger/actions/run.py @@ -3,6 +3,7 @@ import sys import time import signal +from typing import Any from pandas import DataFrame @@ -16,8 +17,13 @@ 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: @@ -25,13 +31,13 @@ def run_n_archive( return # Store system states and other stuff - storage = { + storage: dict[str, Any] = { "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 @@ -41,10 +47,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: " @@ -59,7 +65,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 @@ -86,7 +92,7 @@ def after_evaluate(data: DataFrame): # 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 +120,7 @@ def states_ready(states): pass -def run_routine(args): +def run_routine(args: Any) -> 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 27376bcd..cce43d5c 100644 --- a/src/badger/actions/uninstall.py +++ b/src/badger/actions/uninstall.py @@ -1,11 +1,12 @@ import shutil from os.path import exists import logging +from typing import Any logger = logging.getLogger(__name__) -def plugin_remove(args): +def plugin_remove(args: Any) -> None: print( "This command is currently disabled.\n" "Please refer to the Badger documentation for plugin management.\n\n" diff --git a/src/badger/core.py b/src/badger/core.py index dbb6bd40..0f3b2ccc 100644 --- a/src/badger/core.py +++ b/src/badger/core.py @@ -1,5 +1,5 @@ import time -from typing import Callable +from typing import Any, Callable from pandas import concat, DataFrame @@ -12,7 +12,7 @@ from xopt.vocs import select_best -def check_run_status(active_callback): +def check_run_status(active_callback: Callable[[], int]) -> None: while True: status = active_callback() if status == 2: @@ -24,7 +24,19 @@ def check_run_status(active_callback): break -def convert_to_solution(result: DataFrame, routine: Routine): +def convert_to_solution( + result: DataFrame, routine: Routine +) -> tuple[ + list[Any], + list[Any], + list[Any], + list[Any], + bool, + list[str], + list[str], + list[str], + list[str], +]: vocs = routine.vocs try: best_idx, _, _ = select_best(vocs, routine.sorted_data, n=1) @@ -63,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, + 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: """ diff --git a/src/badger/errors.py b/src/badger/errors.py index b69d6cec..492d7b7e 100644 --- a/src/badger/errors.py +++ b/src/badger/errors.py @@ -4,7 +4,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() @@ -12,7 +12,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. """ @@ -27,7 +27,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. """ @@ -73,7 +73,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) @@ -98,7 +98,7 @@ class BadgerLogbookError(Exception): class BadgerLoadConfigError(Exception): - def __init__(self, message="Config Error!"): + def __init__(self, message: str = "Config Error!") -> None: super().__init__(message) @@ -107,7 +107,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) diff --git a/src/badger/extension.py b/src/badger/extension.py index 1286c4f6..de0a0ed2 100644 --- a/src/badger/extension.py +++ b/src/badger/extension.py @@ -1,5 +1,5 @@ from abc import ABC, abstractmethod -from typing import List +from typing import Any, Callable class Extension(ABC): @@ -9,21 +9,21 @@ def name(self) -> str: pass @abstractmethod - def __init__(self): + def __init__(self) -> None: pass # List all available generators @abstractmethod - def list_generator(self) -> List[str]: + def list_generator(self) -> list[str]: pass # Get config of an generator @abstractmethod - def get_generator_config(self, name: str): + def get_generator_config(self, name: str) -> Any: pass # Run an optimization with an array-style evaluate function # and a configs dict @abstractmethod - def optimize(self, evaluate, configs: dict): + def optimize(self, evaluate: Callable[..., Any], configs: dict[str, Any]) -> Any: pass diff --git a/src/badger/factory.py b/src/badger/factory.py index 4f6ba016..f09bc951 100644 --- a/src/badger/factory.py +++ b/src/badger/factory.py @@ -64,7 +64,7 @@ class BadgerPluginConfig(TypedDict): sys.path.append(BADGER_PLUGIN_ROOT) -def scan_plugins(root: str): +def scan_plugins(root: str) -> dict[str, Any]: factory: dict[str, Any] = {} # Do not scan local generators if option disabled @@ -128,7 +128,7 @@ def load_plugin( _e = BadgerInvalidPluginError( f"{ptype} {pname} is not available due to missing dependencies: {e}" ) - _e.configs = configs # attach information to the exception + _e.configs = configs # type: ignore[attr-defined] # attach information to the exception raise _e if ptype == "generator": @@ -185,7 +185,7 @@ def load_plugin( return plugin -def load_badger_docs(name: str, ptype: str = None) -> str: +def load_badger_docs(name: str, ptype: str | None = None) -> str: """ Load general Badger documentation from Badger/documentation/docs/guides. @@ -231,7 +231,7 @@ def load_badger_docs(name: str, ptype: str = 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: @@ -283,7 +283,7 @@ def load_plugin_docs(pname: str, ptype: str) -> str: if ptype == "environment": docstring = module.Environment.__doc__ - return _format_docs_str(readme, docstring, ptype) + return _format_docs_str(readme, docstring or "", ptype) except: raise BadgerInvalidDocsError( f"Error loading docs for {ptype} {pname}: docs not found" @@ -312,7 +312,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. @@ -345,7 +345,7 @@ def _format_md_docs(text: str): def _md_images_to_html( text: str, - base_prefix: str = None, + base_prefix: str | Path | None = None, width: int = 575, ) -> str: """ @@ -357,15 +357,15 @@ 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[arg-type,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: str) -> tuple[Any, Any]: try: plug = BADGER_FACTORY[ptype][name] if plug is None: # lazy loading @@ -383,25 +383,25 @@ 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") -def list_generators(): +def list_generators() -> list[str]: try: from xopt.generators import try_load_all_generators @@ -417,11 +417,11 @@ def list_generators(): 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 92850db2..9137e2fe 100644 --- a/src/badger/formula.py +++ b/src/badger/formula.py @@ -1,18 +1,20 @@ +from typing import Any, Iterable + import numpy as np import re import ast import difflib -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)} @@ -20,16 +22,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, @@ -49,7 +51,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)}") @@ -78,7 +80,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/interface.py b/src/badger/interface.py index 22e41e5c..3dd97287 100644 --- a/src/badger/interface.py +++ b/src/badger/interface.py @@ -1,14 +1,14 @@ import pickle from abc import ABC, abstractmethod -from typing import Any, ClassVar, Dict, List, TypedDict +from typing import Any, Callable, ClassVar, TypedDict from pydantic import BaseModel 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.keys(): channel_inputs = kwargs["channel_inputs"] @@ -41,8 +41,8 @@ def func_log(*args, **kwargs): class InterfaceInfo(TypedDict): - interface: Dict[str, str] - vars: List[Dict[str, str]] + interface: dict[str, str] + vars: list[dict[str, str]] class Interface(BaseModel, ABC): @@ -51,12 +51,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 @@ -68,7 +68,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 @@ -78,28 +78,28 @@ 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]: + def get_values(self, channel_names: list[str]) -> dict[str, Any]: pass # Environment should only call this method to set channels @abstractmethod - def set_values(self, channel_inputs: Dict[str, Any]): + def set_values(self, channel_inputs: dict[str, Any]) -> Any: pass - 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 """ pass - 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) -> Any: 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/logger/observer.py b/src/badger/logger/observer.py index 27b9bbed..d0f95d4a 100644 --- a/src/badger/logger/observer.py +++ b/src/badger/logger/observer.py @@ -3,26 +3,28 @@ """ from datetime import datetime +from typing import Any + from badger.logger.event import Events class Observer: - def update(self, event, solution): + def update(self, event: Any, solution: Any) -> None: raise NotImplementedError class _Tracker(object): - 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: Any, solution: Any) -> None: if event == Events.OPTIMIZATION_STEP: self._iterations += 1 - def _time_metrics(self): + def _time_metrics(self) -> tuple[str, float, float]: now = datetime.now() 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 c7cbdae6..c126a705 100644 --- a/src/badger/logger/util.py +++ b/src/badger/logger/util.py @@ -14,55 +14,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/stats.py b/src/badger/stats.py index 310a5f22..71a5d8b7 100644 --- a/src/badger/stats.py +++ b/src/badger/stats.py @@ -1,45 +1,48 @@ +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/utils.py b/src/badger/utils.py index cb58342c..6b1335bc 100644 --- a/src/badger/utils.py +++ b/src/badger/utils.py @@ -6,9 +6,10 @@ import pathlib from datetime import datetime from types import TracebackType -from typing import Iterable, Optional +from typing import Any, Iterable, Optional import yaml +from pandas import DataFrame from badger.errors import BadgerLoadConfigError from PyQt5.QtWidgets import QWidget, QLayout @@ -19,13 +20,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( @@ -38,7 +41,7 @@ def __exit__( exc_type: Optional[type[BaseException]], exc_value: Optional[BaseException], exc_traceback: Optional[TracebackType], - ): + ) -> None: for widget in self.widgets: if not widget.signalsBlocked(): logger.warning( @@ -50,34 +53,34 @@ 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) -> Any: return super(Dumper, self).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) -def yprint(content): +def yprint(content: Any) -> None: print(get_yaml_string(content), end="") -def norm(x, lb, ub): +def norm(x: Any, lb: Any, ub: Any) -> Any: return (x - lb) / (ub - lb) -def denorm(x, lb, ub): +def denorm(x: Any, lb: Any, ub: Any) -> Any: return (1 - x) * lb + x * ub -def config_list_to_dict(config_list): +def config_list_to_dict(config_list: list[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 @@ -85,7 +88,7 @@ def config_list_to_dict(config_list): return book -def load_config(fname): +def load_config(fname: str | None) -> Any: configs = None if fname is None: @@ -115,7 +118,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: @@ -128,9 +133,9 @@ def merge_params(default_params, params): return merged_params -def range_to_str(vranges): +def range_to_str(vranges: list[dict[str, Any]]) -> 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] @@ -140,7 +145,7 @@ def range_to_str(vranges): return vranges_str -def ts_to_str(ts, format="lcls-log"): +def ts_to_str(ts: datetime, format: str = "lcls-log") -> str: if format == "lcls-log": return ts.strftime("%d-%b-%Y %H:%M:%S") elif format == "lcls-log-full": @@ -151,7 +156,7 @@ def ts_to_str(ts, format="lcls-log"): return ts.isoformat() -def str_to_ts(timestr, format="lcls-log"): +def str_to_ts(timestr: str, format: str = "lcls-log") -> datetime: if format == "lcls-log": return datetime.strptime(timestr, "%d-%b-%Y %H:%M:%S") elif format == "lcls-log-full": @@ -162,20 +167,20 @@ def str_to_ts(timestr, format="lcls-log"): return datetime.fromisoformat(timestr) -def ts_float_to_str(ts_float, format="lcls-log"): +def ts_float_to_str(ts_float: float, format: str = "lcls-log") -> str: ts = datetime.fromtimestamp(ts_float) return ts_to_str(ts, format) -def curr_ts(): +def curr_ts() -> datetime: return datetime.now() -def curr_ts_to_str(format="lcls-log"): +def curr_ts_to_str(format: str = "lcls-log") -> str: return ts_to_str(datetime.now(), format) -def create_archive_run_filename(routine, format: str = "lcls-fname") -> str: +def create_archive_run_filename(routine: Any, format: str = "lcls-fname") -> str: data = routine.sorted_data env_name = routine.environment.name data_dict = data.to_dict("list") @@ -185,7 +190,7 @@ def create_archive_run_filename(routine, format: str = "lcls-fname") -> str: return fname -def get_header(routine): +def get_header(routine: Any) -> list[str]: try: obj_names = routine.vocs.objective_names except Exception: @@ -206,8 +211,8 @@ def get_header(routine): return obj_names + con_names + var_names + sta_names -def run_names_to_dict(run_names): - runs = {} +def run_names_to_dict(run_names: list[str]) -> dict[str, Any]: + runs: dict[str, Any] = {} for name in run_names: name = os.path.basename( name @@ -239,7 +244,7 @@ def run_names_to_dict(run_names): return runs -def convert_str_to_value(str): +def convert_str_to_value(str: str) -> Any: try: return int(str) except ValueError: @@ -258,7 +263,7 @@ def convert_str_to_value(str): return str -def parse_rule(rule): +def parse_rule(rule: str | dict[str, Any]) -> dict[str, str]: if type(rule) is str: return { "direction": rule, @@ -287,7 +292,7 @@ def parse_rule(rule): } -def get_value_or_none(book, key): +def get_value_or_none(book: dict[Any, Any], key: Any) -> Any: try: value = book[key] except KeyError: @@ -296,7 +301,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: Any, data: DataFrame) -> None: """dump data to file""" if dump_file is not None: output = state_to_dict(generator, data) @@ -305,7 +310,9 @@ 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): +def state_to_dict( + generator: Any, data: DataFrame, include_data: bool = True +) -> dict[str, Any]: # dump data to dict with config metadata output = { "generator": { @@ -321,7 +328,7 @@ def state_to_dict(generator, data, include_data=True): # https://stackoverflow.com/a/18472142 -def strtobool(val): +def strtobool(val: Any) -> Any: """Convert a string representation of truth to true (1) or false (0). True values are 'y', 'yes', 't', 'true', 'on', and '1'; false values are 'n', 'no', 'f', 'false', 'off', and '0'. Raises ValueError if @@ -359,11 +366,12 @@ 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") From e2fd13eabf5e517afc86949f5dab5347db604e11 Mon Sep 17 00:00:00 2001 From: Mitchell Victoriano Date: Wed, 10 Jun 2026 17:27:02 -0700 Subject: [PATCH 06/15] Re-enable mypy hook --- .pre-commit-config.yaml | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 697dc9b0..d1606987 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -22,8 +22,8 @@ repos: - id: ruff-check args: [--fix] - id: ruff-format - # - repo: https://github.com/pre-commit/mirrors-mypy - # rev: v2.1.0 - # hooks: - # - id: mypy - # args: [--strict, --ignore-missing-imports] + - repo: https://github.com/pre-commit/mirrors-mypy + rev: v2.1.0 + hooks: + - id: mypy + args: [--strict, --ignore-missing-imports] From 2e62733da41402f38f0e7c4a1a715202ee2f328b Mon Sep 17 00:00:00 2001 From: Mitchell Victoriano Date: Fri, 12 Jun 2026 01:31:16 -0700 Subject: [PATCH 07/15] Working way through repository adding types --- .pre-commit-config.yaml | 4 + pyproject.toml | 11 +- src/badger/__main__.py | 8 +- src/badger/actions/run.py | 18 ++- src/badger/core.py | 29 ++--- src/badger/errors.py | 5 +- src/badger/gui/__init__.py | 36 +++--- src/badger/gui/components/action_bar.py | 13 +- .../components/bo_visualizer/plotting_area.py | 4 +- src/badger/gui/components/eliding_label.py | 51 ++++---- src/badger/gui/components/env_cbox.py | 73 ++++++----- src/badger/gui/components/navigators.py | 22 ++-- .../gui/components/pf_viewer/pf_widget.py | 119 +++++++++--------- src/badger/gui/components/pf_viewer/types.py | 11 +- src/badger/gui/components/robust_spinbox.py | 46 ++----- src/badger/gui/components/routine_page.py | 106 ++++++++-------- src/badger/gui/components/status_bar.py | 18 +-- src/badger/gui/components/var_table.py | 103 ++++++++------- src/badger/gui/pages/home_page.py | 2 +- src/badger/gui/utils.py | 58 +++++---- .../gui/windows/expandable_message_box.py | 25 ++-- src/badger/gui/windows/settings_dialog.py | 42 ++++--- src/badger/log.py | 54 ++++---- src/badger/logger/__init__.py | 60 +++++---- src/badger/logger/event.py | 18 ++- src/badger/logger/observer.py | 7 +- src/badger/utils.py | 65 ++++++---- 27 files changed, 530 insertions(+), 478 deletions(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 697dc9b0..6a42133f 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -27,3 +27,7 @@ repos: # hooks: # - id: mypy # args: [--strict, --ignore-missing-imports] + # additional_dependencies: + # - pydantic + # - pandas-stubs + # - types-requests diff --git a/pyproject.toml b/pyproject.toml index 1d5c3cf8..540e70d4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -40,6 +40,8 @@ dev = [ "pytest-qt", "pytest-mock", "pyqt5-stubs", + "pandas-stubs", + "types-requests", "mypy", ] @@ -51,12 +53,13 @@ Documentation = "https://xopt-org.github.io/Badger/" [project.scripts] badger = "badger.__main__:main" + [project.license] -file = "LICENSE" +file = "LICENSE" # Deprecated as per PEP 639 -[options] -zip_safe = false -include_package_data = true +[tool.setuptools] +zip-safe = false +include-package-data = true [tool.setuptools.packages.find] where = ["src"] diff --git a/src/badger/__main__.py b/src/badger/__main__.py index 449b9e53..9c3ae31e 100644 --- a/src/badger/__main__.py +++ b/src/badger/__main__.py @@ -2,14 +2,14 @@ import logging from badger.actions import show_info +from badger.actions.config import config_settings from badger.actions.doctor import self_check -from badger.actions.routine import show_routine -from badger.actions.generator import show_generator from badger.actions.env import show_env +from badger.actions.generator import show_generator from badger.actions.install import plugin_install -from badger.actions.uninstall import plugin_remove from badger.actions.intf import show_intf -from badger.actions.config import config_settings +from badger.actions.routine import show_routine +from badger.actions.uninstall import plugin_remove from badger.log import setup_logging logger = logging.getLogger("badger") diff --git a/src/badger/actions/run.py b/src/badger/actions/run.py index 573cdeec..98ec743b 100644 --- a/src/badger/actions/run.py +++ b/src/badger/actions/run.py @@ -1,21 +1,27 @@ import logging import os +import signal import sys import time -import signal -from typing import Any +from typing import Any, TypedDict from pandas import DataFrame -from badger.utils import curr_ts from badger.core import run_routine as run +from badger.errors import BadgerRunTerminated from badger.routine import Routine from badger.settings import init_settings -from badger.errors import BadgerRunTerminated +from badger.utils import curr_ts logger = logging.getLogger(__name__) +class Storage(TypedDict): + states: Any + ts_last_dump: float | None + paused: bool + + def run_n_archive( routine: Routine, yes: bool = False, @@ -31,7 +37,7 @@ def run_n_archive( return # Store system states and other stuff - storage: dict[str, Any] = { + storage: Storage = { "states": None, "ts_last_dump": None, "paused": False, @@ -92,7 +98,7 @@ def after_evaluate(data: DataFrame) -> None: # take a break to let the outside signal to change the status time.sleep(sleep) - def states_ready(states: Any) -> None: + def states_ready(states) -> None: storage["states"] = states try: diff --git a/src/badger/core.py b/src/badger/core.py index 0f3b2ccc..b7aa3391 100644 --- a/src/badger/core.py +++ b/src/badger/core.py @@ -1,16 +1,15 @@ import time from typing import Any, Callable -from pandas import concat, DataFrame +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 -from xopt.vocs import select_best - def check_run_status(active_callback: Callable[[], int]) -> None: while True: @@ -24,25 +23,16 @@ def check_run_status(active_callback: Callable[[], int]) -> None: break -def convert_to_solution( - result: DataFrame, routine: Routine -) -> tuple[ - list[Any], - list[Any], - list[Any], - list[Any], - bool, - list[str], - list[str], - list[str], - list[str], -]: +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: + is_optimal = False + elif best_idx != len(routine.data) - 1: is_optimal = False else: is_optimal = True @@ -58,7 +48,7 @@ def convert_to_solution( cons = list(result[vocs.constraint_names].to_numpy()[0]) stas = list(result[vocs.observable_names].to_numpy()[0]) - solution = ( + solution = Solution( vars, objs, cons, @@ -132,6 +122,7 @@ def run_routine( routine.vocs.constraint_names, routine.vocs.observable_names, ) + solution_meta = Solution(*solution_meta) opt_logger.update(Events.OPTIMIZATION_START, solution_meta) # evaluate initial points: diff --git a/src/badger/errors.py b/src/badger/errors.py index 492d7b7e..39e443ec 100644 --- a/src/badger/errors.py +++ b/src/badger/errors.py @@ -1,6 +1,7 @@ -from PyQt5.QtWidgets import QMessageBox -import traceback import sys +import traceback + +from PyQt5.QtWidgets import QMessageBox class BadgerError(Exception): diff --git a/src/badger/gui/__init__.py b/src/badger/gui/__init__.py index 6257e084..55791d94 100644 --- a/src/badger/gui/__init__.py +++ b/src/badger/gui/__init__.py @@ -1,20 +1,20 @@ -from importlib import resources +import logging import signal import sys import time -from PyQt5.QtWidgets import QApplication -from PyQt5.QtGui import QFont, QIcon -from PyQt5 import QtCore -from qdarkstyle import load_stylesheet, LightPalette, DarkPalette +import traceback +from importlib import resources +from types import TracebackType +from typing import Any, Type -from badger.settings import init_settings -from badger.gui.windows.main_window import BadgerMainWindow +from PyQt5 import QtCore +from PyQt5.QtGui import QFont, QIcon +from PyQt5.QtWidgets import QApplication +from qdarkstyle import DarkPalette, LightPalette, load_stylesheet -import traceback from badger.errors import BadgerError -from types import TracebackType -from typing import Type, NoReturn -import logging +from badger.gui.windows.main_window import BadgerMainWindow +from badger.settings import init_settings logger = logging.getLogger(__name__) @@ -32,10 +32,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() @@ -46,7 +46,7 @@ def on_exit(*args): QApplication.quit() -def on_timeout(): +def on_timeout() -> None: if TIMER["time"] is None: return @@ -57,8 +57,8 @@ def on_timeout(): def error_handler( - etype: Type[BaseException], value: BaseException, tb: TracebackType -) -> NoReturn: + etype: Type[BaseException], value: BaseException, tb: TracebackType | None +) -> None: """ Custom exception handler that formats uncaught exceptions and raises a BadgerError. @@ -68,7 +68,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 @@ -81,7 +81,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 05525f46..9ea3692d 100644 --- a/src/badger/gui/components/action_bar.py +++ b/src/badger/gui/components/action_bar.py @@ -1,8 +1,9 @@ -from PyQt5.QtWidgets import QWidget, QHBoxLayout -from PyQt5.QtWidgets import QToolButton, QMenu, QAction -from PyQt5.QtGui import QIcon, QFont -from PyQt5.QtCore import pyqtSignal, QSize from importlib import resources + +from PyQt5.QtCore import QSize, pyqtSignal +from PyQt5.QtGui import QFont, QIcon +from PyQt5.QtWidgets import QAction, QHBoxLayout, QMenu, QToolButton, QWidget + from badger.gui.utils import create_button from badger.gui.windows.docs_window import BadgerDocsWindow @@ -101,12 +102,12 @@ 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.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: diff --git a/src/badger/gui/components/bo_visualizer/plotting_area.py b/src/badger/gui/components/bo_visualizer/plotting_area.py index 3e1bfe13..01fa8cdf 100644 --- a/src/badger/gui/components/bo_visualizer/plotting_area.py +++ b/src/badger/gui/components/bo_visualizer/plotting_area.py @@ -91,8 +91,8 @@ 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"] diff --git a/src/badger/gui/components/eliding_label.py b/src/badger/gui/components/eliding_label.py index af7acd01..f29e6dd0 100644 --- a/src/badger/gui/components/eliding_label.py +++ b/src/badger/gui/components/eliding_label.py @@ -1,8 +1,12 @@ # https://stackoverflow.com/a/67628976/4263605 -from PyQt5 import QtCore, QtWidgets, QtGui +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: @@ -36,19 +40,22 @@ 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.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.Expanding, QSizePolicy.Preferred) self.setText(text) - def setText(self, text): + def setText(self, text: str) -> None: self._contents = text # This line set for testing. Its value is the return value of @@ -59,20 +66,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: 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: @@ -87,13 +94,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() @@ -102,20 +107,20 @@ 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 = "", parent: QWidget | None = None) -> None: super(SimpleElidedLabel, self).__init__(parent) self._text = text - def setText(self, text): + def setText(self, text: str) -> None: self._text = text super(SimpleElidedLabel, self).setText(self.elidedText()) - def resizeEvent(self, event): + def resizeEvent(self, event: QResizeEvent) -> None: super(SimpleElidedLabel, self).setText(self.elidedText()) super(SimpleElidedLabel, self).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.ElideRight, self.width()) return elided diff --git a/src/badger/gui/components/env_cbox.py b/src/badger/gui/components/env_cbox.py index de11c253..65833dca 100644 --- a/src/badger/gui/components/env_cbox.py +++ b/src/badger/gui/components/env_cbox.py @@ -1,41 +1,38 @@ +import logging from importlib import resources from typing import Any -from PyQt5.QtWidgets import ( - QVBoxLayout, - QHBoxLayout, - QPushButton, - QWidget, - QLineEdit, -) +from gest_api.vocs import ContinuousVariable +from pydantic_core import ValidationError +from PyQt5.QtCore import QPropertyAnimation, QRegExp, pyqtSignal +from PyQt5.QtGui import QFont, QIcon from PyQt5.QtWidgets import ( QCheckBox, - QStyledItemDelegate, + QHBoxLayout, QLabel, + QLineEdit, + QPushButton, QSizePolicy, + QStyledItemDelegate, + QVBoxLayout, + QWidget, ) -from PyQt5.QtGui import QIcon, QFont -from PyQt5.QtCore import QRegExp, QPropertyAnimation, pyqtSignal -from pydantic_core import ValidationError +from xopt.vocs import VOCS from badger.errors import BadgerRoutineError from badger.gui.components.collapsible_box import CollapsibleBox -from badger.gui.components.pydantic_editor import BadgerPydanticEditor -from badger.gui.components.var_table import VariableTable -from badger.gui.components.obj_table import ObjectiveTable from badger.gui.components.con_table import ConstraintTable -from badger.gui.components.obs_table import ObservableTable from badger.gui.components.data_table import init_data_table -from badger.settings import init_settings +from badger.gui.components.obj_table import ObjectiveTable +from badger.gui.components.obs_table import ObservableTable +from badger.gui.components.pydantic_editor import BadgerPydanticEditor +from badger.gui.components.var_table import VariableTable from badger.gui.utils import ( MouseWheelWidgetAdjustmentGuard, NoHoverFocusComboBox, ) +from badger.settings import init_settings from badger.utils import strtobool -from xopt.vocs import VOCS -from gest_api.vocs import ContinuousVariable - -import logging LABEL_WIDTH = 96 ENV_PARAMS_BTN = 1 # use button or collapsible box for env parameters @@ -146,7 +143,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" @@ -499,7 +496,7 @@ def init_ui(self): cbox_more.setContentLayout(vbox_more) - def config_logic(self): + def config_logic(self) -> None: self.dict_con = {} self.edit_var.textChanged.connect(self.filter_var) @@ -518,12 +515,12 @@ def config_logic(self): self.sta_table.data_changed.connect(lambda: self.update_vocs("sta_table")) self.var_table.data_changed.connect(lambda: self.update_vocs("var_table")) - def update_vocs(self, origin: str): + def update_vocs(self, origin: str) -> None: logger.debug(f"Emitting vocs_updated signal from env_cbox: {origin}") 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) @@ -536,18 +533,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) @@ -559,33 +556,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) -> 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) @@ -595,7 +592,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""" diff --git a/src/badger/gui/components/navigators.py b/src/badger/gui/components/navigators.py index 6709e22e..a3647a00 100644 --- a/src/badger/gui/components/navigators.py +++ b/src/badger/gui/components/navigators.py @@ -11,23 +11,25 @@ """ import os + +from PyQt5.QtCore import QDir, Qt, QTimer, QUrl +from PyQt5.QtGui import QCursor, QDesktopServices, QFont from PyQt5.QtWidgets import ( - QWidget, - QVBoxLayout, - QTreeView, - QTreeWidget, - QTreeWidgetItem, - QMenu, QAction, QApplication, - QToolTip, QFileSystemModel, + QMenu, + QToolTip, + QTreeView, + QTreeWidget, + QTreeWidgetItem, + QVBoxLayout, + QWidget, ) -from PyQt5.QtGui import QFont, QDesktopServices, QCursor -from PyQt5.QtCore import Qt, QUrl, QTimer, QDir + from badger.archive import get_base_run_filename, get_runs -from badger.utils import run_names_to_dict from badger.settings import init_settings +from badger.utils import run_names_to_dict class FileContextMenuBase: diff --git a/src/badger/gui/components/pf_viewer/pf_widget.py b/src/badger/gui/components/pf_viewer/pf_widget.py index 51a182a5..5d86254e 100644 --- a/src/badger/gui/components/pf_viewer/pf_widget.py +++ b/src/badger/gui/components/pf_viewer/pf_widget.py @@ -1,59 +1,52 @@ +import logging import time from typing import Optional +import matplotlib.pyplot as plt +import pandas as pd +from matplotlib.axes import Axes +from matplotlib.backends.backend_qt import NavigationToolbar2QT as NavigationToolbar +from matplotlib.backends.backend_qtagg import FigureCanvasQTAgg as FigureCanvas +from matplotlib.colors import Normalize +from matplotlib.figure import Figure +from matplotlib.ticker import MaxNLocator +from PyQt5.QtCore import Qt from PyQt5.QtWidgets import ( - QWidget, - QSizePolicy, - QVBoxLayout, - QHBoxLayout, - QGridLayout, - QPushButton, - QComboBox, QCheckBox, - QLabel, + QComboBox, + QGridLayout, QGroupBox, + QHBoxLayout, + QLabel, + QPushButton, + QSizePolicy, QTabWidget, + QVBoxLayout, + QWidget, ) - -from PyQt5.QtCore import Qt -from badger.gui.components.plot_event_handlers import ( - MatplotlibInteractionHandler, -) -from badger.gui.components.pf_viewer.types import PFUI, ConfigurableOptions -from badger.routine import Routine - -from badger.utils import BlockSignalsContext, create_archive_run_filename -from matplotlib.axes import Axes -from matplotlib.figure import Figure -import matplotlib.pyplot as plt -from matplotlib.colors import Normalize -from matplotlib.backends.backend_qtagg import FigureCanvasQTAgg as FigureCanvas -from matplotlib.ticker import MaxNLocator -from matplotlib.backends.backend_qt import ( - NavigationToolbar2QT as NavigationToolbar, -) -import pandas as pd from torch import Tensor - from xopt.generators.bayesian.mobo import MOBOGenerator -from badger.gui.components.pf_viewer.types import ( - PFUIWidgets, - PFUILayouts, -) - from badger.gui.components.analysis_widget import AnalysisWidget - from badger.gui.components.extension_utilities import ( HandledException, MatplotlibFigureContext, - signal_logger, - requires_update, - clear_tabs, clear_layout, + clear_tabs, + requires_update, + signal_logger, ) - -import logging +from badger.gui.components.pf_viewer.types import ( + PFUI, + ConfigurableOptions, + PFUILayouts, + PFUIWidgets, +) +from badger.gui.components.plot_event_handlers import ( + MatplotlibInteractionHandler, +) +from badger.routine import Routine +from badger.utils import BlockSignalsContext, create_archive_run_filename logger = logging.getLogger(__name__) @@ -73,8 +66,8 @@ class ParetoFrontWidget(AnalysisWidget): - generator: MOBOGenerator # type: ignore - parameters: ConfigurableOptions = DEFAULT_PARAMETERS # type: ignore + generator: MOBOGenerator # pyright: ignore[reportIncompatibleVariableOverride] + parameters: ConfigurableOptions = DEFAULT_PARAMETERS hypervolume_history: pd.DataFrame = pd.DataFrame() pf_1: Optional[Tensor] = None @@ -164,7 +157,7 @@ def update_plots( # Update the last updated time self.last_updated = time.time() - def setup_connections(self): + def setup_connections(self) -> None: self.ui["components"]["update"].clicked.connect( lambda: signal_logger("Update button clicked")( lambda: self.on_button_click() @@ -190,7 +183,7 @@ def setup_connections(self): lambda: signal_logger("Sample checkbox changed")(lambda: self.update_ui())() ) - def create_ui(self): + def create_ui(self) -> None: update_button = QPushButton("Update") variable_1_combo = QComboBox() variable_1_combo.setMinimumWidth(100) @@ -324,7 +317,7 @@ def initialize_widget(self) -> None: variable_1_combo.setCurrentIndex(self.parameters["variable_1"]) variable_2_combo.setCurrentIndex(self.parameters["variable_2"]) - def on_tab_change(self): + def on_tab_change(self) -> None: self.parameters["plot_tab"] = self.ui["components"]["plot"][ "pareto" ].currentIndex() @@ -358,7 +351,7 @@ def on_tab_change(self): # Update the plot self.update_pareto_front_plot() - def on_variable_change(self): + def on_variable_change(self) -> None: plot_tab = self.parameters["plot_tab"] if plot_tab == 0: self.parameters["variable_1"] = self.ui["components"]["variables"][ @@ -379,15 +372,15 @@ def on_variable_change(self): self.update_pareto_front_plot() - def on_button_click(self): + def on_button_click(self) -> None: self.update_extension(self.routine, True) - def update_ui(self): + def update_ui(self) -> None: self.update_plots(requires_rebuild=True) def update_pareto_front_plot( self, - ): + ) -> None: self.update_pareto_front() plot_tab_widget = self.ui["components"]["plot"]["pareto"] @@ -399,8 +392,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"] @@ -425,14 +418,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"] @@ -457,7 +450,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"]) @@ -467,7 +460,7 @@ def update_pareto_front_plot( def update_hypervolume_plot( self, - ): + ) -> None: self.update_hypervolume() plot_hypervolume = self.ui["components"]["plot"]["hypervolume"] @@ -478,14 +471,14 @@ 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): + def create_pareto_plot(self, fig: Figure, ax: Axes) -> tuple[Figure, Axes]: current_tab = self.parameters["plot_tab"] show_only_pareto_front = self.ui["components"]["options"][ "show_only_pareto_front" @@ -590,15 +583,15 @@ def create_pareto_plot(self, fig: Figure, ax: Axes): return fig, ax - def create_hypervolume_plot(self, fig: Figure, ax: Axes): + def create_hypervolume_plot(self, fig: Figure, ax: Axes) -> tuple[Figure, Axes]: data_points = self.hypervolume_history if len(data_points) == 0: raise HandledException( 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"].values.astype(float) + y = data_points["hypervolume"].values.astype(float) # Create a scatter plot ax.plot( x, @@ -616,7 +609,7 @@ def create_hypervolume_plot(self, fig: Figure, ax: Axes): return fig, ax - def update_hypervolume(self): + def update_hypervolume(self) -> None: # Get the hypervolume from the generator self.generator.update_pareto_front_history() pareto_front_history_df = self.generator.pareto_front_history @@ -626,7 +619,7 @@ def update_hypervolume(self): self.hypervolume_history = pareto_front_history_df - def update_pareto_front(self): + def update_pareto_front(self) -> None: pf_1, pf_2, pf_mask, _ = self.generator.get_pareto_front_and_hypervolume() if pf_mask is None or pf_1 is None or pf_2 is None: diff --git a/src/badger/gui/components/pf_viewer/types.py b/src/badger/gui/components/pf_viewer/types.py index 61d6d61d..c91256a6 100644 --- a/src/badger/gui/components/pf_viewer/types.py +++ b/src/badger/gui/components/pf_viewer/types.py @@ -1,12 +1,13 @@ from typing import TypedDict from PyQt5.QtWidgets import ( - QRadioButton, + QCheckBox, QComboBox, - QVBoxLayout, - QHBoxLayout, QGridLayout, + QHBoxLayout, + QPushButton, QTabWidget, + QVBoxLayout, ) @@ -26,7 +27,7 @@ class ConfigurableOptions(TypedDict): class PFOptionsUIWidgets(TypedDict): - show_only_pareto_front: QRadioButton + show_only_pareto_front: QCheckBox class PFVariablesUIWidgets(TypedDict): @@ -42,7 +43,7 @@ class PFPlotUIWidgets(TypedDict): class PFUIWidgets(TypedDict): variables: PFVariablesUIWidgets options: PFOptionsUIWidgets - update: QRadioButton + update: QPushButton plot: PFPlotUIWidgets diff --git a/src/badger/gui/components/robust_spinbox.py b/src/badger/gui/components/robust_spinbox.py index c4e4e099..0335d2fd 100644 --- a/src/badger/gui/components/robust_spinbox.py +++ b/src/badger/gui/components/robust_spinbox.py @@ -1,45 +1,23 @@ -from PyQt5.QtWidgets import QDoubleSpinBox, QAbstractSpinBox from PyQt5.QtCore import Qt +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: - decimals = 6 - - try: - lb = kwargs["lower_bound"] - del kwargs["lower_bound"] - if lb is None: - lb = -1e3 - except: - lb = -1e3 - - try: - ub = kwargs["upper_bound"] - del kwargs["upper_bound"] - if ub is None: - ub = 1e3 - except: - ub = 1e3 - - try: - default_value = kwargs["default_value"] - del kwargs["default_value"] - if default_value is None: - default_value = 0 - except: - 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_page.py b/src/badger/gui/components/routine_page.py index 0a5d1c35..184877f6 100644 --- a/src/badger/gui/components/routine_page.py +++ b/src/badger/gui/components/routine_page.py @@ -1,38 +1,60 @@ -from typing import Any -import warnings -import traceback import copy -from functools import partial +import logging import os -import yaml +import traceback +import warnings +from datetime import datetime +from functools import partial +from typing import Any import numpy as np import pandas as pd -from PyQt5.QtCore import Qt, pyqtSignal, QTimer -from PyQt5.QtWidgets import QLineEdit, QLabel, QPushButton, QFileDialog -from PyQt5.QtWidgets import QMessageBox, QWidget, QTabWidget -from PyQt5.QtWidgets import QVBoxLayout, QHBoxLayout, QScrollArea -from PyQt5.QtWidgets import QTableWidgetItem, QPlainTextEdit +import yaml from coolname import generate_slug -from xopt import VOCS -from xopt.vocs import random_inputs -from xopt.generators import ( - get_generator_defaults, - all_generator_names, - get_generator_dynamic, -) -from xopt.utils import get_local_region from gest_api.vocs import ( BaseConstraint, BaseObjective, GreaterThanConstraint, LessThanConstraint, - MinimizeObjective, MaximizeObjective, + MinimizeObjective, ) from pydantic import ValidationError +from PyQt5.QtCore import Qt, QTimer, pyqtSignal +from PyQt5.QtWidgets import ( + QFileDialog, + QHBoxLayout, + QLabel, + QLineEdit, + QMessageBox, + QPlainTextEdit, + QPushButton, + QScrollArea, + QTableWidgetItem, + QTabWidget, + QVBoxLayout, + QWidget, +) +from xopt import VOCS +from xopt.generators import ( + all_generator_names, + get_generator_defaults, + get_generator_dynamic, +) +from xopt.utils import get_local_region +from xopt.vocs import random_inputs -from badger.gui.components.generator_cbox import BadgerAlgoBox +from badger.archive import update_run +from badger.environment import instantiate_env +from badger.errors import ( + BadgerEnvInstantiationError, + BadgerEnvNotFoundError, + BadgerEnvVarError, + BadgerRoutineError, + VariableRangeError, +) +from badger.factory import get_env, list_env, list_generators +from badger.gui.components.archive_search import ArchiveSearchWidget from badger.gui.components.data_panel import BadgerDataPanel from badger.gui.components.data_table import ( get_table_content_as_dict, @@ -41,42 +63,28 @@ ) from badger.gui.components.env_cbox import BadgerEnvBox from badger.gui.components.filter_cbox import BadgerFilterBox +from badger.gui.components.generator_cbox import BadgerAlgoBox +from badger.gui.utils import filter_generator_config +from badger.gui.windows.add_random_dialog import BadgerAddRandomDialog from badger.gui.windows.docs_window import BadgerDocsWindow from badger.gui.windows.edit_script_dialog import BadgerEditScriptDialog -from badger.gui.windows.lim_vrange_dialog import BadgerLimitVariableRangeDialog from badger.gui.windows.ind_lim_vrange_dialog import ( BadgerIndividualLimitVariableRangeDialog, ) -from badger.gui.windows.review_dialog import BadgerReviewDialog -from badger.gui.windows.add_random_dialog import BadgerAddRandomDialog +from badger.gui.windows.lim_vrange_dialog import BadgerLimitVariableRangeDialog from badger.gui.windows.message_dialog import BadgerScrollableMessageBox -from badger.gui.utils import filter_generator_config -from badger.gui.components.archive_search import ArchiveSearchWidget -from badger.archive import update_run -from badger.environment import instantiate_env -from badger.errors import ( - BadgerEnvNotFoundError, - BadgerRoutineError, - BadgerEnvVarError, - BadgerEnvInstantiationError, - VariableRangeError, -) -from badger.factory import list_generators, list_env, get_env +from badger.gui.windows.review_dialog import BadgerReviewDialog from badger.routine import Routine from badger.settings import init_settings -from datetime import datetime from badger.utils import ( BlockSignalsContext, - load_config, - strtobool, get_badger_version, get_xopt_version, + load_config, + strtobool, ts_float_to_str, ) - -import logging - logger = logging.getLogger(__name__) @@ -123,7 +131,7 @@ class BadgerRoutinePage(QWidget): sig_load_template = pyqtSignal(str) # template path sig_save_template = pyqtSignal(str) # template path - def __init__(self): + def __init__(self) -> None: logger.info("Initializing BadgerRoutinePage.") super().__init__() @@ -168,7 +176,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.") config_singleton = init_settings() @@ -265,8 +273,7 @@ def init_ui(self): self.env_box = BadgerEnvBox(env_dict, None, self.envs) scroll_area = QScrollArea() scroll_area.setFrameShape(QScrollArea.NoFrame) - scroll_area.setStyleSheet( - """ + scroll_area.setStyleSheet(""" QScrollArea { border: none; /* Remove border */ margin: 0px; /* Remove margin */ @@ -275,8 +282,7 @@ def init_ui(self): QScrollArea > QWidget { margin: 0px; /* Remove margin inside */ } - """ - ) + """) scroll_content_env = QWidget() scroll_layout_env = QVBoxLayout(scroll_content_env) scroll_layout_env.setContentsMargins(0, 0, 15, 0) @@ -307,7 +313,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) @@ -340,9 +346,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. diff --git a/src/badger/gui/components/status_bar.py b/src/badger/gui/components/status_bar.py index f0f9a5da..0386de89 100644 --- a/src/badger/gui/components/status_bar.py +++ b/src/badger/gui/components/status_bar.py @@ -1,19 +1,21 @@ from importlib import resources -from PyQt5.QtWidgets import QHBoxLayout, QWidget, QPushButton + +from PyQt5.QtCore import QSize, Qt from PyQt5.QtGui import QIcon -from PyQt5.QtCore import Qt, QSize -from badger.gui.windows.settings_dialog import BadgerSettingsDialog +from PyQt5.QtWidgets import QHBoxLayout, QPushButton, QWidget + from badger.gui.components.eliding_label import SimpleElidedLabel +from badger.gui.windows.settings_dialog import BadgerSettingsDialog 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)) @@ -51,13 +53,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/var_table.py b/src/badger/gui/components/var_table.py index bb88097e..1f3f60f0 100644 --- a/src/badger/gui/components/var_table.py +++ b/src/badger/gui/components/var_table.py @@ -1,41 +1,40 @@ +import logging +import traceback from functools import partial from importlib import resources -import traceback from typing import Any, cast + +from gest_api.vocs import ContinuousVariable +from PyQt5.QtCore import QPoint, QSize, Qt, pyqtSignal +from PyQt5.QtGui import ( + QColor, + QDragEnterEvent, + QDragMoveEvent, + QDropEvent, + QGuiApplication, + QIcon, +) from PyQt5.QtWidgets import ( - QTableWidget, - QTableWidgetItem, - QHeaderView, + QAbstractItemView, QCheckBox, + QDialog, + QGridLayout, + QHBoxLayout, + QHeaderView, + QLabel, + QMenu, QMessageBox, - QAbstractItemView, QPushButton, + QTableWidget, + QTableWidgetItem, QWidget, - QHBoxLayout, - QMenu, - QGridLayout, - QLabel, - QDialog, -) -from PyQt5.QtCore import pyqtSignal, Qt, QSize, QPoint -from PyQt5.QtGui import ( - QColor, - QIcon, - QGuiApplication, - QDropEvent, - QDragMoveEvent, - QDragEnterEvent, ) -from badger.gui.components.robust_spinbox import RobustSpinBox from badger.environment import Environment, instantiate_env from badger.errors import BadgerInterfaceChannelError +from badger.gui.components.robust_spinbox import RobustSpinBox from badger.gui.windows.expandable_message_box import ExpandableMessageBox -from gest_api.vocs import ContinuousVariable - -import logging - logger = logging.getLogger(__name__) @@ -111,14 +110,14 @@ def __init__(self, *args: Any, **kwargs: Any): 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. """ logging.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) @@ -133,7 +132,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: @@ -144,7 +143,7 @@ def is_all_checked(self): return True - def header_clicked(self, idx: int): + def header_clicked(self, idx: int) -> None: if idx: return @@ -161,7 +160,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: @@ -182,7 +181,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. """ @@ -211,7 +210,7 @@ def validate_row(self, row: int): def set_bounds( self, variables: dict[str, tuple[float, float]], signal: bool = True - ): + ) -> None: for name in variables: self.bounds[name] = variables[name] @@ -222,7 +221,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 @@ -239,7 +238,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: @@ -259,21 +258,21 @@ def update_selected(self): if self.checked_only: self.show_checked_only() - def set_selected(self, variable_names: list[str]): + def set_selected(self, variable_names: list[str]) -> None: self.selected: dict[str, bool] = {} 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)) @@ -281,7 +280,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: @@ -315,7 +314,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 @@ -427,10 +426,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() @@ -456,7 +455,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 @@ -541,7 +540,7 @@ def get_bounds(self, name: str) -> tuple[Any, Any]: 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) @@ -557,15 +556,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 @@ -579,7 +578,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 @@ -593,7 +592,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 @@ -619,7 +618,7 @@ def dropEvent(self, event: QDropEvent | None): else: event.ignore() - def flags(self, index): + def flags(self, index) -> Qt.ItemFlags: if not index.isValid(): return Qt.ItemIsEnabled flags = Qt.ItemIsSelectable | Qt.ItemIsEnabled @@ -627,7 +626,7 @@ def flags(self, index): flags |= Qt.ItemIsEditable | Qt.ItemIsDropEnabled return 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. """ @@ -656,17 +655,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/pages/home_page.py b/src/badger/gui/pages/home_page.py index 31e40a35..e70e11fc 100644 --- a/src/badger/gui/pages/home_page.py +++ b/src/badger/gui/pages/home_page.py @@ -361,7 +361,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 diff --git a/src/badger/gui/utils.py b/src/badger/gui/utils.py index 0753d6a4..90cf9257 100644 --- a/src/badger/gui/utils.py +++ b/src/badger/gui/utils.py @@ -1,41 +1,46 @@ +import copy from importlib import resources from typing import Any -from PyQt5.QtWidgets import QAbstractSpinBox, QPushButton, QComboBox, QToolButton -from PyQt5.QtWidgets import QDialog, QVBoxLayout, QLabel -from PyQt5.QtCore import Qt, QObject, QEvent, QSize -from PyQt5.QtGui import QIcon -import copy - -def preventAnnoyingSpinboxScrollBehaviour(self, control: QAbstractSpinBox) -> None: - control.setFocusPolicy(Qt.StrongFocus) - control.installEventFilter(self.MouseWheelWidgetAdjustmentGuard(control)) +from PyQt5.QtCore import QEvent, QObject, QSize, Qt +from PyQt5.QtGui import QFocusEvent, QIcon, QMouseEvent +from PyQt5.QtWidgets import ( + QComboBox, + QDialog, + QLabel, + QPushButton, + QToolButton, + QVBoxLayout, + QWidget, +) 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) @@ -55,7 +60,7 @@ def create_button( return btn -def filter_generator_config(name: str, config: dict[str, Any]): +def filter_generator_config(name: str, config: dict[str, Any]) -> dict[str, Any]: filtered_config: dict[str, Any] = {} if name == "neldermead": filtered_config["adaptive"] = config["adaptive"] @@ -82,28 +87,29 @@ def filter_generator_config(name: str, config: 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: diff --git a/src/badger/gui/windows/expandable_message_box.py b/src/badger/gui/windows/expandable_message_box.py index 06e652e9..1fd8f545 100644 --- a/src/badger/gui/windows/expandable_message_box.py +++ b/src/badger/gui/windows/expandable_message_box.py @@ -1,20 +1,25 @@ +from PyQt5.QtCore import Qt +from PyQt5.QtGui import QFont, QFontDatabase, QTextOption from PyQt5.QtWidgets import ( QDialog, - QMessageBox, - QVBoxLayout, QHBoxLayout, QLabel, + QMessageBox, QPushButton, QTextEdit, + QVBoxLayout, ) -from PyQt5.QtGui import QTextOption, QFont, QFontDatabase -from PyQt5.QtCore import Qt 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) @@ -71,7 +76,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") @@ -79,13 +84,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/settings_dialog.py b/src/badger/gui/windows/settings_dialog.py index a561a908..5eb53f31 100644 --- a/src/badger/gui/windows/settings_dialog.py +++ b/src/badger/gui/windows/settings_dialog.py @@ -1,26 +1,26 @@ import logging import os +from typing import cast + +from PyQt5.QtCore import Qt # from PyQt5.QtCore import QRegExp # from PyQt5.QtGui import QRegExpValidator from PyQt5.QtWidgets import ( + QApplication, QComboBox, + QDialog, + QDialogButtonBox, QGridLayout, - QVBoxLayout, - QWidget, QLabel, QLineEdit, + QVBoxLayout, + QWidget, ) -from PyQt5.QtWidgets import ( - QDialog, - QDialogButtonBox, - QApplication, -) -from PyQt5.QtCore import Qt -from qdarkstyle import load_stylesheet, DarkPalette, LightPalette -from badger.settings import init_settings -from badger.log import get_logging_manager +from qdarkstyle import DarkPalette, LightPalette, load_stylesheet +from badger.log import get_logging_manager +from badger.settings import init_settings logger = logging.getLogger(__name__) @@ -33,7 +33,7 @@ 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) @@ -47,7 +47,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) @@ -178,20 +178,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": @@ -199,14 +203,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( @@ -251,7 +255,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/log.py b/src/badger/log.py index a33dea79..b1e8c52a 100644 --- a/src/badger/log.py +++ b/src/badger/log.py @@ -1,12 +1,12 @@ -import os +import atexit import datetime import logging -import atexit - +import os +from argparse import Namespace from logging.handlers import QueueHandler, QueueListener from multiprocessing import Queue -from badger.settings import get_user_config_folder -from badger.settings import init_settings + +from badger.settings import get_user_config_folder, init_settings logger = logging.getLogger(__name__) @@ -28,12 +28,12 @@ class LoggingManager: 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. @@ -78,12 +78,12 @@ def start_listener(self, log_filepath: str, log_level: str | int): f"Centralized logging listener started with level {logging.getLevelName(log_level)}" ) - 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. @@ -111,7 +111,7 @@ def update_log_level(self, log_level: str | int): f"Log level updated to {logging.getLevelName(log_level)} for badger namespace." ) - def update_logfile_path(self, new_logfile_path: str): + def update_logfile_path(self, new_logfile_path: str) -> None: """ Update logfile path, so logs get written to a different file. @@ -120,11 +120,14 @@ 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 # We want to keep using same formatter and loglevel formatter = file_handler.formatter @@ -145,6 +148,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 ) @@ -152,7 +158,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" @@ -161,11 +167,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. @@ -195,11 +201,11 @@ def create_log_dir(self, log_dir: str): def configure_process_logging( - log_queue: Queue = None, + log_queue: Queue[logging.LogRecord] | None = None, logger_name: str = "badger", log_level: str = "DEBUG", - process_name: str = None, -): + process_name: str | None = None, +) -> None: """ Configure logging in a process to send logs to the shared queue. This should be ran in all subprocesses b4 logging, and also in the main process where @@ -208,10 +214,10 @@ 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: @@ -243,7 +249,7 @@ def configure_process_logging( ) -def setup_logging(args): +def setup_logging(args: Namespace) -> None: """ Init Badger's multiprocess logging system """ @@ -273,7 +279,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/logger/__init__.py b/src/badger/logger/__init__.py index 651de4c5..f15e2666 100644 --- a/src/badger/logger/__init__.py +++ b/src/badger/logger/__init__.py @@ -1,34 +1,32 @@ from __future__ import print_function -import os + import json +import os +from typing import Callable +from badger.logger.event import Events, Solution from badger.logger.observer import _Tracker -from badger.logger.event import Events 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 - super(ScreenLogger, self).__init__() + 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: int | float) -> str: if isinstance(x, int): s = "{x:< {s}}".format( x=x, @@ -48,16 +46,20 @@ 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)) @@ -75,7 +77,10 @@ def _step(self, solution, colour=Colours.black): return "| " + " | ".join(map(colour, cells)) + " |" - def _header(self, solution): + def _header( + self, + solution: Solution, + ) -> str: cells = [] cells.append(self._format_key("iter")) @@ -91,14 +96,21 @@ def _header(self, solution): for sta in solution[8]: cells.append(self._format_key(sta)) - line = "| " + " | ".join(cells) + " |" + line: str = "| " + " | ".join(cells) + " |" self._header_length = len(line) return line + "\n" + ("-" * self._header_length) - def _is_new_max(self, solution): + def _is_new_max( + self, + solution: Solution, + ) -> bool: return solution[4] - def update(self, event, solution): + def update( + self, + event: Events, + solution: Solution, + ) -> None: if event == Events.OPTIMIZATION_START: line = self._header(solution) + "\n" elif event == Events.OPTIMIZATION_STEP: @@ -109,15 +121,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: @@ -126,7 +142,7 @@ def __init__(self, path, reset=True): pass super(JSONLogger, self).__init__() - def update(self, event, solution): + def update(self, event: Events, solution: Solution) -> None: if event == Events.OPTIMIZATION_STEP: data = { "x": solution[0], diff --git a/src/badger/logger/event.py b/src/badger/logger/event.py index ab590c3d..89acf2db 100644 --- a/src/badger/logger/event.py +++ b/src/badger/logger/event.py @@ -1,4 +1,20 @@ -class Events: +from enum import StrEnum +from typing import NamedTuple + + +class Solution(NamedTuple): + variables: list[float] + objectives: list[float] + constraints: list[float] + states: list[float] + 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 d0f95d4a..94ed60fb 100644 --- a/src/badger/logger/observer.py +++ b/src/badger/logger/observer.py @@ -3,13 +3,12 @@ """ from datetime import datetime -from typing import Any -from badger.logger.event import Events +from badger.logger.event import Events, Solution class Observer: - def update(self, event: Any, solution: Any) -> None: + def update(self, event: Events, solution: Solution) -> None: raise NotImplementedError @@ -20,7 +19,7 @@ def __init__(self) -> None: self._start_time: datetime | None = None self._previous_time: datetime | None = None - def _update_tracker(self, event: Any, solution: Any) -> None: + def _update_tracker(self, event: Events, solution: Solution) -> None: if event == Events.OPTIMIZATION_STEP: self._iterations += 1 diff --git a/src/badger/utils.py b/src/badger/utils.py index d723137a..9ac99525 100644 --- a/src/badger/utils.py +++ b/src/badger/utils.py @@ -6,7 +6,7 @@ from datetime import datetime from importlib import metadata from types import TracebackType -from typing import TYPE_CHECKING, Any, Iterable, Optional +from typing import TYPE_CHECKING, Any, Iterable, Optional, TypedDict import yaml from pandas import DataFrame @@ -15,6 +15,8 @@ from badger.errors import BadgerLoadConfigError if TYPE_CHECKING: + from xopt.generators import Generator + from badger.routine import Routine logger = logging.getLogger(__name__) @@ -55,31 +57,33 @@ 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: bool = False, indentless: bool = False) -> Any: - return super(Dumper, self).increase_indent(flow, False) +class Dumper(yaml.Dumper): # type: ignore[misc] + def increase_indent(self, flow: bool = False, indentless: bool = False) -> None: + super(Dumper, self).increase_indent(flow, False) 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: Any) -> None: print(get_yaml_string(content), end="") -def norm(x: Any, lb: Any, ub: Any) -> Any: +def norm(x: float, lb: float, ub: float) -> float: return (x - lb) / (ub - lb) -def denorm(x: Any, lb: Any, ub: Any) -> Any: +def denorm(x: float, lb: float, ub: float) -> float: return (1 - x) * lb + x * ub -def config_list_to_dict(config_list: list[dict[str, Any]] | None) -> dict[str, Any]: +def config_list_to_dict(config_list: list[dict[str, Any]]) -> dict[str, Any]: if not config_list: return {} @@ -91,8 +95,8 @@ def config_list_to_dict(config_list: list[dict[str, Any]] | None) -> dict[str, A return book -def load_config(fname: str | None) -> Any: - configs = None +def load_config(fname: str | None) -> dict[str, Any] | None: + configs: dict[str, Any] | None = None if fname is None: return configs @@ -136,7 +140,7 @@ def merge_params( return merged_params -def range_to_str(vranges: list[dict[str, Any]]) -> list[dict[str, str]]: +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: list[dict[str, str]] = [] for var_dict in vranges: @@ -193,7 +197,12 @@ def create_archive_run_filename(routine: "Routine", format: str = "lcls-fname") return fname -def get_header(routine: Any) -> list[str]: +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 Exception: @@ -214,8 +223,10 @@ def get_header(routine: Any) -> list[str]: return obj_names + con_names + var_names + sta_names -def run_names_to_dict(run_names: list[str]) -> dict[str, Any]: - runs: dict[str, Any] = {} +def run_names_to_dict( + run_names: list[str], +) -> dict[str, dict[str, dict[str, list[str]]]]: + runs: dict[str, dict[str, dict[str, list[str]]]] = {} for name in run_names: name = os.path.basename( name @@ -266,8 +277,14 @@ def convert_str_to_value(str: str) -> Any: return str -def parse_rule(rule: str | dict[str, Any]) -> dict[str, str]: - 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", @@ -288,14 +305,10 @@ def parse_rule(rule: str | dict[str, Any]) -> dict[str, str]: except Exception: reducer = "percentile_80" - return { - "direction": direction, - "filter": filter, - "reducer": reducer, - } + return Rule(direction=direction, filter=filter, reducer=reducer) -def get_value_or_none(book: dict[Any, Any], key: Any) -> Any: +def get_value_or_none(book: dict[str, Any], key: str) -> Any: try: value = book[key] except KeyError: @@ -304,7 +317,7 @@ def get_value_or_none(book: dict[Any, Any], key: Any) -> Any: return value -def dump_state(dump_file: str | None, generator: Any, data: DataFrame) -> None: +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) @@ -314,7 +327,7 @@ def dump_state(dump_file: str | None, generator: Any, data: DataFrame) -> None: def state_to_dict( - generator: Any, data: DataFrame, include_data: bool = True + generator: "Generator", data: DataFrame, include_data: bool = True ) -> dict[str, Any]: # dump data to dict with config metadata output = { @@ -331,7 +344,7 @@ def state_to_dict( # https://stackoverflow.com/a/18472142 -def strtobool(val: Any) -> Any: +def strtobool(val: str) -> bool: """Convert a string representation of truth to true (1) or false (0). True values are 'y', 'yes', 't', 'true', 'on', and '1'; false values are 'n', 'no', 'f', 'false', 'off', and '0'. Raises ValueError if @@ -340,7 +353,7 @@ def strtobool(val: Any) -> Any: try: val = val.lower() except AttributeError: - return val + return bool(val) if val in ("y", "yes", "t", "true", "on", "1"): return True From e4dc21a59dede8346da704b564fee65ed0a325c0 Mon Sep 17 00:00:00 2001 From: Mitchell Victoriano Date: Mon, 15 Jun 2026 20:05:39 -0700 Subject: [PATCH 08/15] Working on typing of root badger directory files --- pyproject.toml | 5 +- src/badger/actions/run.py | 10 +-- src/badger/archive.py | 46 +++++++----- src/badger/core.py | 16 ++-- src/badger/core_subprocess.py | 74 ++++++++++--------- src/badger/environment.py | 72 ++++++++++-------- src/badger/factory.py | 32 ++++---- src/badger/formula.py | 6 +- src/badger/gui/components/create_process.py | 21 +++--- .../gui/components/plot_event_handlers.py | 61 ++++++++------- src/badger/interface.py | 6 +- src/badger/log.py | 4 +- src/badger/logbook.py | 25 ++++--- src/badger/logger/__init__.py | 34 ++++----- src/badger/logger/event.py | 10 +-- src/badger/routine.py | 59 +++++++++------ src/badger/settings.py | 63 +++++++++------- 17 files changed, 300 insertions(+), 244 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 540e70d4..29bbb23c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -42,6 +42,7 @@ dev = [ "pyqt5-stubs", "pandas-stubs", "types-requests", + "types-PyYAML", "mypy", ] @@ -72,6 +73,6 @@ ignore = ["E722"] # Until bare except blocks get fixed [tool.pytest.ini_options] addopts = "--cov=badger/" -log_cli_level = "info" -log_level = "debug" +log_cli_level = "INFO" +log_level = "DEBUG" testpaths = ["src/badger/tests"] diff --git a/src/badger/actions/run.py b/src/badger/actions/run.py index 98ec743b..40d52fcf 100644 --- a/src/badger/actions/run.py +++ b/src/badger/actions/run.py @@ -89,16 +89,16 @@ def after_evaluate(data: DataFrame) -> None: 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: pass # take a break to let the outside signal to change the status time.sleep(sleep) - def states_ready(states) -> None: + def states_ready(states: Any) -> None: storage["states"] = states try: @@ -115,7 +115,7 @@ def states_ready(states) -> None: 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: diff --git a/src/badger/archive.py b/src/badger/archive.py index 3f9d242b..9013abd5 100644 --- a/src/badger/archive.py +++ b/src/badger/archive.py @@ -1,12 +1,13 @@ +import logging import os import time import warnings -import logging +from typing import Any, Hashable, Optional, TypedDict -from badger.utils import ts_float_to_str -from badger.settings import init_settings -from badger.routine import Routine from badger.errors import BadgerConfigError +from badger.routine import Routine +from badger.settings import init_settings +from badger.utils import ts_float_to_str logger = logging.getLogger(__name__) @@ -21,12 +22,21 @@ 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: Optional[Any] # TODO: make this more specific + path: str + + +def archive_run( + routine: Routine, states: Optional[Any] = None +) -> RunArchive: # TODO: make states more specific data = routine.sorted_data 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 @@ -40,7 +50,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, @@ -58,14 +68,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") @@ -77,8 +87,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( [ @@ -128,9 +138,9 @@ def list_run(): return runs -def get_runs(): +def get_runs() -> list[str]: runs = list_run() - run_list = [] + run_list: list[str] = [] for year, months in runs.items(): for month, days in months.items(): for day, files in days.items(): @@ -156,7 +166,7 @@ def load_run(run_fname: str) -> Routine: # TODO: create utility function to catch warnings to remove code # duplication with warnings.catch_warnings(record=True) as caught_warnings: - routine = Routine.from_file(filename) + routine: Routine = Routine.from_file(filename) # Check if any user warnings were caught for warning in caught_warnings: @@ -168,11 +178,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]}" @@ -191,7 +201,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/core.py b/src/badger/core.py index b7aa3391..3f8ab1ea 100644 --- a/src/badger/core.py +++ b/src/badger/core.py @@ -107,16 +107,16 @@ 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, @@ -128,12 +128,11 @@ def run_routine( # 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: @@ -165,8 +164,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 a4743939..b2077cc5 100644 --- a/src/badger/core_subprocess.py +++ b/src/badger/core_subprocess.py @@ -1,32 +1,34 @@ -from copy import deepcopy import logging +import multiprocessing as mp +import os import time import traceback -from typing import Any +from copy import deepcopy +from multiprocessing.connection import Connection +from multiprocessing.synchronize import Event as EventType from queue import Empty +from typing import Any, Optional + from pandas import DataFrame -import multiprocessing as mp -import os +from xopt.errors import FeasibilityError, XoptError +from xopt.vocs import select_best -from badger.settings import ( - init_settings, - apply_pytorch_multiprocess_tensor_sharing_setting, -) from badger.errors import ( - BadgerRunTerminated, - BadgerEnvObsError, - MEASUREMENT_ERROR_TYPE, - MEASUREMENT_ACTION_TYPE, - MEASUREMENT_ACTION_RETRY, MEASUREMENT_ACTION_ABORT, + MEASUREMENT_ACTION_RETRY, + MEASUREMENT_ACTION_TYPE, + MEASUREMENT_ERROR_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.log import configure_process_logging -from xopt.errors import FeasibilityError, XoptError -from xopt.vocs import select_best - +from badger.settings import ( + apply_pytorch_multiprocess_tensor_sharing_setting, + init_settings, +) logger = logging.getLogger(__name__) @@ -34,9 +36,9 @@ def evaluate_measurement_with_retry( routine: Routine, point: Any, - queue: mp.Queue, - stop_process: mp.Event, - dialog_action_queue: mp.Queue, + queue: "mp.Queue[Any]", + stop_process: EventType, + dialog_action_queue: "mp.Queue[Any]", ) -> DataFrame: while True: try: @@ -77,7 +79,7 @@ def evaluate_measurement_with_retry( ) -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. @@ -91,7 +93,7 @@ def convert_to_solution(result: DataFrame, routine: Routine): 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: @@ -113,7 +115,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, @@ -129,14 +131,14 @@ def convert_to_solution(result: DataFrame, routine: Routine): def run_routine_subprocess( - queue: mp.Queue, - evaluate_queue: mp.Pipe, - stop_process: mp.Event, - pause_process: mp.Event, - wait_event: mp.Event, - config_path: str = None, - log_queue: mp.Queue = None, - dialog_action_queue: mp.Queue = None, + queue: "mp.Queue[Any]", + evaluate_queue: tuple[Connection, Connection], + stop_process: EventType, + pause_process: EventType, + wait_event: EventType, + dialog_action_queue: "mp.Queue[Any]", + config_path: Optional[str] = None, + log_queue: "Optional[mp.Queue[Any]]" = None, ) -> None: """ Run the provided routine object using Xopt. This method is run as a subproccess @@ -158,7 +160,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}", ) @@ -173,7 +175,7 @@ def run_routine_subprocess( apply_pytorch_multiprocess_tensor_sharing_setting(config_values) # Now load the archive would use the correct config - from badger.archive import load_run, archive_run + from badger.archive import archive_run, load_run logger.info("Waiting for wait_event to be set...") wait_event.wait() @@ -249,12 +251,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, diff --git a/src/badger/environment.py b/src/badger/environment.py index 6deda3f7..5458d9ea 100644 --- a/src/badger/environment.py +++ b/src/badger/environment.py @@ -1,9 +1,10 @@ from abc import abstractmethod from logging import warning -from typing import TYPE_CHECKING, Any, ClassVar, Dict, List, Optional +from typing import TYPE_CHECKING, Any, Callable, ClassVar, Optional from pydantic import BaseModel, ConfigDict, Field, SerializeAsAny from pydantic._internal._model_construction import ModelMetaclass + from badger.errors import ( BadgerEnvVarError, BadgerNoInterfaceError, @@ -11,12 +12,15 @@ if TYPE_CHECKING: 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 -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] @@ -33,13 +37,15 @@ 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]], +) -> Callable[[Any, list[str]], dict[str, 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]: # get the list of observable names needed by themselves and any formulas formula_observables = [] basic_observables = [] @@ -79,8 +85,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(): @@ -104,7 +112,9 @@ def validate(cls, variable_names: List[str]): class EnvMeta(ModelMetaclass): - def __new__(mcs, name: str, bases: tuple[type, ...], namespace: dict[str, Any]): + def __new__( + 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"]) @@ -116,7 +126,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): @@ -130,11 +140,11 @@ class BaseEnvironment(BaseModel, metaclass=EnvMeta): validate_assignment=True, use_enum_values=True, arbitrary_types_allowed=True ) name: ClassVar[str] = Field(description="environment name") - variables: ClassVar[Dict[str, list[float]]] # bounds list could be empty for var + variables: ClassVar[dict[str, list[float]]] # bounds list could be empty for var observables: ClassVar[list[str]] @abstractmethod - def get_variables(self, variable_names: list[str]) -> Dict[str, float]: + def get_variables(self, variable_names: list[str]) -> dict[str, float]: """ Get the values of the specified variables from the environment. @@ -151,7 +161,7 @@ def get_variables(self, variable_names: list[str]) -> Dict[str, float]: pass @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. @@ -164,8 +174,8 @@ def set_variables(self, variable_inputs: Dict[str, float]): @abstractmethod def get_observables( - self, observable_names: List[str] - ) -> Dict[str, float | List[float]]: + self, observable_names: list[str] + ) -> dict[str, float | list[float]]: """ Get the values of the specified observables from the environment. @@ -174,32 +184,32 @@ 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. """ pass - 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. """ pass - def get_system_states(self) -> Dict[str, Any]: + def get_system_states(self) -> dict[str, Any]: """ Get the current system states from the environment. This method is called to retrieve the current state of the environment. Returns ------- - Dict[str, Any] + dict[str, Any] A dictionary containing the current system states. """ return {} @@ -220,7 +230,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. @@ -255,7 +265,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. @@ -268,7 +278,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. @@ -295,41 +305,43 @@ 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) - 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/factory.py b/src/badger/factory.py index f09bc951..79c41f8c 100644 --- a/src/badger/factory.py +++ b/src/badger/factory.py @@ -1,23 +1,23 @@ -from typing import Any, TypedDict, cast, TYPE_CHECKING -from badger.settings import init_settings -from badger.utils import get_value_or_none +import importlib +import logging +import os +import re +import sys +from pathlib import Path +from typing import TYPE_CHECKING, Any, TypedDict, cast + +import yaml +from xopt.generators import generators, get_generator_defaults + from badger.errors import ( BadgerConfigError, - BadgerInvalidPluginError, BadgerInvalidDocsError, + BadgerInvalidPluginError, BadgerPluginNotFoundError, ) - from badger.interface import Interface as BadgerInterface -import sys -import os -import importlib -import yaml -import re -from pathlib import Path -from xopt.generators import generators, get_generator_defaults - -import logging +from badger.settings import init_settings +from badger.utils import get_value_or_none if TYPE_CHECKING: from badger.environment import Environment as BadgerEnvironment @@ -359,7 +359,7 @@ def _md_images_to_html( def repl(m: re.Match[str]) -> str: url_str = m.group(1).strip().strip("'\"").lstrip("./") - url = Path(base_prefix / url_str) # type: ignore[arg-type,operator] + url = Path(base_prefix / url_str) # type: ignore[operator] return f'' return _MD_IMG.sub(repl, text) @@ -405,7 +405,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()) diff --git a/src/badger/formula.py b/src/badger/formula.py index 9137e2fe..883cafd6 100644 --- a/src/badger/formula.py +++ b/src/badger/formula.py @@ -1,9 +1,9 @@ +import ast +import difflib +import re from typing import Any, Iterable import numpy as np -import re -import ast -import difflib def safe_var_name(var_name: str) -> str: diff --git a/src/badger/gui/components/create_process.py b/src/badger/gui/components/create_process.py index f84ea34f..fe5adea2 100644 --- a/src/badger/gui/components/create_process.py +++ b/src/badger/gui/components/create_process.py @@ -1,12 +1,13 @@ +import logging from multiprocessing import Event, Pipe, Process, Queue +from multiprocessing.connection import Connection +from typing import Any -from PyQt5.QtCore import pyqtSignal, QObject +from PyQt5.QtCore import QObject, pyqtSignal -from badger.settings import init_settings from badger.core_subprocess import run_routine_subprocess - -import logging from badger.log import get_logging_manager +from badger.settings import init_settings logger = logging.getLogger(__name__) @@ -30,11 +31,13 @@ def create_subprocess(self) -> None: """ self.stop_event = Event() self.pause_event = Event() - self.data_queue = Queue() - self.evaluate_queue = Pipe() + self.data_queue: "Queue[Any]" = Queue() + self.evaluate_queue: "tuple[Connection, Connection]" = Pipe() self.wait_event = Event() - self.dialog_action_queue = Queue() - config_path = init_settings()._instance.config_path + self.dialog_action_queue: "Queue[Any]" = 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() @@ -49,9 +52,9 @@ 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() diff --git a/src/badger/gui/components/plot_event_handlers.py b/src/badger/gui/components/plot_event_handlers.py index 37352d5f..f6860c63 100644 --- a/src/badger/gui/components/plot_event_handlers.py +++ b/src/badger/gui/components/plot_event_handlers.py @@ -1,21 +1,19 @@ -from badger.routine import Routine -from matplotlib.backend_bases import MouseEvent, MouseButton, PickEvent +import logging +from typing import Callable, 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 badger.gui.components.types import ConfigurableOptions from badger.gui.components.extension_utilities import ( HandledException, to_precision_float, ) - -from typing import cast - -import logging - -from matplotlib.collections import PathCollection -from matplotlib.text import Annotation -import pandas as pd -from pyparsing import Callable +from badger.gui.components.types import ConfigurableOptions +from badger.routine import Routine logger = logging.getLogger(__name__) @@ -54,7 +52,7 @@ def connect_events(self) -> None: self.canvas.mpl_connect( "button_press_event", lambda event: self.on_click( - event, # type: ignore[call-arg] + cast(MouseEvent, event), self.parameters, self.routine, self.callback, @@ -73,12 +71,12 @@ def connect_events(self) -> None: self.canvas.mpl_connect( "scroll_event", - lambda event: self.on_scroll(event), # type: ignore[call-arg] + lambda event: self.on_scroll(cast(MouseEvent, event)), ) self.canvas.mpl_connect( "pick_event", - lambda event: self.on_pick(event), # type: ignore[call-arg] + lambda event: self.on_pick(cast(PickEvent, event)), ) def update_reference_points( @@ -286,18 +284,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 @@ -315,18 +315,17 @@ def on_pick(self, event: PickEvent) -> None: 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() - ] + x_idx = (routine_data[x_column] - point[0]).abs().idxmin() + y_idx = (routine_data[y_column] - point[1]).abs().idxmin() + searched_row = pd.Series( + routine_data.loc[x_idx if x_idx == y_idx else x_idx] + ) elif x_column: - searched_row = routine_data.loc[ - (routine_data[x_column] - point[0]).abs().idxmin() - ] + idx = (routine_data[x_column] - point[0]).abs().idxmin() + searched_row = pd.Series(routine_data.loc[idx]) elif y_column: - searched_row = routine_data.loc[ - (routine_data[y_column] - point[1]).abs().idxmin() - ] + idx = (routine_data[y_column] - point[1]).abs().idxmin() + searched_row = pd.Series(routine_data.loc[idx]) true_index = searched_row.name @@ -374,7 +373,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. @@ -383,4 +382,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/interface.py b/src/badger/interface.py index 3dd97287..67ec3b71 100644 --- a/src/badger/interface.py +++ b/src/badger/interface.py @@ -78,12 +78,12 @@ def dump_recording(self, filename: str) -> None: # Environment should only call this method to get channels @abstractmethod - def get_values(self, channel_names: list[str]) -> dict[str, Any]: + def get_values(self, channel_names: list[str]) -> dict[str, float | list[float]]: pass # Environment should only call this method to set channels @abstractmethod - def set_values(self, channel_inputs: dict[str, Any]) -> Any: + def set_values(self, channel_inputs: dict[str, float | list[float]]) -> None: pass def reset_interface(self) -> None: @@ -96,7 +96,7 @@ def reset_interface(self) -> None: 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: Any, **kwargs: Any) -> Any: + 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 | None: diff --git a/src/badger/log.py b/src/badger/log.py index b1e8c52a..61fe6713 100644 --- a/src/badger/log.py +++ b/src/badger/log.py @@ -167,7 +167,7 @@ def get_logfile_name(self) -> str: log_filename = f"log_{today.year:04d}_{today.month:02d}_{today.day:02d}.log" return log_filename - def get_queue(self) -> Queue[logging.LogRecord] | None: + def get_queue(self) -> "Queue[logging.LogRecord] | None": """Get the logging queue for use by subprocesses.""" return self.log_queue @@ -201,7 +201,7 @@ def create_log_dir(self, log_dir: str) -> None: def configure_process_logging( - log_queue: Queue[logging.LogRecord] | None = None, + log_queue: "Queue[logging.LogRecord] | None" = None, logger_name: str = "badger", log_level: str = "DEBUG", process_name: str | None = None, diff --git a/src/badger/logbook.py b/src/badger/logbook.py index 99584f3f..661e6d2c 100644 --- a/src/badger/logbook.py +++ b/src/badger/logbook.py @@ -1,10 +1,13 @@ +import logging import os from datetime import datetime -import logging -from badger.settings import init_settings +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__) @@ -18,9 +21,9 @@ logger.info(f"Badger logbook root {BADGER_LOGBOOK_ROOT} created") -def send_to_logbook(routine, widget=None): - from xml.etree import ElementTree +def send_to_logbook(routine: Routine, widget: QWidget | None = None) -> None: from re import sub + from xml.etree import ElementTree log_text = "" routine_name = routine.name @@ -34,6 +37,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"]) @@ -57,7 +62,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") @@ -98,22 +103,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 f15e2666..6019999d 100644 --- a/src/badger/logger/__init__.py +++ b/src/badger/logger/__init__.py @@ -63,16 +63,16 @@ def _step( 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)) + " |" @@ -84,17 +84,17 @@ def _header( 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: str = "| " + " | ".join(cells) + " |" self._header_length = len(line) @@ -104,7 +104,7 @@ def _is_new_max( self, solution: Solution, ) -> bool: - return solution[4] + return solution.is_optimal def update( self, @@ -144,12 +144,12 @@ def __init__(self, path: str, reset: bool = True) -> None: 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 89acf2db..4d834091 100644 --- a/src/badger/logger/event.py +++ b/src/badger/logger/event.py @@ -1,12 +1,12 @@ from enum import StrEnum -from typing import NamedTuple +from typing import NamedTuple, Optional class Solution(NamedTuple): - variables: list[float] - objectives: list[float] - constraints: list[float] - states: list[float] + variables: Optional[list[float]] + objectives: Optional[list[float]] + constraints: Optional[list[float]] + states: Optional[list[float]] is_optimal: bool variable_names: list[str] objective_names: list[str] diff --git a/src/badger/routine.py b/src/badger/routine.py index 6a20ed36..2b6f7ed9 100644 --- a/src/badger/routine.py +++ b/src/badger/routine.py @@ -1,25 +1,27 @@ -import logging import json +import logging from copy import deepcopy -from typing import Any, List, Optional +from typing import Any, Optional + import numpy as np import pandas as pd from pandas import DataFrame from pydantic import ( ConfigDict, Field, - field_validator, - model_validator, SerializeAsAny, ValidationInfo, + field_validator, + model_validator, ) -from xopt import Evaluator, VOCS, Xopt +from xopt import VOCS, Evaluator, Xopt from xopt.generators import get_generator -from xopt.utils import get_local_region from xopt.generators.sequential import SequentialGenerator -from badger.utils import curr_ts -from badger.environment import BaseEnvironment, instantiate_env +from xopt.utils import get_local_region + +from badger.environment import BaseEnvironment, Environment, instantiate_env from badger.factory import get_env +from badger.utils import curr_ts logger = logging.getLogger(__name__) @@ -31,15 +33,15 @@ class Routine(Xopt): description: Optional[str] = Field(None) environment: SerializeAsAny[BaseEnvironment] initial_points: Optional[DataFrame] = Field(None) - critical_constraint_names: Optional[List[str]] = Field([]) - tags: Optional[List] = Field(None) + critical_constraint_names: Optional[list[str]] = Field([]) + tags: Optional[list] = Field(None) script: Optional[str] = Field(None) # Store relative to current params relative_to_current: Optional[bool] = Field(False) vrange_limit_options: Optional[dict] = Field(None) vrange_hard_limit: Optional[dict] = Field({}) # override hard limits - initial_point_actions: Optional[List] = Field(None) - additional_variables: Optional[List[str]] = Field([]) + initial_point_actions: Optional[list] = Field(None) + additional_variables: Optional[list[str]] = Field([]) formulas: Optional[dict[str, dict[str, Any]]] = Field({}) constraint_formulas: Optional[dict[str, dict[str, Any]]] = Field({}) observable_formulas: Optional[dict[str, dict[str, Any]]] = Field({}) @@ -51,7 +53,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())}") @@ -98,8 +100,11 @@ def validate_model(cls, data: Any): 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 = data["data"] + assert isinstance(df, pd.DataFrame) + 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): @@ -123,9 +128,11 @@ 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() env.set_variables(point) @@ -141,7 +148,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: @@ -151,7 +158,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: @@ -159,7 +166,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""" @@ -193,13 +200,15 @@ def json(self, **kwargs) -> str: return json.dumps(dict_result) -def calculate_variable_bounds(limit_options, vocs, env): +def calculate_variable_bounds( + limit_options: dict[str, Any], vocs: VOCS, env: Environment +) -> 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] @@ -231,10 +240,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: Environment +) -> 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 4273bb91..11fd2fea 100644 --- a/src/badger/settings.py +++ b/src/badger/settings.py @@ -1,13 +1,15 @@ +import logging import os import platform -import yaml import shutil from importlib import resources -from badger.utils import get_datadir +from typing import Any, Optional, Union + +import yaml from pydantic import BaseModel, Field, ValidationError -from typing import Any, Dict, Optional, Union + from badger.errors import BadgerLoadConfigError -import logging +from badger.utils import get_datadir logger = logging.getLogger(__name__) @@ -132,9 +134,14 @@ class BadgerConfig(BaseModel): class ConfigSingleton: - _instance = None - - def __new__(cls, config_path: str = None, user_flag: bool = False): + _instance: "Optional[ConfigSingleton]" = None + user_flag: bool + _config: BadgerConfig + config_path: str | None + + def __new__( + cls, config_path: str | None = None, user_flag: bool = False + ) -> "ConfigSingleton": if cls._instance is None: cls._instance = super(ConfigSingleton, cls).__new__(cls) cls._instance.user_flag = user_flag @@ -143,14 +150,14 @@ def __new__(cls, config_path: str = None, user_flag: bool = False): return 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 @@ -159,7 +166,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) @@ -191,7 +198,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) @@ -201,7 +208,7 @@ def load_or_create_config(cls, config_path: str) -> BadgerConfig: def config(self) -> BadgerConfig: return self._config - def update_and_save_config(self, updates: Dict[str, Any]) -> None: + def update_and_save_config(self, updates: dict[str, Any]) -> None: """Saves changes to the config file. Parameters @@ -216,13 +223,14 @@ 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) def _update_config_by_dot_key( - self, config_data: Dict[str, Any], dot_key: str, value: Any + self, config_data: dict[str, Any], dot_key: str, value: Any ) -> None: """Update the config data with the provided value using dot-separated keys.""" keys = dot_key.split(":") @@ -235,23 +243,23 @@ def _update_config_by_dot_key( else: d[last_key] = value - def list_settings(self) -> Dict[str, Any]: + def list_settings(self) -> dict[str, Any]: """List all the settings in Badger 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. """ return self._config.model_dump(by_alias=True) - def list_path_settings(self) -> Dict[str, Any]: + def list_path_settings(self) -> dict[str, Any]: """List all the path-related settings in Badger 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. @@ -388,7 +396,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 key in keys[:-1]: @@ -407,13 +415,13 @@ def reset_settings(self) -> None: ) -def init_settings(config_arg: str = None) -> ConfigSingleton: +def init_settings(config_arg: str | None = None) -> ConfigSingleton: """ Builds and returns an instance of the ConfigSingleton class. Parameters ---------- - config_arg: str + config_arg: str | None a path to a config file passed through the --config__filepath argument Returns @@ -468,7 +476,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 @@ -478,7 +486,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" @@ -488,7 +496,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") @@ -528,9 +536,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": From edd0083ba02ba33f76c4b4d60e66733c555abfe0 Mon Sep 17 00:00:00 2001 From: Mitchell Victoriano Date: Thu, 24 Sep 2026 03:32:06 -0700 Subject: [PATCH 09/15] Addresses many errors with extensions and supporting modules --- .pre-commit-config.yaml | 3 +- src/badger/actions/__init__.py | 4 +- src/badger/actions/install.py | 3 +- src/badger/actions/run.py | 3 +- src/badger/actions/uninstall.py | 3 +- .../environments/sphere_2d/__init__.py | 8 +- .../interfaces/default/__init__.py | 12 +-- .../gui/components/analysis_extensions.py | 6 +- src/badger/gui/components/analysis_widget.py | 9 ++- .../components/bax_visualizer/bax_widget.py | 4 +- .../gui/components/bax_visualizer/controls.py | 2 +- .../gui/components/bo_visualizer/bo_widget.py | 2 +- .../components/bo_visualizer/plotting_area.py | 2 +- .../gui/components/bo_visualizer/types.py | 8 +- .../gui/components/pf_viewer/pf_widget.py | 10 +-- src/badger/gui/components/pf_viewer/types.py | 7 +- .../gui/components/plot_event_handlers.py | 77 ++++++------------- src/badger/gui/components/types.py | 23 ++++-- src/badger/gui/mini/__init__.py | 21 ++--- src/badger/gui/mini/pages/routine_page.py | 2 +- src/badger/gui/utils.py | 9 ++- src/badger/interface.py | 5 +- src/badger/tests/test_gui_basic.py | 54 ++++++------- src/badger/utils.py | 19 ++--- 24 files changed, 143 insertions(+), 153 deletions(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index d79d02c2..e9b23562 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -31,8 +31,9 @@ repos: - id: ruff-check args: [--fix] - id: ruff-format + # - repo: https://github.com/pre-commit/mirrors-mypy - # rev: v2.1.0 + # rev: v2.3.1 # hooks: # - id: mypy # args: [--strict, --ignore-missing-imports] diff --git a/src/badger/actions/__init__.py b/src/badger/actions/__init__.py index 100a9314..c8e4efe9 100644 --- a/src/badger/actions/__init__.py +++ b/src/badger/actions/__init__.py @@ -25,9 +25,9 @@ def show_info(args: 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/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/run.py b/src/badger/actions/run.py index 9edccf2d..0c66fc97 100644 --- a/src/badger/actions/run.py +++ b/src/badger/actions/run.py @@ -13,6 +13,7 @@ import signal import sys import time +from argparse import Namespace from typing import Any, TypedDict from pandas import DataFrame @@ -138,7 +139,7 @@ def states_ready(states: Any) -> None: @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/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..ed2b96ba 100644 --- a/src/badger/built_in_plugins/interfaces/default/__init__.py +++ b/src/badger/built_in_plugins/interfaces/default/__init__.py @@ -2,6 +2,8 @@ Used by test environments and as the fallback when no hardware interface is configured.""" +from typing import Any + from badger import interface @@ -10,14 +12,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 +32,6 @@ def get_values(self, channel_names): return channel_outputs - def set_values(self, channel_inputs): + def set_values(self, channel_inputs: dict[str, float | list[float]]) -> None: for channel, value in channel_inputs.items(): self._states[channel] = value diff --git a/src/badger/gui/components/analysis_extensions.py b/src/badger/gui/components/analysis_extensions.py index 8246a19e..bc3886f6 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: 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/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..bfc95721 100644 --- a/src/badger/gui/components/bax_visualizer/controls.py +++ b/src/badger/gui/components/bax_visualizer/controls.py @@ -190,7 +190,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) diff --git a/src/badger/gui/components/bo_visualizer/bo_widget.py b/src/badger/gui/components/bo_visualizer/bo_widget.py index bfba76d6..26008635 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") diff --git a/src/badger/gui/components/bo_visualizer/plotting_area.py b/src/badger/gui/components/bo_visualizer/plotting_area.py index 91403818..b511e676 100644 --- a/src/badger/gui/components/bo_visualizer/plotting_area.py +++ b/src/badger/gui/components/bo_visualizer/plotting_area.py @@ -101,7 +101,7 @@ def update_plot( 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/pf_viewer/pf_widget.py b/src/badger/gui/components/pf_viewer/pf_widget.py index d04abbc0..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 # pyright: ignore[reportIncompatibleVariableOverride] - parameters: ConfigurableOptions = DEFAULT_PARAMETERS +class ParetoFrontWidget(AnalysisWidget[ConfigurableOptions]): + generator: MOBOGenerator + parameters: ConfigurableOptions = DEFAULT_PARAMETERS.copy() hypervolume_history: pd.DataFrame = pd.DataFrame() pf_1: Tensor | None = None @@ -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.astype(float) - y = data_points["hypervolume"].values.astype(float) + 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 52f5a51c..2b30f388 100644 --- a/src/badger/gui/components/pf_viewer/types.py +++ b/src/badger/gui/components/pf_viewer/types.py @@ -13,16 +13,15 @@ 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 diff --git a/src/badger/gui/components/plot_event_handlers.py b/src/badger/gui/components/plot_event_handlers.py index d1a2424a..e5a1cad4 100644 --- a/src/badger/gui/components/plot_event_handlers.py +++ b/src/badger/gui/components/plot_event_handlers.py @@ -16,7 +16,7 @@ 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__) @@ -28,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], @@ -46,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( - cast(MouseEvent, event), - 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(cast(MouseEvent, event)), - ) - - self.canvas.mpl_connect( - "pick_event", - lambda event: self.on_pick(cast(PickEvent, event)), - ) + 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 ( @@ -120,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") @@ -268,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 @@ -321,14 +291,13 @@ def on_pick(self, event: PickEvent) -> None: distances = (routine_data[x_column] - point[0]) ** 2 + ( routine_data[y_column] - point[1] ) ** 2 - idx = distances.idxmin() - searched_row = pd.Series(routine_data.loc[idx]) + searched_row = routine_data.iloc[int(distances.to_numpy().argmin())] elif x_column: - idx = (routine_data[x_column] - point[0]).abs().idxmin() - searched_row = pd.Series(routine_data.loc[idx]) + distances = (routine_data[x_column] - point[0]).abs() + searched_row = routine_data.iloc[int(distances.to_numpy().argmin())] elif y_column: - idx = (routine_data[y_column] - point[1]).abs().idxmin() - searched_row = pd.Series(routine_data.loc[idx]) + distances = (routine_data[y_column] - point[1]).abs() + searched_row = routine_data.iloc[int(distances.to_numpy().argmin())] true_index = searched_row.name 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/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/pages/routine_page.py b/src/badger/gui/mini/pages/routine_page.py index ebf52b69..2d163151 100644 --- a/src/badger/gui/mini/pages/routine_page.py +++ b/src/badger/gui/mini/pages/routine_page.py @@ -1075,7 +1075,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/utils.py b/src/badger/gui/utils.py index 0f6c83a6..020dae6e 100644 --- a/src/badger/gui/utils.py +++ b/src/badger/gui/utils.py @@ -8,7 +8,7 @@ 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 @@ -38,6 +38,9 @@ QWidget, ) +T = TypeVar("T") +P = ParamSpec("P") + class MouseWheelWidgetAdjustmentGuard(QObject): def __init__(self, parent: QObject): @@ -184,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/interface.py b/src/badger/interface.py index 84200206..eed21b6f 100644 --- a/src/badger/interface.py +++ b/src/badger/interface.py @@ -84,18 +84,19 @@ def dump_recording(self, filename: str) -> None: # Environment should only call this method to get channels @abstractmethod def get_values(self, channel_names: list[str]) -> dict[str, float | list[float]]: - pass + raise NotImplementedError() # Environment should only call this method to set channels @abstractmethod def set_values(self, channel_inputs: dict[str, float | list[float]]) -> None: - pass + raise NotImplementedError() 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 """ + raise NotImplementedError() def get_value(self, channel_name: str, **kwargs: Any) -> Any: return self.get_values([channel_name], **kwargs)[channel_name] diff --git a/src/badger/tests/test_gui_basic.py b/src/badger/tests/test_gui_basic.py index bfd3711a..106fd9f4 100644 --- a/src/badger/tests/test_gui_basic.py +++ b/src/badger/tests/test_gui_basic.py @@ -3,16 +3,17 @@ from unittest.mock import patch import pytest -from PyQt5.QtCore import QEventLoop, Qt, QTimer +from PyQt5.QtCore import QEventLoop, QTimer 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): +def test_gui_main(qtbot: QtBot, init_multiprocessing): from badger.gui.windows.main_window import BadgerMainWindow from badger.tests.utils import fix_path_issues @@ -79,7 +80,7 @@ def test_close_main(qtbot, init_multiprocessing): window.process_manager.close_proccesses() -def test_traceback_during_run(qtbot, init_multiprocessing): +def test_traceback_during_run(qtbot: QtBot) -> None: with patch("badger.core.run_routine") as run_routine_mock: run_routine_mock.side_effect = Exception("Test exception") @@ -131,7 +132,7 @@ def inner(ins, event): window.process_manager.close_proccesses() -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 @@ -194,7 +195,7 @@ def test_measurement_retry_dialog_in_app_flow(qtbot, init_multiprocessing): # 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 @@ -234,31 +235,32 @@ def test_default_low_noise_prior_in_bo(qtbot, init_multiprocessing): window.process_manager.close_proccesses() -def test_default_turbo_in_bo(qtbot): +# TODO: Fix default turbo in bo test +def test_default_turbo_in_bo(qtbot: QtBot) -> None: return - import yaml - from xopt.generators import all_generator_names + # 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 + # from badger.gui.windows.main_window import BadgerMainWindow + # from badger.tests.utils import fix_db_path_issue - fix_db_path_issue() + # fix_db_path_issue() - window = BadgerMainWindow() - qtbot.addWidget(window) + # 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 + # # 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) + # 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/utils.py b/src/badger/utils.py index 57e8c964..b5e44342 100644 --- a/src/badger/utils.py +++ b/src/badger/utils.py @@ -20,7 +20,7 @@ from badger.errors import BadgerLoadConfigError if TYPE_CHECKING: - from xopt.generators import Generator + from xopt.generator import Generator from badger.routine import Routine @@ -28,11 +28,6 @@ from gest_api.vocs import ContinuousVariable -if TYPE_CHECKING: - from xopt.generators import Generator - - from badger.routine import Routine - logger = logging.getLogger(__name__) @@ -60,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( @@ -204,14 +199,16 @@ 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 @@ -367,7 +364,7 @@ def state_to_dict( "vocs": json.loads(generator.vocs.model_dump_json()), } if include_data: - output["data"] = json.loads(data.to_json()) # type: ignore + output["data"] = json.loads(data.to_json()) return output From c3c9567df4c85f5c7375fadf99796d617f5b378f Mon Sep 17 00:00:00 2001 From: Mitchell Victoriano Date: Thu, 24 Sep 2026 20:40:36 -0700 Subject: [PATCH 10/15] Removed pyqt5-stubs as it has not been maintained; preferring native pyqt types --- pyproject.toml | 1 - src/badger/gui/components/action_bar.py | 94 +++++++------ .../gui/components/analysis_extensions.py | 2 +- src/badger/gui/components/archive_search.py | 108 +++++++++------ .../gui/components/bax_visualizer/controls.py | 6 +- .../gui/components/bo_visualizer/bo_widget.py | 3 +- .../components/bo_visualizer/plotting_area.py | 3 + src/badger/gui/components/bounds_preview.py | 4 +- src/badger/gui/components/collapsible_box.py | 101 +++++++++----- src/badger/gui/components/constraint_item.py | 42 ++++-- src/badger/gui/components/data_panel.py | 92 +++++++++---- src/badger/gui/components/data_table.py | 125 +++++++++++++----- src/badger/gui/components/editable_table.py | 40 +++--- src/badger/gui/components/env_cbox.py | 3 +- src/badger/gui/components/navigators.py | 51 ++++--- .../gui/windows/load_data_from_run_dialog.py | 5 +- 16 files changed, 440 insertions(+), 240 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index d4d5fe43..ca9c9e79 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -42,7 +42,6 @@ dev = [ "pytest-cov", "pytest-qt", "pytest-mock", - "pyqt5-stubs", "pandas-stubs", "types-requests", "types-PyYAML", diff --git a/src/badger/gui/components/action_bar.py b/src/badger/gui/components/action_bar.py index 85bd5918..d3eee3a1 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 Any, 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, @@ -25,7 +26,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 +36,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) @@ -189,13 +198,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 +301,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 +328,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 +338,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 +348,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 +373,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 +399,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 +410,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,54 +420,54 @@ 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): + def update_run_tooltip(self, tc: dict[str, Any] | None = None) -> None: """Update btn_stop tooltip: tc dict for run-until mode, or None.""" if tc is None: self.run_action.setToolTip("Run") diff --git a/src/badger/gui/components/analysis_extensions.py b/src/badger/gui/components/analysis_extensions.py index bc3886f6..c9b034c4 100644 --- a/src/badger/gui/components/analysis_extensions.py +++ b/src/badger/gui/components/analysis_extensions.py @@ -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/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/controls.py b/src/badger/gui/components/bax_visualizer/controls.py index bfc95721..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("") @@ -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 26008635..a5aacb4f 100644 --- a/src/badger/gui/components/bo_visualizer/bo_widget.py +++ b/src/badger/gui/components/bo_visualizer/bo_widget.py @@ -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 b511e676..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) 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/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/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..3759e01b 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): @@ -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: @@ -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( diff --git a/src/badger/gui/components/env_cbox.py b/src/badger/gui/components/env_cbox.py index 7a664e52..ed81c630 100644 --- a/src/badger/gui/components/env_cbox.py +++ b/src/badger/gui/components/env_cbox.py @@ -31,6 +31,7 @@ QHBoxLayout, QLabel, QLineEdit, + QListWidget, QPushButton, QSizePolicy, QStyledItemDelegate, @@ -598,7 +599,7 @@ def toggle_sta_show_mode(self, _: Any) -> None: def filter_sta(self) -> None: self.sta_table.update_keyword(self.edit_sta.text()) - def _fit_content(self, list) -> None: + 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) diff --git a/src/badger/gui/components/navigators.py b/src/badger/gui/components/navigators.py index e661684b..2cce16f0 100644 --- a/src/badger/gui/components/navigators.py +++ b/src/badger/gui/components/navigators.py @@ -3,8 +3,10 @@ TemplateNavigator shows routine template files. Both support right-click context menus (open, copy path, etc.).""" +import logging import os +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 +26,8 @@ from badger.settings import init_settings from badger.utils import run_names_to_dict +logger = logging.getLogger(__name__) + class FileContextMenuBase: """ @@ -34,7 +38,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,7 +47,7 @@ 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() clip.setText(fullpath) # Display a little popup box a bit after user's click @@ -56,10 +60,10 @@ 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 @@ -82,8 +86,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) @@ -106,14 +110,16 @@ def __init__(self): 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. """ @@ -129,7 +135,7 @@ def _firstSelectableItem(self, parent=None): 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: @@ -172,24 +178,31 @@ 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) for k in range(month_item.childCount()): @@ -200,19 +213,19 @@ def _selectItemByRun(self, 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 +249,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 +265,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 @@ -303,7 +316,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 +325,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/windows/load_data_from_run_dialog.py b/src/badger/gui/windows/load_data_from_run_dialog.py index e9346c2a..28fd064a 100644 --- a/src/badger/gui/windows/load_data_from_run_dialog.py +++ b/src/badger/gui/windows/load_data_from_run_dialog.py @@ -57,7 +57,7 @@ class BadgerLoadDataFromRunDialog(QDialog): def __init__( self, parent: QWidget, - env_vocs: VOCS = None, + env_vocs: VOCS | None = None, on_set: Callable[[Routine], None] | None = None, ): """ @@ -65,7 +65,8 @@ def __init__( Args: parent (QWidget): The parent widget. - data_table (QTableWidget, optional): The data table to update with loaded run data. + env_vocs (VOCS | None, optional): 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. From aa810fb880cf10f683d1e2fcb758f67efcf5c260 Mon Sep 17 00:00:00 2001 From: Mitchell Victoriano Date: Mon, 28 Sep 2026 18:34:13 -0700 Subject: [PATCH 11/15] working through components folder --- .pre-commit-config.yaml | 18 ++-- src/badger/archive.py | 4 +- src/badger/environment.py | 4 +- src/badger/gui/components/con_table.py | 10 +- src/badger/gui/components/editable_table.py | 36 +++---- src/badger/gui/components/eliding_label.py | 21 ++-- src/badger/gui/components/env_cbox.py | 16 ++-- src/badger/gui/components/filter_cbox.py | 12 ++- src/badger/gui/components/generator_cbox.py | 14 +-- src/badger/gui/components/labeled_lineedit.py | 8 +- src/badger/gui/components/navigators.py | 57 +++++++++-- src/badger/gui/components/obj_table.py | 14 ++- src/badger/gui/components/obs_table.py | 12 +-- src/badger/gui/components/process_manager.py | 42 +++++--- src/badger/gui/components/pydantic_editor.py | 47 +++++---- .../gui/components/reorderable_table.py | 64 +++++++++---- src/badger/gui/components/routine_editor.py | 20 ++-- src/badger/gui/components/routine_item.py | 43 +++++---- src/badger/gui/components/routine_page.py | 14 +-- src/badger/gui/components/routine_runner.py | 95 ++++++++++++------- src/badger/gui/components/search_bar.py | 2 +- src/badger/gui/components/state_item.py | 9 +- src/badger/gui/components/status_bar.py | 2 +- src/badger/gui/components/syntax.py | 42 +++++--- src/badger/gui/components/var_table.py | 47 +++++---- src/badger/routine.py | 2 +- 26 files changed, 419 insertions(+), 236 deletions(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index e9b23562..0f3d1a06 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -32,12 +32,12 @@ repos: 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 + - 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/src/badger/archive.py b/src/badger/archive.py index e8b7f61c..afed6741 100644 --- a/src/badger/archive.py +++ b/src/badger/archive.py @@ -45,6 +45,8 @@ def archive_run( ) -> 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") and routine.creation_ts is not None: suffix = routine.creation_ts @@ -175,7 +177,7 @@ def load_run(run_fname: str) -> Routine: # TODO: create utility function to catch warnings to remove code # duplication with warnings.catch_warnings(record=True) as caught_warnings: - routine: Routine = Routine.from_file(filename) + routine = Routine.from_file(filename) # Check if any user warnings were caught for warning in caught_warnings: diff --git a/src/badger/environment.py b/src/badger/environment.py index 52d00592..1605be0c 100644 --- a/src/badger/environment.py +++ b/src/badger/environment.py @@ -225,9 +225,7 @@ def get_system_states(self) -> dict[str, Any]: """ 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 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/editable_table.py b/src/badger/gui/components/editable_table.py index 3759e01b..f34fbd59 100644 --- a/src/badger/gui/components/editable_table.py +++ b/src/badger/gui/components/editable_table.py @@ -88,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 @@ -110,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: """ @@ -142,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()) @@ -167,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. @@ -191,7 +191,7 @@ def dragMoveEvent(self, e: QDragMoveEvent | None) -> None: Parameters ---------- - event : QDragMoveEvent + e : QDragMoveEvent | None The drag move event. """ @@ -221,7 +221,7 @@ def dropEvent(self, event: QDropEvent | None) -> None: Parameters ---------- - event : QDropEvent + event : QDropEvent | None The drop event. """ @@ -334,7 +334,7 @@ def insert_item( self, row: int, name: str, - info: list[Any], + info: tuple[str, float, bool], selected: bool = False, ) -> None: """ @@ -346,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. """ @@ -637,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, @@ -649,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: @@ -691,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 b6a281d9..63b41951 100644 --- a/src/badger/gui/components/eliding_label.py +++ b/src/badger/gui/components/eliding_label.py @@ -48,17 +48,20 @@ class ElidingLabel(QLabel): def __init__( self, text: str = "", - mode: Qt.TextElideMode = Qt.ElideMiddle, + mode: Qt.TextElideMode = Qt.TextElideMode.ElideMiddle, **kwargs: Any, ) -> None: super().__init__(**kwargs) self._mode = mode self.is_elided = False - self.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Preferred) + self.setSizePolicy(QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Preferred) self.setText(text) - def setText(self, text: str) -> None: + 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 @@ -72,7 +75,7 @@ def setText(self, text: str) -> None: def text(self) -> str: return self._contents - def paintEvent(self, event: QPaintEvent) -> None: + def paintEvent(self, event: QPaintEvent | None) -> None: super().paintEvent(event) did_elide = False @@ -111,19 +114,21 @@ def paintEvent(self, event: QPaintEvent) -> None: class SimpleElidedLabel(QLabel): - def __init__(self, text: str = "", parent: QWidget | None = None) -> None: + def __init__(self, text: str | None = None, parent: QWidget | None = None) -> None: super().__init__(parent) self._text = text - def setText(self, text: str) -> None: + def setText(self, text: str | None) -> None: self._text = text super().setText(self.elidedText()) - def resizeEvent(self, event: QResizeEvent) -> None: + def resizeEvent(self, event: QResizeEvent | None) -> None: super().setText(self.elidedText()) super().resizeEvent(event) def elidedText(self) -> str: metrics = QFontMetrics(self.font()) - elided = metrics.elidedText(self._text, Qt.ElideRight, self.width()) + 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 ed81c630..36b820ff 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 @@ -38,7 +38,6 @@ QVBoxLayout, QWidget, ) -from xopt.vocs import VOCS from badger.errors import BadgerRoutineError from badger.gui.components.collapsible_box import CollapsibleBox @@ -522,7 +521,6 @@ def init_ui(self) -> None: cbox_more.setContentLayout(vbox_more) def config_logic(self) -> None: - self.dict_con = {} self.edit_var.textChanged.connect(self.filter_var) self.check_only_var.stateChanged.connect(self.toggle_var_show_mode) @@ -635,18 +633,18 @@ 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, tuple[str, float]] = {} critical_constraints: list[str] = [] for constraint in self.con_table.export_data(): con_name = next(iter(constraint)) relation, threshold, critical = constraint[con_name] - constraints[con_name] = [CONS_RELATION_DICT[relation], threshold] + constraints[con_name] = (CONS_RELATION_DICT[relation], threshold) if critical: critical_constraints.append(con_name) @@ -657,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 2cce16f0..88fda3ef 100644 --- a/src/badger/gui/components/navigators.py +++ b/src/badger/gui/components/navigators.py @@ -5,6 +5,7 @@ import logging import os +from typing import cast from PyQt5 import QtCore from PyQt5.QtCore import QDir, Qt, QTimer, QUrl @@ -49,6 +50,9 @@ def show_context_menu(self, widget: QWidget, fullpath: str) -> None: # Functions to execute the context-menu actions 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( @@ -67,8 +71,12 @@ 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) @@ -99,11 +107,14 @@ def __init__(self, parent: QWidget | None = None) -> None: # 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 ) @@ -125,10 +136,15 @@ def _firstSelectableItem( """ 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: @@ -147,21 +163,36 @@ def updateItems(self, runs: list[str] | None = None) -> 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) @@ -205,10 +236,16 @@ def _selectItemByRun(self, run: str) -> 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 @@ -283,7 +320,9 @@ def __init__(self, parent: QWidget | None = None) -> 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 ) 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/process_manager.py b/src/badger/gui/components/process_manager.py index 8beee914..5551217b 100644 --- a/src/badger/gui/components/process_manager.py +++ b/src/badger/gui/components/process_manager.py @@ -1,9 +1,25 @@ """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 dataclasses import dataclass +from multiprocessing import Process, Queue +from multiprocessing.synchronize import Event + from PyQt5.QtCore import QObject, pyqtSignal +@dataclass +class ProcessWithArgs: + process: Process + stop_event: Event + pause_event: Event + args_queue: Queue + data_queue: Queue + evaluate_queue: Queue + wait_event: Event + dialog_action_queue: Queue + + class ProcessManager(QObject): """ The ProcessManager class is for holding an array of live processes @@ -14,26 +30,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 +58,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 99573e79..7c342c15 100644 --- a/src/badger/gui/components/pydantic_editor.py +++ b/src/badger/gui/components/pydantic_editor.py @@ -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)) @@ -206,7 +208,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 +286,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): @@ -645,8 +648,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) @@ -1205,8 +1208,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() @@ -1246,8 +1255,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) @@ -1366,9 +1379,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 }') @@ -1378,9 +1392,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/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 88292288..282c7725 100644 --- a/src/badger/gui/components/routine_page.py +++ b/src/badger/gui/components/routine_page.py @@ -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) @@ -232,7 +234,7 @@ def init_ui(self) -> None: 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) -> None: "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) @@ -409,7 +411,7 @@ def load_template_yaml(self, template_path: str | None = None) -> None: 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..bcbe8e70 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 dataclasses import dataclass +from multiprocessing import Process +from multiprocessing.synchronize import Event import pandas as pd from PyQt5.QtCore import QObject, QTimer, pyqtSignal @@ -39,6 +42,22 @@ logger = logging.getLogger(__name__) +@dataclass +class ArgumentDict: + routine_id: int | None + routine_filename: str | None + routine_name: str + variable_ranges: dict + initial_points: pd.DataFrame + evaluate: bool + archive: bool + termination_condition: dict + start_time: float + testing: bool + run_data: bool + init_points: bool + + class BadgerRoutineSignals(QObject): env_ready = pyqtSignal(list) finished = pyqtSignal() @@ -58,7 +77,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 +87,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 """ @@ -92,12 +111,14 @@ def __init__( self.termination_condition = ( 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.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 = None - self.stop_event = None - self.pause_event = None - self.routine_process = 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 @@ -125,6 +146,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 +161,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 @@ -167,29 +192,33 @@ def run(self, run_data_flag: bool = False, init_points_flag: bool = False) -> No 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 = ArgumentDict( + 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() 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 64ceafbf..f4f45549 100644 --- a/src/badger/gui/components/status_bar.py +++ b/src/badger/gui/components/status_bar.py @@ -33,7 +33,7 @@ def init_ui(self) -> None: 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;") 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/var_table.py b/src/badger/gui/components/var_table.py index 6bbed00a..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,12 +124,12 @@ 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 @@ -226,7 +229,7 @@ def validate_row(self, row: int) -> None: def set_bounds( self, variables: dict[str, tuple[float, float]], signal: bool = True - ): + ) -> None: for name, bounds in variables.items(): self.bounds[name] = bounds @@ -278,7 +281,7 @@ def update_selected(self) -> None: self.show_checked_only() def set_selected(self, variable_names: list[str]) -> None: - self.selected: dict[str, bool] = {} + self.selected = {} for vname in variable_names: self.selected[vname] = True @@ -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) @@ -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) @@ -577,6 +582,8 @@ def try_insert_variable(self, name: str) -> None: 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) @@ -661,19 +668,23 @@ def dropEvent(self, event: QDropEvent | None) -> None: else: event.ignore() - def flags(self, index) -> Qt.ItemFlags: + 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) -> 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) -> 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 diff --git a/src/badger/routine.py b/src/badger/routine.py index 5d1fc0ae..7cada305 100644 --- a/src/badger/routine.py +++ b/src/badger/routine.py @@ -32,7 +32,6 @@ field_validator, model_validator, ) -from xopt import VOCS, Evaluator, Xopt from xopt.generators import get_generator from xopt.generators.sequential import SequentialGenerator from xopt.vocs import get_local_region @@ -40,6 +39,7 @@ from badger.environment import BaseEnvironment, Environment, instantiate_env from badger.factory import get_env from badger.utils import curr_ts +from xopt import VOCS, Evaluator, Xopt logger = logging.getLogger(__name__) From 793e1eb14e62be011a454ee7a1547babf8956aef Mon Sep 17 00:00:00 2001 From: Mitchell Victoriano Date: Tue, 29 Sep 2026 08:53:57 -0700 Subject: [PATCH 12/15] Working through routine and run monitor --- mypy_deferred_errors.md | 327 ------------------ .../interfaces/default/__init__.py | 3 +- src/badger/environment.py | 19 +- src/badger/gui/components/process_manager.py | 15 +- src/badger/gui/components/routine_page.py | 2 +- src/badger/gui/components/routine_runner.py | 86 +++-- src/badger/gui/components/run_monitor.py | 6 +- .../gui/windows/measurement_retry_dialog.py | 12 +- .../gui/windows/termination_reached_dialog.py | 34 +- src/badger/interface.py | 4 +- src/badger/log.py | 21 +- src/badger/routine.py | 52 +-- src/badger/tests/conftest.py | 54 +-- .../multiobjective_test/__init__.py | 4 +- .../plugins/environments/test/__init__.py | 6 +- .../mock/plugins/interfaces/test/__init__.py | 8 +- src/badger/tests/multiprocess_logging.py | 8 +- src/badger/tests/test_cli_basic.py | 29 +- src/badger/tests/test_core.py | 2 +- src/badger/tests/test_core_subprocess.py | 107 ++---- src/badger/tests/test_create_process.py | 11 +- src/badger/tests/test_env.py | 16 +- src/badger/tests/test_environment.py | 38 +- 23 files changed, 281 insertions(+), 583 deletions(-) delete mode 100644 mypy_deferred_errors.md diff --git a/mypy_deferred_errors.md b/mypy_deferred_errors.md deleted file mode 100644 index 51f2c226..00000000 --- a/mypy_deferred_errors.md +++ /dev/null @@ -1,327 +0,0 @@ -# Mypy Deferred Errors - -This document captures the full state of `mypy --strict --ignore-missing-imports` on the badger repository as of 2026-06-09. The first run produced **1506 errors in 116 files**. After this PR's mechanical fixes (annotation-only changes, no runtime behavior changes) the count is **1473 errors in 98 files**. - -The remaining errors fall into two groups: - -1. **Mechanical work not yet done** — straightforward signature additions across GUI components and tests. Listed in §1 with file-by-file scopes for follow-up PRs. -2. **Real issues requiring business-logic review** — not mechanical fixes. Listed in §2 with concrete file:line references and what each one means. - ---- - -## §1. Remaining mechanical work (deferred to follow-up PRs) - -These are pure annotation additions. No runtime behavior changes required. Work can proceed file-by-file in independent PRs. - -### 1a. Test files — add `-> None` to test functions, type fixture parameters as `Any` - -| File | Approx. errors | -| --- | --- | -| `src/badger/tests/test_environment.py` | 40 | -| `src/badger/tests/test_formulas.py` | 39 | -| `src/badger/tests/test_routine_page.py` | 31 | -| `src/badger/tests/test_run_monitor.py` | 27 | -| `src/badger/tests/utils.py` | 17 | -| `src/badger/tests/test_settings.py` | 17 | -| `src/badger/tests/test_routine_runner.py` | 12 | -| `src/badger/tests/test_gui_basic.py` | 11 | -| `src/badger/tests/conftest.py` | 11 | -| `src/badger/tests/test_env.py` | 9 | -| `src/badger/tests/test_routine.py` | 4 | -| `src/badger/tests/test_intf.py` | 7 | -| `src/badger/tests/test_core_subprocess.py` | 8 | -| `src/badger/tests/test_cli_basic.py` | 7 | -| `src/badger/tests/test_core.py` | 6 | -| `src/badger/tests/test_lib_basic.py` | 2 | -| `src/badger/tests/test_factory.py` | 1 | -| `src/badger/tests/test_home_page.py` | 2 | -| `src/badger/tests/test_create_process.py` | 2 | -| `src/badger/tests/test_routine_id.py` (`x-test_routine_id.py`) | 4 | -| `src/badger/tests/x-test_db.py` | 2 | -| `src/badger/tests/test_settings.py` (continued) | — | -| `src/badger/tests/test_measurement_retry_dialog.py` | 2 | -| `src/badger/tests/test_process_manager.py` | 3 | -| `src/badger/tests/multiprocess_logging.py` | 3 | -| `src/badger/tests/mock/plugins/...` | ~10 | - -**Pattern**: -- `def test_xxx(...)` → `def test_xxx(...) -> None` -- `def fixture_xxx(mocker, qtbot, request)` → `def fixture_xxx(mocker: Any, qtbot: Any, request: Any) -> ...` -- `points_eval_list = []` → `points_eval_list: list[Any] = []` -- `out, err = "", ""` after `out: str` annotation in `test_cli_basic.py:8-9`. - -### 1b. GUI windows / dialogs — add signatures to `__init__`, `init_ui`, `config_logic`, slot handlers - -| File | Approx. errors | -| --- | --- | -| `src/badger/gui/windows/load_data_from_run_dialog.py` | 17 | -| `src/badger/gui/windows/settings_dialog.py` | 18 | -| `src/badger/gui/windows/ind_lim_vrange_dialog.py` | 14 | -| `src/badger/gui/windows/lim_vrange_dialog.py` | 11 | -| `src/badger/gui/windows/add_random_dialog.py` | 11 | -| `src/badger/gui/windows/main_window.py` | 12 | -| `src/badger/gui/windows/docs_window.py` | 11 | -| `src/badger/gui/windows/expandable_message_box.py` | 7 | -| `src/badger/gui/windows/edit_script_dialog.py` | 9 | -| `src/badger/gui/windows/var_dialog.py` | 7 | -| `src/badger/gui/windows/terminition_condition_dialog.py` | 11 | -| `src/badger/gui/windows/measurement_retry_dialog.py` | 1 | -| `src/badger/gui/windows/message_dialog.py` | 4 | -| `src/badger/gui/windows/review_dialog.py` | 4 | - -**Pattern**: -- `def __init__(self, ...):` → `def __init__(self, ...) -> None:` -- `def init_ui(self):` → `def init_ui(self) -> None:` -- `def config_logic(self):` → `def config_logic(self) -> None:` -- All event-handler slots → `-> None`. - -### 1c. GUI components/pages — same mechanical patterns - -| File | Approx. errors | -| --- | --- | -| `src/badger/gui/pages/home_page.py` | 40 (mostly logic — see §2) | -| `src/badger/gui/components/run_monitor.py` | 121 (mostly logic — see §2) | -| `src/badger/gui/components/routine_page.py` | 119 (logic-heavy — see §2) | -| `src/badger/gui/components/pydantic_editor.py` | 88 (some logic — see §2) | -| `src/badger/gui/components/var_table.py` | 60 (logic-heavy — see §2) | -| `src/badger/gui/components/data_table.py` | 26 (mostly mechanical) | -| `src/badger/gui/components/navigators.py` | 35 | -| `src/badger/gui/components/archive_search.py` | 26 | -| `src/badger/gui/components/action_bar.py` | 24 | -| `src/badger/gui/components/env_cbox.py` | 22 | -| `src/badger/gui/components/routine_runner.py` | 20 (mostly logic — see §2) | -| `src/badger/gui/components/editable_table.py` | 19 | -| `src/badger/gui/components/collapsible_box.py` | 17 | -| `src/badger/gui/components/bo_visualizer/bo_widget.py` | 17 | -| `src/badger/gui/components/plot_event_handlers.py` | 16 | -| `src/badger/gui/components/bo_visualizer/ui_components.py` | 16 | -| `src/badger/gui/components/syntax.py` | 15 | -| `src/badger/gui/components/routine_item.py` | 15 | -| `src/badger/gui/components/data_panel.py` | 14 (some logic — see §2) | -| `src/badger/gui/components/routine_editor.py` | 12 | -| `src/badger/gui/components/eliding_label.py` | 12 | -| `src/badger/gui/components/pf_viewer/pf_widget.py` | 40 | -| `src/badger/gui/components/process_manager.py` | 4 | -| `src/badger/gui/components/extension_utilities.py` | 2 | -| `src/badger/gui/components/extensions_palette.py` | 11 | -| `src/badger/gui/components/create_process.py` | 4 | -| `src/badger/gui/components/state_item.py` | 3 | -| `src/badger/gui/components/status_bar.py` | 6 | -| `src/badger/gui/components/search_bar.py` | 1 | -| `src/badger/gui/components/labeled_lineedit.py` | 1 | -| `src/badger/gui/components/robust_spinbox.py` | 1 | -| `src/badger/gui/components/reorderable_table.py` | 4 | -| `src/badger/gui/components/obs_table.py` | 1 | -| `src/badger/gui/components/obj_table.py` | 1 | -| `src/badger/gui/components/con_table.py` | — | -| `src/badger/gui/components/filter_cbox.py` | 3 | -| `src/badger/gui/components/generator_cbox.py` | 5 | -| `src/badger/gui/components/constraint_item.py` | 3 | -| `src/badger/gui/components/bo_visualizer/plotting_area.py` | 4 | -| `src/badger/gui/components/analysis_widget.py` | 1 | -| `src/badger/gui/components/analysis_extensions.py` | 1 | -| `src/badger/gui/utils.py` | 9 | -| `src/badger/gui/__init__.py` | 4 (one is logic — see §2) | - -### 1d. Plugin scaffolding - -| File | Errors | -| --- | --- | -| `src/badger/built_in_plugins/interfaces/default/__init__.py` | 5 | -| `src/badger/built_in_plugins/environments/sphere_2d/__init__.py` | 4 | -| `src/badger/tests/mock/plugins/interfaces/test/__init__.py` | 5 | -| `src/badger/tests/mock/plugins/environments/test/__init__.py` | 3 | -| `src/badger/tests/mock/plugins/environments/multiobjective_test/__init__.py` | 2 | - -### 1e. Other core/library modules with remaining mechanical work - -| File | Errors | -| --- | --- | -| `src/badger/db.py` | 27 (some logic — see §2.G) | -| `src/badger/log.py` | 23 | -| `src/badger/settings.py` | 18 (mostly logic — see §2.J) | -| `src/badger/routine.py` | 17 (some logic — see §2.E) | -| `src/badger/environment.py` | 29 (some logic — see §2.F) | -| `src/badger/core_subprocess.py` | 29 (mostly logic — see §2.H) | -| `src/badger/archive.py` | 9 | -| `src/badger/logger/__init__.py` | 29 | -| `src/badger/logbook.py` | 4 (one logic — see §2.K) | - ---- - -## §2. Errors requiring business-logic review (NOT addressed in this PR) - -These are real type-correctness issues — fixing them changes runtime behavior or requires a design decision. Each entry: file, line(s), what mypy says, why it's deferred. - -### A. Truthy-function checks in `core.py` - -`src/badger/core.py:108,132,136,165,169` — code like: -```python -if states_callback and (states is not None): -if evaluate_callback: -if dump_file_callback: -``` -Mypy: `[truthy-function]`. A function object is always truthy — these checks always pass. The intent appears to be `is not None`, but the parameters are typed `Callable[..., Any]` (non-Optional) except `dump_file_callback`. **Decision needed**: does each call site pass `None` for these callbacks, or always a callable? If never `None`, drop the checks; if sometimes `None`, change type to `Optional` and check `is not None`. - -### B. Unguarded `Optional` attribute access — `run_monitor.py` - -`src/badger/gui/components/run_monitor.py` — `self.routine` is initialized to `None` and accessed without None checks throughout. Specific lines: -- 109 (`vocs` property) -- 467, 469, 471, 473, 475 (`start()`) -- 503, 528, 538, 547 (data accessors) -- 602, 605, 617 (constraint plotting) -- 654, 655 (run_filename access) -- 688, 693 (routine + environment access) -- 795, 901, 909 (more environment / vocs) -- 956, 968, 969, 970 (vrange_hard_limit) -- 1053, 1107, 1109, 1115–1118, 1149–1152, 1161–1165 (more cascading) - -**Decision needed**: either -- Restructure so `self.routine` becomes non-Optional after a setter is called (and all access goes through that setter), or -- Add `if self.routine is None: return` guards at the top of every method that touches it. - -Also: `run_monitor.py:281–350` has `Cannot determine type of "plot_con"/"plot_obs"` errors — class attributes referenced before being assigned in `__init__`. Same issue as below in `routine_page.py`. These are forward-reference bugs in the class init path. - -### C. Unguarded `Optional` attribute access — `data_panel.py` - -`src/badger/gui/components/data_panel.py:139,170,226,267,272,347,382` — similar pattern: `self.run_monitor` or similar set to `None` then accessed. Plus: -- Line 181: `get_data_from_dialog` declared `-> None` but caller uses return value (`func-returns-value`). -- Line 267, 382: missing dict generic args. - -### D. Unguarded `Optional` access — `home_page.py`, `var_table.py` - -`src/badger/gui/pages/home_page.py:306,546,547,548,596,990` — `Item "None" of "Any | None" has no attribute "data"` etc. -`src/badger/gui/components/var_table.py:540,541,637,645` — `"None" has no attribute "get_variable"/"get_bounds"/"get_info"` and `Item "None" of "Any | None" has no attribute "text"`. - -Same pattern: attribute initialized to `None`, accessed unconditionally. - -### E. `routine.py:101,102` — wrong type annotation on `data` field - -```python -@validate_data -def evaluate_data(self, data: dict) -> dict: - ... - new_index = max(self.data.index, default=-1) + 1 - ... - self.data.sort_index(...) -``` -But `self.data` is annotated as `dict[Any, Any]`. The code expects a pandas DataFrame. **Decision needed**: the field type is wrong; should be `DataFrame | None`. Fixing reveals downstream callers expect dict-style access. - -### F. `environment.py` — multiple typing issues - -- **Line 223**: `def add_observation(self, callable=None) -> bool:` — uses bare `callable` (the builtin) as a type. Replace with `Callable[..., Any] | None`. -- **Line 285**: `__call__` declared `-> float` but actually returns `float | list[float]`. Need to widen return type or split into two methods. -- **Lines 298–324**: `Returning Any from function declared to return "dict"` — internal return-type drift. Several functions have explicit `-> dict[Any, Any]` but actually return `Any`. -- Lines 18, 36, 42, 82–83, 154, 188, 258, 304, 316, 327, 331 — missing annotations (mechanical). - -### G. `db.py:109,128,148,149,174,339` — untyped `@db.commit`-style decorators - -`Untyped decorator makes function "save_routine" untyped`. The SQLAlchemy / decorator library lacks type stubs. **Decision needed**: either accept `# type: ignore[misc]` per decorated function, or wait for upstream stubs. - -Plus `db.py:33,38,56,61,79,92,210,303,358,375,388,400,432` — straightforward missing annotations (mechanical). - -### H. `core_subprocess.py` — multiprocessing primitives used as types - -```python -event: multiprocessing.Event # Event is a factory, not a type -queue: multiprocessing.Queue # missing generic param -pipe: multiprocessing.Pipe # Pipe returns tuple, not a type -``` -Lines 37–39, 132–139, 179, 281, 285, 287, 289, 291, 329–349, 359–375. - -**Fix**: import the actual types from `multiprocessing.synchronize` (`Event`), `multiprocessing.connection` (`Connection`), `queue.Queue`. Non-trivial — affects many call sites and signatures across `log.py`, `routine_runner.py`, `create_process.py`. - -### I. `routine_page.py` — multiple structural issues - -`src/badger/gui/components/routine_page.py`: -- Lines 102, 104, 106: function declared `-> str` but returns tuples. Pick one. -- Lines 477, 824: `Incompatible types in assignment (expression has type "list[dict]", variable has type "dict")` — variable reused with mismatched types. -- Lines 512, 545, 584, 861, 897, 940, 1050, 1098, 1378: `Cannot determine type of "configs" / "archive_search"` — class attributes referenced before assignment in init flow. Init method order is fragile. -- Lines 551, 903: `Unpacking a string is disallowed` — code does `a, b, c = some_string` where the runtime value happens to be a tuple sometimes. Genuine bug or unclear intent. -- Lines 1176, 1280–1282: `List item N has incompatible type "float"/"bool"; expected "str"` — heterogeneous list passed where homogeneous expected. -- Lines 1245, 1273, 1308, 1394: `Item "None" of "Any | None"` not iterable / indexable. -- Lines 1903, 1907, 1911: more None-attribute access. - -### J. `settings.py` — `ConfigSingleton` pattern not modeled - -Lines 137, 140, 142, 194, 202, 212, 219, 247, 391, 406, 481, 492, 541, 1683, 1684, 1685, 1686, 1687, 1688, 1689, 1690, 1691, 1694: -- `"ConfigSingleton" has no attribute "user_flag"` etc. — singleton attributes set dynamically. -- `Cannot determine type of "_config"`. -- `Argument 1 to "copytree" has incompatible type "Traversable"; expected "str | PathLike[str]"`. - -**Decision needed**: refactor `ConfigSingleton` with explicit `__slots__` / class-level annotations, or sprinkle `cast()`/`# type: ignore`. - -### K. `logbook.py:60` — `Element` type-var bound - -`Value of type variable "_Tag" of "Element" cannot be "None"` — passing `None` where `xml.etree.ElementTree` expects a string tag. Likely a real bug if `name` is ever None at runtime. - -### L. `gui/__init__.py:85` — `sys.excepthook` signature - -```python -sys.excepthook = my_hook -``` -where `my_hook`'s third parameter is `TracebackType` (non-Optional) but `sys.excepthook` expects `Optional[TracebackType]`. Need to add a None guard inside the hook. - -### M. `pydantic_editor.py` — list/tuple variance and BaseModel typing - -Lines 150, 154, 171, 189, 254 in `src/badger/gui/components/pydantic_editor.py`: -- `Incompatible types in assignment (expression has type "list[Any]", variable has type "tuple[Any, ...]")` — type narrowing breaks because variable is reused. -- `Argument "main" to "BadgerResolvedType" has incompatible type "Any | type[Any] | | None"; expected "type[Any] | None"` — pydantic generic resolution returns more types than annotated. -- `Incompatible types in assignment (expression has type "BadgerResolvedType | None", variable has type "BadgerResolvedType")`. - -These need careful review of the pydantic-resolution machinery. - -### N. `pydantic_editor.py:1110` — redundant cast - -Already addressed in code; the `cast` was redundant but mypy says so as `[redundant-cast]`. Just delete the call. - -### O. Untyped pydantic decorators in `routine.py` - -Lines 52, 143: `@validate_model`, `@validate_data` make decorated methods untyped. Cascade is the same as §G. Wait for upstream pydantic typing fix or add `# type: ignore[misc]`. - -### P. Stale `# type: ignore` comments (mechanical) - -Files with `[unused-ignore]` errors — just delete the comments: -- `src/badger/gui/components/pydantic_editor.py`: lines 668, 674, 931 -- `src/badger/gui/components/plot_event_handlers.py`: lines 57, 76, 81, 318, 319 -- `src/badger/gui/components/pf_viewer/pf_widget.py`: lines 76, 77 -- `src/badger/gui/components/bo_visualizer/bo_widget.py`: lines 51, 52, 120, 129, 241 -- `src/badger/gui/windows/settings_dialog.py`: line 181 - -### Q. `editable_table.py:263` — unbound `ParamSpec` - -`ParamSpec "P" is unbound`. The class uses `ParamSpec` outside a generic context. Needs refactor to either remove the `ParamSpec` or make the surrounding class generic. - -### R. `var_table.py:263,333` — duplicate attribute assignment (`[no-redef]`) - -Class attributes `selected` and `addtl_vars` assigned in two places. Needs single-source declaration in `__init__`. - -### S. Test file logic mismatches - -- `src/badger/tests/test_environment.py:126` — dict literal type mismatch (`Dict entry 0 has incompatible type "str": "list[int]"; expected "str": "float"`). The test is feeding wrong-shape data. -- `src/badger/tests/test_core.py:64` — `Incompatible types in assignment (expression has type "dict[Any, Any]", variable has type "None")` — variable initialized as None then assigned dict; should be `dict | None`. -- `src/badger/tests/x-test_db.py` — file prefixed `x-`, not collected by pytest. Likely dead code; could be deleted. -- `src/badger/tests/utils.py:25,53,94,124,162,164` — `Argument after ** must be a mapping, not "Collection[str]"` — repeatedly unpacks a Collection as kwargs. Actual bug if executed; the file may be unused or only partially exercised. - ---- - -## §3. Suggested follow-up plan - -1. **Mechanical PR for tests** (§1a) — adds ~200 `-> None` annotations. Low risk, high signal; gates future strictness on test code. -2. **Mechanical PR for GUI windows/dialogs** (§1b) — ~120 annotations. -3. **Mechanical PR for GUI components leaf files** (the smaller files in §1c). -4. **Logic PR: callbacks & Optional core types** (§A, §B, §C, §D) — change `self.routine` etc. to proper `Optional` and add guards. -5. **Logic PR: multiprocessing types** (§H) — `multiprocessing.synchronize.Event` etc. across `core_subprocess.py`, `log.py`, `routine_runner.py`. -6. **Logic PR: settings.py + ConfigSingleton** (§J) — restructure or cast. -7. **Logic PR: routine_page.py, run_monitor.py, var_table.py large files** (§B, §I, §R) — biggest, most fragile; do last. - -Each can be a separate PR; none of them block each other except where noted. - ---- - -## §4. Counts at end of this PR - -- **Errors before this PR**: 1506 -- **Errors after this PR**: 1473 (down 33) -- **Errors after PyQt5-stubs install**: jumped to 1621 first (Qt now visible to mypy), then dropped to 1473 after fixing core/library + stubs benefit. Net effect positive vs. baseline despite the stub-revealed new errors. -- **Files fully clean now**: `stats.py`, `errors.py`, `extension.py`, `formula.py`, `interface.py`, `logger/util.py`, `logger/observer.py`, `actions/__init__.py`, `actions/doctor.py`, `actions/run.py`, `actions/routine.py`, `actions/intf.py`, `actions/generator.py`, `actions/env.py`, `actions/config.py`, `actions/install.py`, `actions/uninstall.py`, `__main__.py`, plus partial in `core.py`, `factory.py`, `utils.py`. diff --git a/src/badger/built_in_plugins/interfaces/default/__init__.py b/src/badger/built_in_plugins/interfaces/default/__init__.py index ed2b96ba..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,7 @@ 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 @@ -32,6 +33,6 @@ def get_values(self, channel_names: list[str]) -> dict[str, float | list[float]] return channel_outputs - def set_values(self, channel_inputs: dict[str, float | list[float]]) -> None: + 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/environment.py b/src/badger/environment.py index 1605be0c..23746fe7 100644 --- a/src/badger/environment.py +++ b/src/badger/environment.py @@ -13,7 +13,7 @@ import logging from abc import abstractmethod from collections.abc import Callable -from typing import TYPE_CHECKING, Any, ClassVar, Self +from typing import TYPE_CHECKING, Any, ClassVar, cast from pydantic import BaseModel, ConfigDict, Field, SerializeAsAny from pydantic._internal._model_construction import ModelMetaclass @@ -53,14 +53,16 @@ def validate(cls: Any, variable_inputs: dict[str, float]) -> Any: def process_formulas( - func: Callable[[Any, list[str]], dict[str, float]], -) -> Callable[[Any, list[str]], dict[str, float]]: + 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 = [] @@ -127,9 +129,10 @@ def validate(cls: Any, variable_names: list[str]) -> dict[str, tuple[float, floa class EnvMeta(ModelMetaclass): - def __new__( + # 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] - ) -> Self: + ) -> "EnvMeta": # Wrap get_bounds with validate_bounds if defined if "get_bounds" in namespace: namespace["get_bounds"] = validate_bounds(namespace["get_bounds"]) @@ -318,7 +321,9 @@ 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]) -> None: if not self.interface: diff --git a/src/badger/gui/components/process_manager.py b/src/badger/gui/components/process_manager.py index 5551217b..0da73b98 100644 --- a/src/badger/gui/components/process_manager.py +++ b/src/badger/gui/components/process_manager.py @@ -1,23 +1,30 @@ """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 dataclasses import dataclass from multiprocessing import Process, Queue +from multiprocessing.connection import Connection from multiprocessing.synchronize import Event +from typing import TYPE_CHECKING, Any from PyQt5.QtCore import QObject, pyqtSignal +if TYPE_CHECKING: + from badger.gui.components.routine_runner import ArgumentDict + @dataclass class ProcessWithArgs: process: Process stop_event: Event pause_event: Event - args_queue: Queue - data_queue: Queue - evaluate_queue: Queue + args_queue: Queue[ArgumentDict] + data_queue: Queue[dict[str, Any] | tuple[str, str]] + evaluate_queue: tuple[Connection, Connection] wait_event: Event - dialog_action_queue: Queue + dialog_action_queue: Queue[dict[str, str]] class ProcessManager(QObject): diff --git a/src/badger/gui/components/routine_page.py b/src/badger/gui/components/routine_page.py index 282c7725..a8100e43 100644 --- a/src/badger/gui/components/routine_page.py +++ b/src/badger/gui/components/routine_page.py @@ -64,7 +64,6 @@ QVBoxLayout, QWidget, ) -from xopt import VOCS from xopt.generators import ( all_generator_names, get_generator_defaults, @@ -112,6 +111,7 @@ strtobool, ts_float_to_str, ) +from xopt import VOCS logger = logging.getLogger(__name__) diff --git a/src/badger/gui/components/routine_runner.py b/src/badger/gui/components/routine_runner.py index bcbe8e70..a755e446 100644 --- a/src/badger/gui/components/routine_runner.py +++ b/src/badger/gui/components/routine_runner.py @@ -11,8 +11,9 @@ import time import traceback from dataclasses import dataclass -from multiprocessing import Process +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 @@ -44,14 +45,14 @@ @dataclass class ArgumentDict: - routine_id: int | None + routine_id: str | None routine_filename: str | None routine_name: str - variable_ranges: dict - initial_points: pd.DataFrame + variable_ranges: dict[str, Any] + initial_points: pd.DataFrame | None evaluate: bool archive: bool - termination_condition: dict + termination_condition: dict[str, Any] | None start_time: float testing: bool run_data: bool @@ -103,19 +104,19 @@ 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: dict[str, Any] | None = ( None # additional option to control the optimization flow ) 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 = None + self.data_and_error_queue: Queue[dict[str, Any] | tuple[str, str]] | None = None self.stop_event: Event | None = None self.pause_event: Event | None = None self.routine_process: Process | None = None @@ -124,7 +125,7 @@ def __init__( self.testing = testing self.config_singleton = init_settings() - def set_termination_condition(self, termination_condition: dict) -> None: + def set_termination_condition(self, termination_condition: dict[str, Any]) -> None: logger.info(f"Setting termination condition: {termination_condition}") """ Setter method for the termination condition. @@ -169,8 +170,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, ) @@ -178,16 +186,16 @@ 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() @@ -234,8 +242,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, ) @@ -243,23 +258,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) @@ -281,6 +297,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() @@ -319,7 +341,7 @@ 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: dict[str, Any]) -> str: dialog = BadgerMeasurementRetryDialog( text=msg.get("title", "Measurement failed."), detailedText=msg.get("traceback", ""), @@ -329,24 +351,26 @@ def handle_measurement_error(self, msg: dict) -> str: return MEASUREMENT_ACTION_RETRY return MEASUREMENT_ACTION_ABORT - def handle_termination_reached(self, msg: dict) -> str: + def handle_termination_reached(self, msg: dict[str, Any]) -> str: # update status - tc_condition = msg.get("tc_condition") + tc_condition = cast(dict[str, Any], msg.get("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.get("title", ""), ) result = dialog.exec_() if result == QDialog.Accepted: - self.signals.sig_status.emit(f"Running routine {self.routine.name}...") + 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" @@ -354,7 +378,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.") @@ -373,6 +399,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) @@ -384,6 +412,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) @@ -403,6 +433,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..3232bc5d 100644 --- a/src/badger/gui/components/run_monitor.py +++ b/src/badger/gui/components/run_monitor.py @@ -240,7 +240,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 +251,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. 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/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/interface.py b/src/badger/interface.py index eed21b6f..e6e81c03 100644 --- a/src/badger/interface.py +++ b/src/badger/interface.py @@ -4,7 +4,7 @@ import pickle from abc import ABC, abstractmethod -from collections.abc import Callable +from collections.abc import Callable, Mapping from typing import Any, ClassVar, TypedDict from pydantic import BaseModel @@ -88,7 +88,7 @@ def get_values(self, channel_names: list[str]) -> dict[str, float | list[float]] # Environment should only call this method to set channels @abstractmethod - def set_values(self, channel_inputs: dict[str, float | list[float]]) -> None: + def set_values(self, channel_inputs: Mapping[str, float | list[float]]) -> None: raise NotImplementedError() def reset_interface(self) -> None: diff --git a/src/badger/log.py b/src/badger/log.py index f892902f..9879a236 100644 --- a/src/badger/log.py +++ b/src/badger/log.py @@ -24,6 +24,16 @@ 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, @@ -51,7 +61,7 @@ def start_listener(self, log_filepath: str, log_level: str | int) -> None: # Queue for sending all the logs to self.log_queue = Queue() - self.handlers: list[logging.Handler] = [] + self.handlers = [] # File handler file_handler = logging.FileHandler(log_filepath, mode="a") file_formatter = logging.Formatter( @@ -76,7 +86,7 @@ def start_listener(self, log_filepath: str, log_level: str | int) -> None: ) # '*' unpacks the array for us (unpacking operator) self.listener.start() - log_level_name = logging.getLevelName(log_level) + log_level_name = _get_level_name(log_level) logger.info(f"Centralized logging listener started with level {log_level_name}") @@ -109,8 +119,7 @@ def update_log_level(self, log_level: str | int) -> None: if isinstance(logger_obj, logging.Logger) and name.startswith("badger"): logger_obj.setLevel(log_level) - # deprecated, getLevelNamesMapping()[log_level] is the replacement - log_level_name = logging.getLevelName(log_level) + log_level_name = _get_level_name(log_level) logger.info(f"Log level updated to {log_level_name} for badger namespace.") @@ -209,7 +218,7 @@ def create_log_dir(self, log_dir: str) -> None: def configure_process_logging( - log_queue: Queue[str] | None = None, + log_queue: Queue[logging.LogRecord] | None = None, logger_name: str = "badger", log_level: str | int = "DEBUG", process_name: str | None = None, @@ -250,7 +259,7 @@ def configure_process_logging( if isinstance(logger_obj, logging.Logger) and name.startswith(logger_name): logger_obj.setLevel(log_level) - log_level_name = logging.getLevelName(log_level) + log_level_name = _get_level_name(log_level) logger.info(f"process logger configured with level {log_level_name}") diff --git a/src/badger/routine.py b/src/badger/routine.py index 7cada305..e2a3da66 100644 --- a/src/badger/routine.py +++ b/src/badger/routine.py @@ -14,7 +14,7 @@ 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 @@ -52,13 +52,13 @@ class Routine(Xopt): environment: SerializeAsAny[BaseEnvironment] 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({}) @@ -113,23 +113,20 @@ def validate_model(cls, data: Any) -> 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]) - - df = data["data"] - assert isinstance(df, pd.DataFrame) + df = pd.DataFrame(data["data"], index=[0]) df.index = df.index.astype(int) df.sort_index(inplace=True) - data["data"] = df + 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): @@ -151,7 +148,9 @@ 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 +159,7 @@ def evaluate_point( logger.debug(f"Evaluation result: {obs}") return obs - data["evaluator"] = Evaluator(function=evaluate_point) + data["evaluator"] = Evaluator(function=evaluate_point) # type: ignore[call-arg] return data @@ -205,20 +204,23 @@ def json(self, **kwargs: Any) -> 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: dict[str, Any], vocs: VOCS, env: Environment + limit_options: dict[str, dict[str, Any]], vocs: VOCS, env: BaseEnvironment ) -> dict[str, list[float]]: logger.info("Calculating variable bounds.") vnames = vocs.variable_names @@ -258,7 +260,7 @@ def calculate_variable_bounds( def calculate_initial_points( - init_actions: list[dict[str, Any]], vocs: VOCS, env: Environment + init_actions: list[dict[str, Any]], vocs: VOCS, env: BaseEnvironment ) -> dict[str, list[float]]: logger.info("Calculating initial points.") vnames = vocs.variable_names diff --git a/src/badger/tests/conftest.py b/src/badger/tests/conftest.py index 1fc25c09..63c237f4 100644 --- a/src/badger/tests/conftest.py +++ b/src/badger/tests/conftest.py @@ -1,12 +1,15 @@ 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 @pytest.fixture(autouse=True) -def suppress_popups(mocker): +def suppress_popups(mocker: MockerFixture) -> None: mocker.patch( "badger.gui.windows.expandable_message_box.ExpandableMessageBox.exec_", return_value=None, @@ -19,14 +22,15 @@ 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_plugin_root: str, + 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() @@ -71,12 +75,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 +99,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..a2d8327b 100644 --- a/src/badger/tests/mock/plugins/environments/multiobjective_test/__init__.py +++ b/src/badger/tests/mock/plugins/environments/multiobjective_test/__init__.py @@ -1,9 +1,9 @@ from typing import ClassVar import torch +from badger.errors import BadgerNoInterfaceError from badger import environment -from badger.errors import BadgerNoInterfaceError class Environment(environment.Environment): @@ -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..e8fcb316 100644 --- a/src/badger/tests/mock/plugins/environments/test/__init__.py +++ b/src/badger/tests/mock/plugins/environments/test/__init__.py @@ -2,9 +2,9 @@ from typing import ClassVar import torch +from badger.errors import BadgerNoInterfaceError from badger import environment -from badger.errors import BadgerNoInterfaceError class Environment(environment.Environment): @@ -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..151dd5ce 100644 --- a/src/badger/tests/test_core_subprocess.py +++ b/src/badger/tests/test_core_subprocess.py @@ -1,9 +1,13 @@ 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 ArgumentDict + class TestCore: """ @@ -11,7 +15,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 +26,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 +39,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 @@ -57,7 +61,7 @@ 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 + self, process_manager: ProcessManager, init_multiprocessing: None ) -> None: """ A unit test to ensure the core functionality @@ -77,22 +81,27 @@ def test_run_routine_subprocess( "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 = ArgumentDict( + 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..4b724432 100644 --- a/src/badger/tests/test_create_process.py +++ b/src/badger/tests/test_create_process.py @@ -1,21 +1,28 @@ import multiprocessing as mp import sys +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"), 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..6da7e4d3 100644 --- a/src/badger/tests/test_environment.py +++ b/src/badger/tests/test_environment.py @@ -19,12 +19,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 +39,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 +58,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 +90,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 +109,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 +137,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 +170,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 +204,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 +233,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 +249,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 +269,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 +293,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 +328,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 +350,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 +394,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 From ba282cb41361406a42da3a6bcb9fe5a1de2f3121 Mon Sep 17 00:00:00 2001 From: Mitchell Victoriano Date: Tue, 29 Sep 2026 19:39:12 -0700 Subject: [PATCH 13/15] continued progress with typing --- src/badger/core_subprocess.py | 147 +++++++++++------- src/badger/gui/components/create_process.py | 7 +- src/badger/gui/components/process_manager.py | 21 +-- src/badger/gui/components/routine_page.py | 2 +- src/badger/gui/components/routine_runner.py | 20 +-- src/badger/gui/components/run_monitor.py | 8 +- src/badger/gui/mini/pages/home_page.py | 63 ++++---- src/badger/gui/mini/pages/routine_page.py | 17 +- src/badger/gui/mini/windows/main_window.py | 36 +++-- .../gui/windows/load_data_from_run_dialog.py | 45 +++--- src/badger/gui/windows/main_window.py | 40 ++--- src/badger/gui/windows/message_dialog.py | 35 +++-- src/badger/gui/windows/review_dialog.py | 6 +- src/badger/gui/windows/settings_dialog.py | 2 +- .../windows/terminition_condition_dialog.py | 68 +++++--- src/badger/gui/windows/var_dialog.py | 24 ++- src/badger/routine.py | 7 +- .../multiobjective_test/__init__.py | 2 +- .../plugins/environments/test/__init__.py | 2 +- src/badger/tests/test_core_subprocess.py | 4 +- src/badger/tests/test_factory.py | 2 +- src/badger/tests/test_formulas.py | 84 +++++----- src/badger/tests/test_gui_basic.py | 28 ++-- src/badger/tests/test_home_page.py | 2 +- src/badger/tests/test_routine_runner.py | 4 +- src/badger/tests/test_run_monitor.py | 4 +- src/badger/types.py | 37 +++++ 27 files changed, 411 insertions(+), 306 deletions(-) create mode 100644 src/badger/types.py diff --git a/src/badger/core_subprocess.py b/src/badger/core_subprocess.py index 7e4fd9e5..f53e14ed 100644 --- a/src/badger/core_subprocess.py +++ b/src/badger/core_subprocess.py @@ -18,8 +18,9 @@ import time import traceback from copy import deepcopy +from dataclasses import dataclass from multiprocessing.connection import Connection -from multiprocessing.synchronize import Event as EventType +from multiprocessing.synchronize import Event from queue import Empty from typing import Any @@ -47,15 +48,23 @@ apply_pytorch_multiprocess_tensor_sharing_setting, init_settings, ) +from badger.types import ArgumentQueueType logger = logging.getLogger(__name__) +@dataclass +class ErrorQueueType: + type: str + title: str + traceback: str + + def evaluate_measurement_with_retry( routine: Routine, point: Any, - queue: mp.Queue[Any], - stop_process: EventType, + queue: mp.Queue[ErrorQueueType], + stop_process: Event, dialog_action_queue: mp.Queue[Any], ) -> DataFrame: while True: @@ -66,11 +75,11 @@ def evaluate_measurement_with_retry( 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, - } + ErrorQueueType( + type=MEASUREMENT_ERROR_TYPE, + title=error_title, + traceback=error_traceback, + ) ) # Wait for response back from retry dialog @@ -97,19 +106,26 @@ def evaluate_measurement_with_retry( ) +@dataclass +class TerminationCondition: + type: str + condition: dict[str, Any] + + def pause_for_termination_dialog_action( - queue: mp.Queue, - stop_process: EventType, - pause_process: EventType, - dialog_action_queue: mp.Queue, - tc_condition: dict, + queue: mp.Queue[ErrorQueueType], + stop_process: Event, + pause_process: Event, + dialog_action_queue: mp.Queue[Any], + 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, - } + ErrorQueueType( + type=TERMINATION_REACHED_TYPE, + title="Termination condition reached", + traceback=str(tc_condition.condition), + ) ) while True: @@ -148,6 +164,8 @@ def convert_to_solution(result: DataFrame, routine: Routine) -> Solution: 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) @@ -190,28 +208,30 @@ def convert_to_solution(result: DataFrame, routine: Routine) -> Solution: def run_routine_subprocess( - args_queue: mp.Queue, - queue: mp.Queue, + args_queue: mp.Queue[ArgumentQueueType], + queue: mp.Queue[ErrorQueueType], evaluate_queue: tuple[Connection, Connection], - stop_process: EventType, - pause_process: EventType, - wait_event: EventType, + stop_process: Event, + pause_process: Event, + wait_event: Event, + dialog_action_queue: mp.Queue[dict[str, str]], config_path: str | None = None, - log_queue: mp.Queue | None = None, - dialog_action_queue: mp.Queue[Any] | 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[Any] | None = None + queue: mp.Queue[ErrorQueueType] | None = None + args_queue: mp.Queue[ArgumentQueueType], evaluate_queue: tuple[Connection, Connection], - stop_process: EventType - pause_process: EventType - wait_event: EventType - config_path: str | None = None - log_queue: mp.Queue[Any] | None = None + stop_process: Event, + pause_process: Event, + wait_event: Event, + config_path: str | None = None, + log_queue: mp.Queue[logging.LogRecord] | None = None, + dialog_action_queue: mp.Queue[dict[str, str]] | None = None, """ # Setup logging for this subprocess if log_queue is not None: @@ -240,17 +260,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: @@ -268,7 +288,11 @@ 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)) + queue.put( + ErrorQueueType( + type="InitializationError", title=error_title, traceback=error_traceback + ) + ) raise # TODO look into this bug with serializing of turbo. Fix might be needed in Xopt @@ -285,18 +309,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 @@ -325,12 +349,16 @@ 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 + routine, + ele.to_dict(), + queue, + stop_process, + dialog_action_queue, ) solution = convert_to_solution(result, routine) opt_logger.update(Events.OPTIMIZATION_STEP, solution) @@ -343,6 +371,7 @@ def run_routine_subprocess( 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...") @@ -372,16 +401,20 @@ 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, 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 @@ -400,11 +433,13 @@ def run_routine_subprocess( 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 diff --git a/src/badger/gui/components/create_process.py b/src/badger/gui/components/create_process.py index c62ddf30..f4965dff 100644 --- a/src/badger/gui/components/create_process.py +++ b/src/badger/gui/components/create_process.py @@ -10,6 +10,7 @@ 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 logger = logging.getLogger(__name__) @@ -33,11 +34,11 @@ 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[dict[str, Any] | tuple[str, str]] = Queue() self.evaluate_queue = Pipe() self.wait_event = Event() - self.dialog_action_queue: Queue[Any] = Queue() + self.dialog_action_queue: Queue[dict[str, str]] = Queue() config_instance = init_settings()._instance config_path = config_instance.config_path if config_instance else None diff --git a/src/badger/gui/components/process_manager.py b/src/badger/gui/components/process_manager.py index 0da73b98..551e05d6 100644 --- a/src/badger/gui/components/process_manager.py +++ b/src/badger/gui/components/process_manager.py @@ -3,28 +3,9 @@ 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 TYPE_CHECKING, Any - from PyQt5.QtCore import QObject, pyqtSignal -if TYPE_CHECKING: - from badger.gui.components.routine_runner import ArgumentDict - - -@dataclass -class ProcessWithArgs: - process: Process - stop_event: Event - pause_event: Event - args_queue: Queue[ArgumentDict] - data_queue: Queue[dict[str, Any] | tuple[str, str]] - evaluate_queue: tuple[Connection, Connection] - wait_event: Event - dialog_action_queue: Queue[dict[str, str]] +from badger.types import ProcessWithArgs class ProcessManager(QObject): diff --git a/src/badger/gui/components/routine_page.py b/src/badger/gui/components/routine_page.py index a8100e43..282c7725 100644 --- a/src/badger/gui/components/routine_page.py +++ b/src/badger/gui/components/routine_page.py @@ -64,6 +64,7 @@ QVBoxLayout, QWidget, ) +from xopt import VOCS from xopt.generators import ( all_generator_names, get_generator_defaults, @@ -111,7 +112,6 @@ strtobool, ts_float_to_str, ) -from xopt import VOCS logger = logging.getLogger(__name__) diff --git a/src/badger/gui/components/routine_runner.py b/src/badger/gui/components/routine_runner.py index a755e446..0935426b 100644 --- a/src/badger/gui/components/routine_runner.py +++ b/src/badger/gui/components/routine_runner.py @@ -10,7 +10,6 @@ import logging import time import traceback -from dataclasses import dataclass from multiprocessing import Process, Queue from multiprocessing.synchronize import Event from typing import Any, cast @@ -39,26 +38,11 @@ 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 logger = logging.getLogger(__name__) -@dataclass -class ArgumentDict: - routine_id: str | None - routine_filename: str | None - routine_name: str - variable_ranges: dict[str, Any] - initial_points: pd.DataFrame | None - evaluate: bool - archive: bool - termination_condition: dict[str, Any] | None - start_time: float - testing: bool - run_data: bool - init_points: bool - - class BadgerRoutineSignals(QObject): env_ready = pyqtSignal(list) finished = pyqtSignal() @@ -213,7 +197,7 @@ def run(self, run_data_flag: bool = False, init_points_flag: bool = False) -> No self.wait_event = process_with_args.wait_event self.dialog_action_queue = process_with_args.dialog_action_queue - arg_dict = ArgumentDict( + arg_dict = ArgumentQueueType( routine_id=self.routine.id, routine_filename=self.routine_filename, routine_name=self.routine.name, diff --git a/src/badger/gui/components/run_monitor.py b/src/badger/gui/components/run_monitor.py index 3232bc5d..83334ab0 100644 --- a/src/badger/gui/components/run_monitor.py +++ b/src/badger/gui/components/run_monitor.py @@ -39,6 +39,7 @@ from badger.gui.components.pydantic_editor import BadgerPydanticEditor from badger.gui.components.routine_runner import BadgerRoutineSubprocess from badger.gui.windows.message_dialog import BadgerScrollableMessageBox +from badger.gui.windows.terminition_condition_dialog import TerminationConditionConfig # from ...utils import AURORA_PALETTE, FROST_PALETTE from badger.logbook import BADGER_LOGBOOK_ROOT, send_to_logbook @@ -488,7 +489,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: @@ -659,6 +660,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"] @@ -693,6 +697,8 @@ def routine_finished(self) -> None: # self.reset_routine_runner() def destroy_unused_env(self) -> None: + if self.routine is None: + raise ValueError("Routine is not set") if not self.running: try: del self.routine_runner.routine.environment diff --git a/src/badger/gui/mini/pages/home_page.py b/src/badger/gui/mini/pages/home_page.py index 350b9d06..eb6ecbd0 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 @@ -76,7 +78,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 +92,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 +121,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 +151,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 +177,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 +214,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,7 +293,7 @@ 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 @@ -301,7 +303,7 @@ def _configure_default_run_action(self): 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 +314,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 +327,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 +384,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 +397,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 +428,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 +437,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 @@ -489,7 +491,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 +564,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 +611,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 +641,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 +659,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 +668,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 +687,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 +708,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 +722,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 2d163151..665359e6 100644 --- a/src/badger/gui/mini/pages/routine_page.py +++ b/src/badger/gui/mini/pages/routine_page.py @@ -22,6 +22,7 @@ import yaml from coolname import generate_slug from gest_api.vocs import ( + BaseConstraint, BaseObjective, ContinuousVariable, GreaterThanConstraint, @@ -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() @@ -246,7 +249,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 +306,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 +347,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 +671,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 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/windows/load_data_from_run_dialog.py b/src/badger/gui/windows/load_data_from_run_dialog.py index 28fd064a..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,15 +58,15 @@ class BadgerLoadDataFromRunDialog(QDialog): def __init__( self, parent: QWidget, - env_vocs: VOCS | None = 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. - env_vocs (VOCS | None, optional): The environment VOCS object containing variable and objective names. + 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: @@ -73,7 +74,7 @@ def __init__( 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 @@ -147,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. """ @@ -159,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, @@ -234,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 @@ -302,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 = {} @@ -345,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. @@ -379,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. @@ -398,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/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..24a1e0b7 100644 --- a/src/badger/gui/windows/review_dialog.py +++ b/src/badger/gui/windows/review_dialog.py @@ -9,7 +9,7 @@ class BadgerReviewDialog(QDialog): - def __init__(self, parent, routine: Routine): + def __init__(self, routine: Routine, parent: QDialog | None = None): super().__init__(parent) self.routine = routine @@ -17,7 +17,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 +31,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 ab2f95bf..5ff25738 100644 --- a/src/badger/gui/windows/settings_dialog.py +++ b/src/badger/gui/windows/settings_dialog.py @@ -40,7 +40,7 @@ class BadgerSettingsDialog(QDialog): 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() diff --git a/src/badger/gui/windows/terminition_condition_dialog.py b/src/badger/gui/windows/terminition_condition_dialog.py index d0e8e640..93cb6e66 100644 --- a/src/badger/gui/windows/terminition_condition_dialog.py +++ b/src/badger/gui/windows/terminition_condition_dialog.py @@ -2,6 +2,10 @@ 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 dataclasses import dataclass + +from PyQt5.QtGui import QCloseEvent from PyQt5.QtWidgets import ( QComboBox, QDialog, @@ -34,20 +38,37 @@ """ +@dataclass +class TerminationConditionConfig: + tc_idx: int + max_eval: int + max_time: float + ftol: float + + 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 +90,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 +108,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 +120,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 +133,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 +143,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 +164,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/routine.py b/src/badger/routine.py index e2a3da66..5e256d50 100644 --- a/src/badger/routine.py +++ b/src/badger/routine.py @@ -23,6 +23,7 @@ # 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,6 +33,7 @@ field_validator, model_validator, ) +from xopt import Evaluator, Xopt from xopt.generators import get_generator from xopt.generators.sequential import SequentialGenerator from xopt.vocs import get_local_region @@ -39,7 +41,6 @@ from badger.environment import BaseEnvironment, Environment, instantiate_env from badger.factory import get_env from badger.utils import curr_ts -from xopt import VOCS, Evaluator, Xopt logger = logging.getLogger(__name__) @@ -148,9 +149,7 @@ def evaluate_point( point: dict[str, float], ) -> dict[str, float | list[float]]: logger.debug(f"Evaluating point: {point}") - point = cast( - dict[str, float], 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() 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 a2d8327b..d648081c 100644 --- a/src/badger/tests/mock/plugins/environments/multiobjective_test/__init__.py +++ b/src/badger/tests/mock/plugins/environments/multiobjective_test/__init__.py @@ -1,9 +1,9 @@ from typing import ClassVar import torch -from badger.errors import BadgerNoInterfaceError from badger import environment +from badger.errors import BadgerNoInterfaceError class Environment(environment.Environment): diff --git a/src/badger/tests/mock/plugins/environments/test/__init__.py b/src/badger/tests/mock/plugins/environments/test/__init__.py index e8fcb316..2db19d4f 100644 --- a/src/badger/tests/mock/plugins/environments/test/__init__.py +++ b/src/badger/tests/mock/plugins/environments/test/__init__.py @@ -2,9 +2,9 @@ from typing import ClassVar import torch -from badger.errors import BadgerNoInterfaceError from badger import environment +from badger.errors import BadgerNoInterfaceError class Environment(environment.Environment): diff --git a/src/badger/tests/test_core_subprocess.py b/src/badger/tests/test_core_subprocess.py index 151dd5ce..407c62a0 100644 --- a/src/badger/tests/test_core_subprocess.py +++ b/src/badger/tests/test_core_subprocess.py @@ -6,7 +6,7 @@ import pytest from badger.gui.components.process_manager import ProcessManager -from badger.gui.components.routine_runner import ArgumentDict +from badger.gui.components.routine_runner import ArgumentQueueType class TestCore: @@ -88,7 +88,7 @@ def test_run_routine_subprocess( routine_process = process_with_args.process evaluate_queue = process_with_args.evaluate_queue - arg_dict = ArgumentDict( + arg_dict = ArgumentQueueType( routine_id=self.routine.id, routine_filename=tmp_filename, routine_name=self.routine.name, 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 106fd9f4..b95017f7 100644 --- a/src/badger/tests/test_gui_basic.py +++ b/src/badger/tests/test_gui_basic.py @@ -13,7 +13,7 @@ def init_multiprocessing() -> None: multiprocessing.set_start_method("fork", force=True) -def test_gui_main(qtbot: QtBot, init_multiprocessing): +def test_gui_main(qtbot: QtBot, init_multiprocessing) -> None: from badger.gui.windows.main_window import BadgerMainWindow from badger.tests.utils import fix_path_issues @@ -37,10 +37,10 @@ def test_gui_main(qtbot: QtBot, init_multiprocessing): # TODO: Test if generator tab has been filled # assert window.stacks.currentWidget().tabs.currentIndex() == 1 - window.process_manager.close_proccesses() + window.process_manager.close_processes() -def test_close_main(qtbot, init_multiprocessing): +def test_close_main(qtbot, init_multiprocessing) -> 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 @@ -77,10 +77,10 @@ def test_close_main(qtbot, init_multiprocessing): with pytest.raises(AttributeError): _ = home_page.run_monitor.routine.environment - window.process_manager.close_proccesses() + window.process_manager.close_processes() -def test_traceback_during_run(qtbot: QtBot) -> None: +def test_traceback_during_run(qtbot: QtBot, init_multiprocessing) -> None: with patch("badger.core.run_routine") as run_routine_mock: run_routine_mock.side_effect = Exception("Test exception") @@ -129,7 +129,7 @@ def inner(ins, event): while home_page.run_monitor.running: qtbot.wait(100) - window.process_manager.close_proccesses() + window.process_manager.close_processes() def test_measurement_retry_dialog_in_app_flow(qtbot: QtBot) -> None: @@ -154,10 +154,10 @@ def test_measurement_retry_dialog_in_app_flow(qtbot: QtBot) -> 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"] + 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_", @@ -187,9 +187,9 @@ def test_measurement_retry_dialog_in_app_flow(qtbot: QtBot) -> None: "action": "retry", } - window.process_manager.close_proccesses() - process_with_args["process"].terminate() - process_with_args["process"].join() + window.process_manager.close_processes() + process_with_args.process.terminate() + process_with_args.process.join() # TODO: Check the use_low_noise_prior parameter in the routine @@ -232,7 +232,7 @@ def test_default_low_noise_prior_in_bo(qtbot: QtBot) -> None: else: # that part of params is hidden so we need to dig deeper pass - window.process_manager.close_proccesses() + window.process_manager.close_processes() # TODO: Fix default turbo in bo test diff --git a/src/badger/tests/test_home_page.py b/src/badger/tests/test_home_page.py index b5bd6192..8ca3301a 100644 --- a/src/badger/tests/test_home_page.py +++ b/src/badger/tests/test_home_page.py @@ -67,4 +67,4 @@ def test_home_page_run_routine(qtbot, init_multiprocessing): # 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() + main_page.process_manager.close_processes() diff --git a/src/badger/tests/test_routine_runner.py b/src/badger/tests/test_routine_runner.py index f4b3fe0b..9fd8f207 100644 --- a/src/badger/tests/test_routine_runner.py +++ b/src/badger/tests/test_routine_runner.py @@ -30,7 +30,7 @@ def process_manager(self): yield process_manager - process_manager.close_proccesses() + process_manager.close_processes() @pytest.fixture def instance(self, process_manager, init_multiprocessing): @@ -170,4 +170,4 @@ def inner(ins, event): assert len(monitor.routine.data) == 2 - window.process_manager.close_proccesses() + window.process_manager.close_processes() diff --git a/src/badger/tests/test_run_monitor.py b/src/badger/tests/test_run_monitor.py index fd220c6d..dfee1d2d 100644 --- a/src/badger/tests/test_run_monitor.py +++ b/src/badger/tests/test_run_monitor.py @@ -32,7 +32,7 @@ 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): @@ -400,7 +400,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/types.py b/src/badger/types.py new file mode 100644 index 00000000..2e07a6b3 --- /dev/null +++ b/src/badger/types.py @@ -0,0 +1,37 @@ +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 + + +@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: dict[str, Any] | 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[dict[str, Any] | tuple[str, str]] + evaluate_queue: tuple[Connection, Connection] + wait_event: Event + dialog_action_queue: Queue[dict[str, str]] From 032c49f8217f7fe5b3287082b1003646edb03d0c Mon Sep 17 00:00:00 2001 From: Mitchell Victoriano Date: Wed, 30 Sep 2026 16:23:46 -0700 Subject: [PATCH 14/15] Refactor tests and enhance type annotations with new dataclasses --- src/badger/core_subprocess.py | 138 +++++++++--------- src/badger/db.py | 22 +-- src/badger/errors.py | 3 + src/badger/gui/__init__.py | 2 +- src/badger/gui/components/action_bar.py | 16 +- src/badger/gui/components/create_process.py | 35 +++-- src/badger/gui/components/routine_page.py | 2 +- src/badger/gui/components/routine_runner.py | 60 ++++---- src/badger/gui/components/run_monitor.py | 62 +++++--- src/badger/gui/mini/pages/home_page.py | 5 +- src/badger/gui/mini/pages/routine_page.py | 4 +- src/badger/gui/pages/home_page.py | 32 ++-- .../windows/terminition_condition_dialog.py | 11 +- src/badger/routine.py | 5 - src/badger/tests/test_core_subprocess.py | 14 +- src/badger/tests/test_create_process.py | 19 +-- src/badger/tests/test_environment.py | 3 +- src/badger/tests/test_gui_basic.py | 76 ++++++---- src/badger/tests/test_home_page.py | 15 +- src/badger/tests/test_intf.py | 12 +- src/badger/tests/test_lib_basic.py | 4 +- .../tests/test_measurement_retry_dialog.py | 7 +- src/badger/tests/test_process_manager.py | 11 +- src/badger/tests/test_routine.py | 9 +- src/badger/tests/test_routine_page.py | 28 ++-- src/badger/tests/test_routine_runner.py | 102 +++---------- src/badger/tests/test_run_monitor.py | 86 ++++++----- src/badger/tests/test_settings.py | 40 ++--- src/badger/tests/utils.py | 2 +- src/badger/tests/x_test_db.py | 10 +- src/badger/types.py | 98 ++++++++++++- src/badger/utils.py | 2 +- 32 files changed, 514 insertions(+), 421 deletions(-) diff --git a/src/badger/core_subprocess.py b/src/badger/core_subprocess.py index f53e14ed..7d2c9000 100644 --- a/src/badger/core_subprocess.py +++ b/src/badger/core_subprocess.py @@ -31,12 +31,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, BadgerEnvObsError, BadgerRunTerminated, ) @@ -48,24 +44,26 @@ apply_pytorch_multiprocess_tensor_sharing_setting, init_settings, ) -from badger.types import ArgumentQueueType +from badger.types import ( + ArgumentQueueType, + DataQueueMessage, + DialogActionMessage, + MeasurementActionMessage, + MeasurementErrorMessage, + RoutineErrorMessage, + TerminationActionMessage, + TerminationReachedMessage, +) logger = logging.getLogger(__name__) -@dataclass -class ErrorQueueType: - type: str - title: str - traceback: str - - def evaluate_measurement_with_retry( routine: Routine, point: Any, - queue: mp.Queue[ErrorQueueType], + data_and_error_queue: mp.Queue[DataQueueMessage], stop_process: Event, - dialog_action_queue: mp.Queue[Any], + dialog_action_queue: mp.Queue[DialogActionMessage], ) -> DataFrame: while True: try: @@ -74,9 +72,8 @@ 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( - ErrorQueueType( - type=MEASUREMENT_ERROR_TYPE, + data_and_error_queue.put( + MeasurementErrorMessage( title=error_title, traceback=error_traceback, ) @@ -93,13 +90,11 @@ 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}" @@ -113,18 +108,16 @@ class TerminationCondition: def pause_for_termination_dialog_action( - queue: mp.Queue[ErrorQueueType], + data_and_error_queue: mp.Queue[DataQueueMessage], stop_process: Event, pause_process: Event, - dialog_action_queue: mp.Queue[Any], + 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( - ErrorQueueType( - type=TERMINATION_REACHED_TYPE, - title="Termination condition reached", - traceback=str(tc_condition.condition), + data_and_error_queue.put( + TerminationReachedMessage( + tc_condition={"type": tc_condition.type, **tc_condition.condition}, ) ) @@ -139,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 @@ -209,12 +200,12 @@ def convert_to_solution(result: DataFrame, routine: Routine) -> Solution: def run_routine_subprocess( args_queue: mp.Queue[ArgumentQueueType], - queue: mp.Queue[ErrorQueueType], + 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[dict[str, str]], + dialog_action_queue: mp.Queue[DialogActionMessage], config_path: str | None = None, log_queue: mp.Queue[logging.LogRecord] | None = None, ) -> None: @@ -223,15 +214,15 @@ def run_routine_subprocess( Parameters ---------- - queue: mp.Queue[ErrorQueueType] | None = None 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, - dialog_action_queue: mp.Queue[dict[str, str]] | None = None, """ # Setup logging for this subprocess if log_queue is not None: @@ -280,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 @@ -288,10 +279,8 @@ 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( - ErrorQueueType( - type="InitializationError", title=error_title, traceback=error_traceback - ) + data_and_error_queue.put( + RoutineErrorMessage(title=error_title, traceback=error_traceback) ) raise @@ -351,20 +340,23 @@ def run_routine_subprocess( # initial sampling 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: @@ -379,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 @@ -410,7 +402,7 @@ def run_routine_subprocess( ) 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, @@ -420,7 +412,7 @@ def run_routine_subprocess( 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: @@ -429,7 +421,7 @@ 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, @@ -460,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) @@ -483,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: @@ -493,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 b47643ae..ac20fb45 100644 --- a/src/badger/db.py +++ b/src/badger/db.py @@ -95,7 +95,7 @@ def func_safe(*args: P.args, **kwargs: P.kwargs) -> T: 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: @@ -108,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: @@ -226,7 +226,7 @@ def load_routine(id: str) -> tuple[Routine, Any]: @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") @@ -318,7 +318,7 @@ def list_routine(keyword="", tags: dict[str, str] | None = None): @ensure_runs_db_exists -def save_run(run: dict[str, Any]) -> int: +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) @@ -335,6 +335,8 @@ def save_run(run: dict[str, Any]) -> int: 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 = ?", @@ -355,7 +357,7 @@ def save_run(run: dict[str, Any]) -> int: @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) @@ -373,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) @@ -390,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) @@ -403,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) @@ -415,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() @@ -447,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/errors.py b/src/badger/errors.py index 93967072..d05408cb 100644 --- a/src/badger/errors.py +++ b/src/badger/errors.py @@ -132,3 +132,6 @@ def __init__(self, message: str = "Optimization run has been terminated!") -> No 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/gui/__init__.py b/src/badger/gui/__init__.py index 6eb7df13..c5f25a91 100644 --- a/src/badger/gui/__init__.py +++ b/src/badger/gui/__init__.py @@ -60,7 +60,7 @@ def on_timeout() -> None: 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. diff --git a/src/badger/gui/components/action_bar.py b/src/badger/gui/components/action_bar.py index d3eee3a1..2fd50353 100644 --- a/src/badger/gui/components/action_bar.py +++ b/src/badger/gui/components/action_bar.py @@ -2,7 +2,7 @@ docs access, and the extensions palette launcher.""" from importlib import resources -from typing import Any, cast +from typing import cast from PyQt5.QtCore import QEvent, QPoint, QSize, pyqtSignal from PyQt5.QtGui import QFont, QHelpEvent, QIcon @@ -18,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): @@ -467,15 +468,14 @@ def env_ready(self) -> None: self.btn_log.setDisabled(False) self.btn_opt.setDisabled(False) - def update_run_tooltip(self, tc: dict[str, Any] | None = None) -> 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/create_process.py b/src/badger/gui/components/create_process.py index f4965dff..325a2319 100644 --- a/src/badger/gui/components/create_process.py +++ b/src/badger/gui/components/create_process.py @@ -3,14 +3,18 @@ import logging from multiprocessing import Event, Pipe, Process, Queue -from typing import Any from PyQt5.QtCore import QObject, pyqtSignal 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 +from badger.types import ( + ArgumentQueueType, + DataQueueMessage, + DialogActionMessage, + ProcessWithArgs, +) logger = logging.getLogger(__name__) @@ -35,10 +39,10 @@ def create_subprocess(self) -> None: self.stop_event = Event() self.pause_event = Event() self.args_queue: Queue[ArgumentQueueType] = Queue() - self.data_queue: Queue[dict[str, Any] | tuple[str, str]] = Queue() + self.data_queue: Queue[DataQueueMessage] = Queue() self.evaluate_queue = Pipe() self.wait_event = Event() - self.dialog_action_queue: Queue[dict[str, str]] = Queue() + self.dialog_action_queue: Queue[DialogActionMessage] = Queue() config_instance = init_settings()._instance config_path = config_instance.config_path if config_instance else None @@ -63,16 +67,17 @@ def create_subprocess(self) -> None: ), ) 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/routine_page.py b/src/badger/gui/components/routine_page.py index 282c7725..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, diff --git a/src/badger/gui/components/routine_runner.py b/src/badger/gui/components/routine_runner.py index 0935426b..a9315f47 100644 --- a/src/badger/gui/components/routine_runner.py +++ b/src/badger/gui/components/routine_runner.py @@ -21,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, ) @@ -38,7 +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 +from badger.types import ( + ArgumentQueueType, + DataQueueMessage, + MeasurementActionMessage, + MeasurementErrorMessage, + TerminationActionMessage, + TerminationConditionConfig, + TerminationReachedMessage, +) logger = logging.getLogger(__name__) @@ -93,14 +97,14 @@ def __init__( self.save = save self.verbose = verbose self.use_full_ts = use_full_ts - self.termination_condition: dict[str, Any] | None = ( + self.termination_condition: TerminationConditionConfig | None = ( None # additional option to control the optimization flow ) 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[dict[str, Any] | tuple[str, str]] | None = None + 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 @@ -109,14 +113,16 @@ def __init__( self.testing = testing self.config_singleton = init_settings() - def set_termination_condition(self, termination_condition: dict[str, Any]) -> 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 @@ -296,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 @@ -325,29 +321,29 @@ def check_queue(self) -> None: self.close() self.evaluate_queue[1].close() - def handle_measurement_error(self, msg: dict[str, Any]) -> 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, Any]) -> str: + def handle_termination_reached(self, msg: TerminationReachedMessage) -> str: # update status - tc_condition = cast(dict[str, Any], 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: + 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 diff --git a/src/badger/gui/components/run_monitor.py b/src/badger/gui/components/run_monitor.py index 83334ab0..b86e97c5 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,23 +33,21 @@ 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 -from badger.gui.windows.terminition_condition_dialog import TerminationConditionConfig # from ...utils import AURORA_PALETTE, FROST_PALETTE 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__) @@ -74,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 @@ -98,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 @@ -107,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() @@ -115,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 @@ -162,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) @@ -407,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)] @@ -453,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) @@ -520,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 @@ -697,9 +700,10 @@ def routine_finished(self) -> None: # self.reset_routine_runner() def destroy_unused_env(self) -> None: - if self.routine is None: - raise ValueError("Routine is not set") 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 @@ -1079,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") @@ -1106,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/mini/pages/home_page.py b/src/badger/gui/mini/pages/home_page.py index eb6ecbd0..eddae1fd 100644 --- a/src/badger/gui/mini/pages/home_page.py +++ b/src/badger/gui/mini/pages/home_page.py @@ -54,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__) @@ -299,7 +300,9 @@ def _configure_default_run_action(self) -> None: 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) diff --git a/src/badger/gui/mini/pages/routine_page.py b/src/badger/gui/mini/pages/routine_page.py index 665359e6..f01cf2e9 100644 --- a/src/badger/gui/mini/pages/routine_page.py +++ b/src/badger/gui/mini/pages/routine_page.py @@ -22,6 +22,7 @@ import yaml from coolname import generate_slug from gest_api.vocs import ( + VOCS, BaseConstraint, BaseObjective, ContinuousVariable, @@ -45,7 +46,6 @@ QVBoxLayout, QWidget, ) -from xopt import VOCS from xopt.generators import ( all_generator_names, get_generator_defaults, @@ -706,7 +706,7 @@ def save_template_yaml(self) -> None: 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})" ) diff --git a/src/badger/gui/pages/home_page.py b/src/badger/gui/pages/home_page.py index 0679dcfd..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) @@ -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/windows/terminition_condition_dialog.py b/src/badger/gui/windows/terminition_condition_dialog.py index 93cb6e66..750d19b9 100644 --- a/src/badger/gui/windows/terminition_condition_dialog.py +++ b/src/badger/gui/windows/terminition_condition_dialog.py @@ -3,7 +3,6 @@ number of evaluations or after a maximum elapsed time.""" from collections.abc import Callable -from dataclasses import dataclass from PyQt5.QtGui import QCloseEvent from PyQt5.QtWidgets import ( @@ -21,6 +20,8 @@ QWidget, ) +from badger.types import TerminationConditionConfig + stylesheet_run = """ QPushButton:hover:pressed { @@ -38,14 +39,6 @@ """ -@dataclass -class TerminationConditionConfig: - tc_idx: int - max_eval: int - max_time: float - ftol: float - - class BadgerTerminationConditionDialog(QDialog): def __init__( self, diff --git a/src/badger/routine.py b/src/badger/routine.py index 5e256d50..a8d7d6f6 100644 --- a/src/badger/routine.py +++ b/src/badger/routine.py @@ -18,11 +18,6 @@ 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 ( diff --git a/src/badger/tests/test_core_subprocess.py b/src/badger/tests/test_core_subprocess.py index 407c62a0..db53cc78 100644 --- a/src/badger/tests/test_core_subprocess.py +++ b/src/badger/tests/test_core_subprocess.py @@ -7,6 +7,7 @@ from badger.gui.components.process_manager import ProcessManager from badger.gui.components.routine_runner import ArgumentQueueType +from badger.types import TerminationConditionConfig class TestCore: @@ -60,9 +61,8 @@ def test_core_setup(self, *args: object, **kwargs: object) -> None: self.points_eval_target = pd.DataFrame(data_eval_target) - def test_run_routine_subprocess( - self, process_manager: ProcessManager, init_multiprocessing: None - ) -> 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. @@ -76,10 +76,10 @@ 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() assert process_with_args is not None pause_event = process_with_args.pause_event diff --git a/src/badger/tests/test_create_process.py b/src/badger/tests/test_create_process.py index 4b724432..e2118d0e 100644 --- a/src/badger/tests/test_create_process.py +++ b/src/badger/tests/test_create_process.py @@ -1,5 +1,6 @@ import multiprocessing as mp import sys +from dataclasses import fields from typing import TYPE_CHECKING from unittest.mock import MagicMock, patch @@ -41,9 +42,9 @@ def test_create_subprocess_emits_signals( 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", @@ -54,10 +55,10 @@ def test_create_subprocess_emits_signals( "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_environment.py b/src/badger/tests/test_environment.py index 6da7e4d3..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 @@ -413,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_gui_basic.py b/src/badger/tests/test_gui_basic.py index b95017f7..a5f3c574 100644 --- a/src/badger/tests/test_gui_basic.py +++ b/src/badger/tests/test_gui_basic.py @@ -4,6 +4,7 @@ import pytest from PyQt5.QtCore import QEventLoop, QTimer +from PyQt5.QtGui import QShowEvent from PyQt5.QtWidgets import QDialog from pytestqt.qtbot import QtBot @@ -13,7 +14,8 @@ def init_multiprocessing() -> None: multiprocessing.set_start_method("fork", force=True) -def test_gui_main(qtbot: QtBot, init_multiprocessing) -> None: +@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 @@ -40,10 +42,12 @@ def test_gui_main(qtbot: QtBot, init_multiprocessing) -> None: window.process_manager.close_processes() -def test_close_main(qtbot, init_multiprocessing) -> None: +@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() @@ -61,10 +65,10 @@ def test_close_main(qtbot, init_multiprocessing) -> None: 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.run_monitor.termination_condition = TerminationConditionConfig( + tc_idx=0, + max_eval=3, + ) home_page.go_run(-1) home_page.run_monitor.start(True) @@ -73,6 +77,9 @@ def test_close_main(qtbot, init_multiprocessing) -> None: loop.exec_() window.close() # this action should release the env + + assert home_page.run_monitor.routine is not None + # So we expect an AttributeError here with pytest.raises(AttributeError): _ = home_page.run_monitor.routine.environment @@ -80,13 +87,15 @@ def test_close_main(qtbot, init_multiprocessing) -> None: window.process_manager.close_processes() -def test_traceback_during_run(qtbot: QtBot, init_multiprocessing) -> None: +@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() @@ -102,27 +111,26 @@ def test_traceback_during_run(qtbot: QtBot, init_multiprocessing) -> None: 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.run_monitor.termination_condition = TerminationConditionConfig( + 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 + # Patch showEvent to auto-accept the traceback dialog when it appears. + original_show_event = BadgerScrollableMessageBox.showEvent - assert ins # make sure the dialog is created - assert ins.detailedTextWidget.toPlainText() # make sure it's not empty + def patched_show_event( + ins: BadgerScrollableMessageBox, event: QShowEvent | None + ) -> None: + original_show_event(ins, event) # Call the original showEvent - QTimer.singleShot(100, ins.accept) # Close the dialog after 100 ms + assert ins # make sure the dialog is created + assert ins.detailedTextWidget.toPlainText() # make sure it's not empty - return inner + QTimer.singleShot(100, ins.accept) # Close the dialog after 100 ms - BadgerScrollableMessageBox.showEvent = patched_showEvent( - BadgerScrollableMessageBox.showEvent - ) + monkeypatch.setattr(BadgerScrollableMessageBox, "showEvent", patched_show_event) home_page.run_monitor.start(True) # Wait until the run is done @@ -135,6 +143,11 @@ def inner(ins, event): 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() @@ -147,9 +160,14 @@ def test_measurement_retry_dialog_in_app_flow(qtbot: QtBot) -> None: 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() @@ -164,11 +182,10 @@ def test_measurement_retry_dialog_in_app_flow(qtbot: QtBot) -> None: return_value=QDialog.Accepted, ) as exec_mock: runner.data_and_error_queue.put( - { - "type": "measurement_error", - "title": "Injected measurement error", - "traceback": "traceback details", - } + MeasurementErrorMessage( + title="Injected measurement error", + traceback="traceback details", + ) ) # Wait up to 1 second in 10 for queue to receive the error. @@ -182,10 +199,7 @@ def test_measurement_retry_dialog_in_app_flow(qtbot: QtBot) -> None: exec_mock.assert_called_once() response = runner.dialog_action_queue.get(timeout=1) - assert response == { - "type": "measurement_action", - "action": "retry", - } + assert response == MeasurementActionMessage(action="retry") window.process_manager.close_processes() process_with_args.process.terminate() diff --git a/src/badger/tests/test_home_page.py b/src/badger/tests/test_home_page.py index 8ca3301a..6221676a 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,6 +25,7 @@ def test_home_page_run_routine(qtbot, init_multiprocessing): create_routine, fix_path_issues, ) + from badger.types import TerminationConditionConfig fix_path_issues() @@ -45,10 +48,10 @@ def test_home_page_run_routine(qtbot, init_multiprocessing): 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.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) 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..63fa8644 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 @@ -269,28 +269,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 9fd8f207..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() @@ -33,7 +35,7 @@ def process_manager(self): 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_processes() + assert instance.termination_condition diff --git a/src/badger/tests/test_run_monitor.py b/src/badger/tests/test_run_monitor.py index dfee1d2d..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() @@ -35,9 +43,10 @@ def process_manager(self): 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) 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 2db4a906..8377a4c5 100644 --- a/src/badger/tests/utils.py +++ b/src/badger/tests/utils.py @@ -3,7 +3,7 @@ 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 diff --git a/src/badger/tests/x_test_db.py b/src/badger/tests/x_test_db.py index 65067db7..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) -> 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/types.py b/src/badger/types.py index 2e07a6b3..9eeef765 100644 --- a/src/badger/types.py +++ b/src/badger/types.py @@ -8,6 +8,98 @@ 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: @@ -18,7 +110,7 @@ class ArgumentQueueType: initial_points: pd.DataFrame | None evaluate: bool archive: bool - termination_condition: dict[str, Any] | None + termination_condition: TerminationConditionConfig | None start_time: float testing: bool run_data: bool @@ -31,7 +123,7 @@ class ProcessWithArgs: stop_event: Event pause_event: Event args_queue: Queue[ArgumentQueueType] - data_queue: Queue[dict[str, Any] | tuple[str, str]] + data_queue: Queue[DataQueueMessage] evaluate_queue: tuple[Connection, Connection] wait_event: Event - dialog_action_queue: Queue[dict[str, str]] + dialog_action_queue: Queue[DialogActionMessage] diff --git a/src/badger/utils.py b/src/badger/utils.py index b5e44342..9bea5ccb 100644 --- a/src/badger/utils.py +++ b/src/badger/utils.py @@ -469,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 From e6c29b14cb6a3066a067e3649baadedf799a0ef1 Mon Sep 17 00:00:00 2001 From: Mitchell Victoriano Date: Wed, 30 Sep 2026 20:53:34 -0700 Subject: [PATCH 15/15] Enhance type annotations and refactor code across multiple components and fixed tests --- src/badger/gui/components/env_cbox.py | 4 +- src/badger/gui/components/pydantic_editor.py | 7 +- src/badger/gui/components/run_monitor.py | 18 +- src/badger/gui/mini/components/var_table.py | 12 +- src/badger/gui/mini/pages/home_page.py | 8 +- src/badger/gui/mini/pages/routine_page.py | 4 +- src/badger/gui/windows/add_random_dialog.py | 47 ++- src/badger/gui/windows/review_dialog.py | 10 +- src/badger/interface.py | 3 +- src/badger/routine.py | 4 +- src/badger/tests/conftest.py | 58 +++- src/badger/tests/test_gui_basic.py | 301 ++++++++++--------- src/badger/tests/test_home_page.py | 70 ++--- src/badger/tests/test_routine_page.py | 23 +- src/badger/tests/utils.py | 5 +- 15 files changed, 334 insertions(+), 240 deletions(-) diff --git a/src/badger/gui/components/env_cbox.py b/src/badger/gui/components/env_cbox.py index 36b820ff..a77de27e 100644 --- a/src/badger/gui/components/env_cbox.py +++ b/src/badger/gui/components/env_cbox.py @@ -639,12 +639,12 @@ def compose_vocs(self) -> tuple[VOCS, list[str]]: rule = objective[obj_name][0] objectives[obj_name] = rule - constraints: dict[str, tuple[str, float]] = {} + constraints: dict[str, list[str | float]] = {} critical_constraints: list[str] = [] for constraint in self.con_table.export_data(): con_name = next(iter(constraint)) relation, threshold, critical = constraint[con_name] - constraints[con_name] = (CONS_RELATION_DICT[relation], threshold) + constraints[con_name] = [CONS_RELATION_DICT[relation], threshold] if critical: critical_constraints.append(con_name) diff --git a/src/badger/gui/components/pydantic_editor.py b/src/badger/gui/components/pydantic_editor.py index 7c342c15..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, @@ -166,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) @@ -183,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, ) diff --git a/src/badger/gui/components/run_monitor.py b/src/badger/gui/components/run_monitor.py index b86e97c5..f532558c 100644 --- a/src/badger/gui/components/run_monitor.py +++ b/src/badger/gui/components/run_monitor.py @@ -673,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") - - self.sig_run_name.emit(run["filename"]) - self.sig_status.emit( - f"Archive success: Run data archived to {BADGER_ARCHIVE_ROOT}" - ) + 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}" + ) # if not self.testing: # QMessageBox.information( # self, 'Success!', 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 eddae1fd..7f9cf726 100644 --- a/src/badger/gui/mini/pages/home_page.py +++ b/src/badger/gui/mini/pages/home_page.py @@ -478,12 +478,14 @@ def validate_loaded_data_keys(self, vocs: VOCS, open_dialog: bool = True) -> Non 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") diff --git a/src/badger/gui/mini/pages/routine_page.py b/src/badger/gui/mini/pages/routine_page.py index f01cf2e9..5c08a4c7 100644 --- a/src/badger/gui/mini/pages/routine_page.py +++ b/src/badger/gui/mini/pages/routine_page.py @@ -193,7 +193,9 @@ def init_ui(self) -> None: 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() 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/review_dialog.py b/src/badger/gui/windows/review_dialog.py index 24a1e0b7..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, routine: Routine, parent: QDialog | None = None): + def __init__(self, parent: QWidget | None, routine: Routine) -> None: super().__init__(parent) self.routine = routine diff --git a/src/badger/interface.py b/src/badger/interface.py index e6e81c03..43abff78 100644 --- a/src/badger/interface.py +++ b/src/badger/interface.py @@ -94,9 +94,8 @@ def set_values(self, channel_inputs: Mapping[str, float | list[float]]) -> None: 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. """ - raise NotImplementedError() def get_value(self, channel_name: str, **kwargs: Any) -> Any: return self.get_values([channel_name], **kwargs)[channel_name] diff --git a/src/badger/routine.py b/src/badger/routine.py index a8d7d6f6..02c0a3c7 100644 --- a/src/badger/routine.py +++ b/src/badger/routine.py @@ -46,6 +46,8 @@ 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: dict[str, Any] | None = Field(None) @@ -153,7 +155,7 @@ def evaluate_point( logger.debug(f"Evaluation result: {obs}") return obs - data["evaluator"] = Evaluator(function=evaluate_point) # type: ignore[call-arg] + data["evaluator"] = Evaluator(function=evaluate_point) return data diff --git a/src/badger/tests/conftest.py b/src/badger/tests/conftest.py index 63c237f4..2f5041f7 100644 --- a/src/badger/tests/conftest.py +++ b/src/badger/tests/conftest.py @@ -7,6 +7,59 @@ 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 _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: @@ -22,7 +75,6 @@ def suppress_popups(mocker: MockerFixture) -> None: @pytest.fixture(scope="module", autouse=True) def config_test_settings( - mock_plugin_root: str, mock_template_root: str, mock_logbook_root: str, mock_archive_root: str, @@ -38,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") @@ -50,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) @@ -61,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) diff --git a/src/badger/tests/test_gui_basic.py b/src/badger/tests/test_gui_basic.py index a5f3c574..18955278 100644 --- a/src/badger/tests/test_gui_basic.py +++ b/src/badger/tests/test_gui_basic.py @@ -22,24 +22,24 @@ def test_gui_main(qtbot: QtBot) -> None: 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 - - window.process_manager.close_processes() + # TODO: Test if generator tab has been filled + # assert window.stacks.currentWidget().tabs.currentIndex() == 1 + finally: + window.process_manager.close_processes() @pytest.mark.usefixtures("init_multiprocessing") @@ -52,39 +52,39 @@ def test_close_main(qtbot: QtBot) -> None: fix_path_issues() window = BadgerMainWindow() + try: + qtbot.addWidget(window) - qtbot.addWidget(window) - - 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) + loop = QEventLoop() + QTimer.singleShot(5000, loop.quit) # 1000 ms pause + loop.exec_() - loop = QEventLoop() - QTimer.singleShot(1000, 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) - window.close() # this action should release the env + loop = QEventLoop() + QTimer.singleShot(1000, loop.quit) # 1000 ms pause + loop.exec_() - assert home_page.run_monitor.routine is not None + 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_processes() + # So we expect an AttributeError here + with pytest.raises(AttributeError): + _ = home_page.run_monitor.routine.environment + finally: + window.process_manager.close_processes() @pytest.mark.usefixtures("init_multiprocessing") @@ -100,44 +100,46 @@ def test_traceback_during_run(qtbot: QtBot, monkeypatch: pytest.MonkeyPatch) -> fix_path_issues() window = BadgerMainWindow() + 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) - 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 + # Patch showEvent to auto-accept the traceback dialog when it appears. + original_show_event = BadgerScrollableMessageBox.showEvent - assert ins # make sure the dialog is created - assert ins.detailedTextWidget.toPlainText() # make sure it's not empty + def patched_show_event( + ins: BadgerScrollableMessageBox, event: QShowEvent | None + ) -> None: + original_show_event(ins, event) # Call the original showEvent - QTimer.singleShot(100, ins.accept) # Close the dialog after 100 ms + assert ins # make sure the dialog is created + assert ins.detailedTextWidget.toPlainText() # make sure it's not empty - monkeypatch.setattr(BadgerScrollableMessageBox, "showEvent", patched_show_event) + QTimer.singleShot(100, ins.accept) # Close the dialog after 100 ms - 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_processes() + 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: QtBot) -> None: @@ -151,59 +153,62 @@ def test_measurement_retry_dialog_in_app_flow(qtbot: QtBot) -> None: 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.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", - ) - ) + 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_() - # Explicitly check if any errors were thrown - runner.check_queue() + 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", + ) + ) + + # 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 == MeasurementActionMessage(action="retry") + # Explicitly check if any errors were thrown + runner.check_queue() - window.process_manager.close_processes() - 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 @@ -219,34 +224,34 @@ def test_default_low_noise_prior_in_bo(qtbot: QtBot) -> None: 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) - - 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 + qtbot.addWidget(window) - window.process_manager.close_processes() + # 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 diff --git a/src/badger/tests/test_home_page.py b/src/badger/tests/test_home_page.py index 6221676a..c096ad0c 100644 --- a/src/badger/tests/test_home_page.py +++ b/src/badger/tests/test_home_page.py @@ -30,44 +30,44 @@ def test_home_page_run_routine(qtbot: QtBot) -> None: 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 = 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) + 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_processes() + # 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_routine_page.py b/src/badger/tests/test_routine_page.py index 63fa8644..1aa49cd5 100644 --- a/src/badger/tests/test_routine_page.py +++ b/src/badger/tests/test_routine_page.py @@ -241,12 +241,20 @@ def test_add_random_points(qtbot: QtBot) -> None: 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] diff --git a/src/badger/tests/utils.py b/src/badger/tests/utils.py index 8377a4c5..0b934778 100644 --- a/src/badger/tests/utils.py +++ b/src/badger/tests/utils.py @@ -59,7 +59,6 @@ def create_multiobjective_routine() -> Routine: return Routine( name="test", - vocs=vocs, generator=generator, environment={"name": "multiobjective_test"}, initial_points=pd.DataFrame(test_routine["init_points"]), @@ -99,7 +98,6 @@ def create_routine_turbo() -> Routine: return Routine( name="test-turbo", - vocs=vocs, generator=generator, environment={"name": "test"}, initial_points=pd.DataFrame(test_routine["config"]["init_points"]), @@ -128,7 +126,6 @@ def create_routine_critical() -> Routine: return Routine( name="test", - vocs=vocs, generator=generator, environment={"name": "test"}, initial_points=pd.DataFrame(test_routine["init_points"]), @@ -184,7 +181,7 @@ def get_vars_in_row(routine: Routine, idx: int = 0) -> np.ndarray: var_names = routine.vocs.variable_names if routine.data is None: raise ValueError("Routine data is None. Unable to get variables in row.") - output = routine.data.iloc[idx][var_names].to_numpy() + output: np.ndarray = routine.data.iloc[idx][var_names].to_numpy() # output = routine.data.iat[idx][var_names].to_numpy() return output