diff --git a/.deepsource.toml b/.deepsource.toml index 74a47e2b1..95b03efc0 100644 --- a/.deepsource.toml +++ b/.deepsource.toml @@ -2,7 +2,9 @@ version = 1 test_patterns = [ "tests/**", - "tests/**/**" + "tests/**/**", + "server/tests/**", + "server/tests/**/**", ] exclude_patterns = [ @@ -17,6 +19,8 @@ exclude_patterns = [ "tools/**", "tests/**", "tests/**/**", + "server/tests/**", + "server/tests/**/**", ] [[analyzers]] @@ -24,4 +28,4 @@ name = "python" [analyzers.meta] runtime_version = "3.x.x" - max_line_length = 120 \ No newline at end of file + max_line_length = 120 diff --git a/.gitignore b/.gitignore index 15d2b60d3..b1b0ab76a 100644 --- a/.gitignore +++ b/.gitignore @@ -143,3 +143,4 @@ users.sqlite3 testerror.txt /server/server_state /server/pssm_gremlin_server.egg-info +/.tmp diff --git a/CHANGELOG.md b/CHANGELOG.md index 07bb0a4fb..55ded4bb8 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -130,6 +130,25 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - Renamed Simplified Chinese label from 中文 to 简体中文 in language registry (`language.json`). ### Fixed +- **Package manager self-bootstrap**: manager UI and extras registry now + bootstrap into a writable runtime directory instead of writing into `src/` or + an installed package directory. Bootstrap fetches use bounded timeouts and + retry/backoff, and downloaded bootstrap assets must match the HMAC manifest + before they are written. The Refresh button remains the explicit + user-triggered network update path. `fetch_tags` degrades silently on network + failure instead of blocking startup with an error popup. +- **Self-upgrade HMAC integrity**: the self-upgrade flow now verifies downloaded + assets against an HMAC-SHA256 manifest before applying changes. The HMAC key + is embedded in the plugin source and never leaves the client. `make upload-gists` + generates `manifest.json` automatically. +- **Download registry integrity checks**: file registries now reject entries + without hashes, Monaco editor downloads verify npm tarball integrity metadata, + and remote ESM weight downloads verify configured MD5 hashes. +- **Monaco editor token handling**: generated editor backend authentication + tokens are no longer written to application logs. +- **Generated socket authentication keys**: `generate_strong_password()` now uses + `secrets.choice` instead of the default pseudo-random generator when creating + authentication/password strings. - **Uploaded FASTA filename collision**: two users uploading files with the same name (e.g. `seqs.fasta`) would overwrite each other in the shared upload directory. Files are now saved as `{owner_scoped_md5}.fasta` so on-disk names diff --git a/Makefile b/Makefile index 5036d08f6..e76ac3123 100644 --- a/Makefile +++ b/Makefile @@ -81,6 +81,19 @@ upload-gists: gh gist edit c1e8bfe0fc0b9c60bf49ea04a550a044 -f REvoDesign-PyMOL-entry.ui src/REvoDesign/UI/REvoDesign-PyMOL-entry.ui # JSONs for installer gh gist edit c1e8bfe0fc0b9c60bf49ea04a550a044 -f REvoDesignExtrasTableRich.json jsons/REvoDesignExtrasTableRich.json + # HMAC manifest — key is extracted from the uploaded source file + @python -c '\n\ +import hmac, hashlib, json, re;\n\ +src = open("src/REvoDesign/tools/package_manager.py").read();\n\ +m = re.search(r"_MANAGER_HMAC_KEY\s*=\s*bytes\.fromhex\(\"([a-f0-9]+)\"\)", src);\n\ +key = bytes.fromhex(m.group(1));\n\ +files = {"REvoDesign_PyMOL.py": "src/REvoDesign/tools/package_manager.py", "REvoDesign-PyMOL-entry.ui": "src/REvoDesign/UI/REvoDesign-PyMOL-entry.ui", "REvoDesignExtrasTableRich.json": "jsons/REvoDesignExtrasTableRich.json"};\n\ +manifest = {name: hmac.new(key, open(path, "rb").read(), "sha256").hexdigest() for name, path in files.items()};\n\ +json.dump(manifest, open("/tmp/revodesign-manifest.json", "w"), indent=2);\n\ +print("Manifest:", json.dumps(manifest, indent=2))\n\ +' + gh gist edit c1e8bfe0fc0b9c60bf49ea04a550a044 -f manifest.json /tmp/revodesign-manifest.json + rm /tmp/revodesign-manifest.json install-pymol-plugin: cp ./src/REvoDesign/tools/package_manager.py ~/.pymol/startup/REvoDesign_PyMOL.py @@ -112,7 +125,7 @@ install-pytorch-cpu-non-mac: # local dev reinstall: make clean - make black;rm -r /Users/yyy/.REvoDesign/config/; pip install . -U + make black; rm -rf "$${HOME}/.REvoDesign/config/"; python -m pip install . -U compile-ui: stage=compile bash tools/translate.sh @@ -145,7 +158,7 @@ test: # ui-test-pymol: # # Run a tmp folder to make sure the tests are run on the installed version # mkdir -p $(TESTDIR) -# cd $(TESTDIR); pymol -ckqy /Users/yyy/miniconda_py39_arm64/lib/python3.10/site-packages/REvoDesign/tests/cases/tabs/; +# cd $(TESTDIR); pymol -ckqy $(PYTEST_CASES_PATH)/tabs/; # all test all-test: diff --git a/dev/TODO.md b/dev/TODO.md index 85e50a29a..937cb4506 100644 --- a/dev/TODO.md +++ b/dev/TODO.md @@ -2,61 +2,19 @@ 1. **High – Harden QtSocketConnector serialization** `src/REvoDesign/clients/QtSocketConnector.py:160-192` still base64-encodes arbitrary `pickle` payloads that are deserialized on receipt, so any collaborator can trigger arbitrary code execution. Replace pickle with a safe schema (msgpack/JSON) and add signature or ACL checks before calling `broadcast`/`digest_dict`. -2. **High – Package manager bootstrapping assumes writable, online site-packages** - `src/REvoDesign/tools/package_manager.py:358-404` and `:970-993` always download UI/assets from GitHub gists during startup and write them back into the installed package directory. On read-only installs (system pip, CI wheels) this raises permission errors, and there is no offline cache fallback or signature/timeout on the download. Vendor these assets with the package and gate network calls behind explicit user consent with timeouts. -3. **High – Config bootstrap requires GUI prompts** - `src/REvoDesign/bootstrap/set_config.py:18-88` shows Qt modal dialogs while deciding whether to reset config trees. The bootstrap runs during import (tests call `set_REvoDesign_config_file()`), so headless runs and automation hang waiting for GUI input. Refactor these prompts behind a CLI/flag interface and avoid Qt when `PYMOL_GIT_TEST`/headless. 4. **High – Module-level PyMOL imports make the library unusable outside PyMOL** Modules such as `src/REvoDesign/evaluate/evaluator.py:7-18` and `src/REvoDesign/tools/mutant_tools.py:5-26` import `from pymol import cmd` at import time. Any script, CLI, or test that touches these modules without a PyMOL runtime crashes before code executes. Defer PyMOL imports to call sites or inject an abstraction so non-PyMOL contexts can still use the logic. -5. **High – REvoDesignPlugin changes the global working directory** - `src/REvoDesign/REvoDesign.py:103-145` calls `os.chdir(self.PWD)` every time a window opens, mutating the process-wide CWD for all other PyMOL plugins and threads. Persisting state should rely on explicit paths stored in the bus, not on moving the interpreter’s CWD. 6. **Medium – ConfigBus is a mutable global singleton storing Qt widget handles** `src/REvoDesign/driver/ui_driver.py:365-381` binds the singleton instance to the active UI and exposes widgets to every module. This leaks Qt objects across tests, prevents multiple windows, and requires manual resets (see `tests/conftest.reset_singletons`). Break ConfigBus into a pure configuration store plus scoped UI adapters so state does not persist globally. -7. **Medium – Developer-only absolute paths are committed** - The Makefile and tests refer to `/Users/yyy/...` (e.g., `Makefile:113`, `tests/data/test_data.py:393`), so running scripts on any other machine wipes or looks for that specific user’s home directory. Replace these with `$HOME`/config values or derive from repo-relative paths. 8. **Medium – PIPPack runner calls private APIs of its dependency** `src/REvoDesign/sidechain/mutate_runner/PIPPack.py:64-105` invokes `_initialize_with_a_model`, `_run_repack_single`, and `_run_repack_batch` on the third-party `PIPPack` object. These underscore-prefixed methods are not part of the public API and break whenever upstream refactors; switch to supported entrypoints or wrap the CLI. -9. **High – Shell pipeline script evaluates user input** - `server/REvoDesign_PSSM_GREMLIN.sh:158-220` builds shell commands containing user-provided file paths and runs them via `eval "$cmd"`. A crafted FASTA path or output directory injects arbitrary shell code into cluster jobs. Use arrays/`"$@"` instead of eval and properly quote user input. -10. **High – WebSocket keys use non-cryptographic randomness** - `src/REvoDesign/tools/utils.py:743-765` uses `random.choice` to generate passwords, and `src/REvoDesign/REvoDesign.py:1158-1167` feeds those into `ui.socket.input.key`. The default PRNG is predictable and unsuitable for authentication secrets; switch to `secrets.choice`/`secrets.token_urlsafe`. 11. **High – Importing the logger mutates user state** `src/REvoDesign/logger/logger.py:1-120` configures logging, creates directories under `user_log_path`, and spins up a background `QueueListener` during import. Simply touching `REvoDesign.logger` writes to the host filesystem and leaves live threads around; lazily initialize logging from the entrypoint instead. -12. **High – WorkerThread drops exceptions** - `src/REvoDesign/tools/package_manager.py:2296-2345` executes `self.func` in `QThread.run()` without trapping exceptions or emitting an error signal. Failures just print to stderr while the UI thinks work succeeded. Catch exceptions, forward them through a signal, and surface them via `notify_signal`. -13. **Medium – Conformer preview leaks PyMOL processes** - `src/REvoDesign/shortcuts/function_utils.py:17-57` launches a brand-new `pymol -xi` process for every preview and never registers or terminates it, so repeated clicks accumulate orphan PyMOL instances. Reuse the existing PyMOL session or track the subprocess via `RunningProcessRegistry` so abort/exit can kill it. -14. **High – Tests still touch the real config tree** - `tests/conftest.py:53-124` calls `set_REvoDesign_config_file()` at import time, which may pop up Qt dialogs and can delete `/Users/yyy/Library/Application Support/REvoDesign/config` if the user confirms. Mock `platformdirs.user_data_dir`/`user_cache_dir` so tests operate entirely inside the repo. -15. **High – RFdiffusion helper forces Qt5Agg backend globally** - `src/REvoDesign/tools/rfdiffusion_tools.py:1-7` unconditionally runs `matplotlib.use("Qt5Agg")` on import, breaking headless builds and notebooks that rely on Agg. Only set the backend when a GUI feature is invoked and honor `MPLBACKEND`. -16. **High – Download registry accepts unsigned payloads** - `src/REvoDesign/tools/download_registry.py:133-152` leaves registry entries as `None`, so `pooch` skips hash verification and trusts whatever the server returns. Require hashes for every asset or refuse to download; otherwise a MITM can ship arbitrary binaries. -17. **High – RFdiffusion weight downloads block the UI** +16. **High – RFdiffusion weight downloads block the UI** `src/REvoDesign/shortcuts/tools/rfdiffusion_tasks.py:90-151` fetches multi‑GB checkpoints synchronously inside `RfDiffusion.pick_model`, freezing PyMOL while `pooch` downloads. Move the download to `run_worker_thread_in_pool` and provide progress/abort hooks. -18. **High – DglSolver silently runs pip install** - `src/REvoDesign/shortcuts/tools/rfdiffusion_tasks.py:33-70` detects CUDA and immediately executes `python -m pip install dgl==2.2.1` in the current interpreter without prompting or isolating the environment. Prompt the user and install inside a managed env/container so you don’t mutate the host Python unexpectedly. -19. **High – Server keeps plaintext credentials in repo** - `server/pssm_gremlin/pssm_gremlin.py:17-36` parses `users.txt` where each line is `username:password` and only hashes them at runtime, so the file on disk stores the raw passwords. Store hashed secrets (or use a proper auth backend) to avoid leaking credentials if the repo/server is compromised. -20. **High – GREMLIN web service hardcodes cluster-only paths** - `server/pssm_gremlin/pssm_gremlin.py:38-120` assumes `/mnt/data/...`, `/mnt/db/...`, and `redis://localhost` exist and derives Docker mounts from `/home/{os.getlogin()}`. With no configuration hooks the service cannot run off that exact cluster. Externalize broker paths, DBs, and output dirs via a config file/env vars. -21. **Medium – Config-driven env overrides clobber user variables** - `src/REvoDesign/driver/environ_register.py:8-39` iterates `environ.yaml` and unconditionally sets each key in `os.environ`, overwriting whatever the shell provided (proxy settings, CUDA flags, etc.). Respect existing values or mark overrides explicitly in config. -22. **High – Archive extraction is vulnerable to path traversal** - `src/REvoDesign/tools/utils.py:592-627` calls `tarfile.ZipFile(...).extractall()` on user-selected archives with no member filtering, so a malicious tar can overwrite arbitrary files. Validate members (reject `..`/absolute paths) or use Python’s newer `tarfile` filters. -23. **High – SSL certificate manager writes weak certs** - `src/REvoDesign/tools/ssl_certificates.py:34-115` generates self-signed certs with a fixed serial number, week-long validity, and default filesystem permissions under a shared cache dir. Anyone on the box can read or overwrite the keys. Randomize serials, tighten permissions, and allow the validity to be configured. -24. **High – “Background” workers busy-wait on the GUI thread** +22. **High – “Background” workers busy-wait on the GUI thread** `src/REvoDesign/tools/package_manager.py:2635-2661` starts a `WorkerThread` but then loops on `work_thread.isFinished()` calling `refresh_window()`/`time.sleep`, so the main thread stays blocked and abort buttons can’t repaint. Switch to signal-driven completion instead of polling. -25. **High – Aborting a worker interrupts the whole PyMOL session** - `src/REvoDesign/tools/package_manager.py:2223-2244` handles every abort click by calling `cmd.interrupt()`, which cancels whatever PyMOL is currently doing—even unrelated jobs. Terminate only the subprocesses registered for that worker rather than nuking the global interpreter. -26. **High – Monaco editor logs its auth token** - `src/REvoDesign/editor/monaco/server.py:71-78` generates a `SECRET_TOKEN` and immediately logs it, so anyone with read access to the log file can hijack the editor backend. Never log secrets and rotate them when the server restarts. -27. **High – WebSocket clients hand over the shared auth key on request** +24. **High – WebSocket clients hand over the shared auth key on request** `src/REvoDesign/clients/QtSocketConnector.py:813-863` responds to a `RequireKey` message by sending `self.authentication_key` before the server proves its identity. A malicious host can simply ask for the key and impersonate everyone. Authenticate the server (TLS + pinned cert) and only send a key through an already-authenticated channel. -28. **Medium – Git tag fetches are synchronous and unauthenticated** - `src/REvoDesign/tools/package_manager.py:2349-2380` hits the public GitHub tags API on startup without auth, caching, or error UI. When the rate limit is exceeded the manager freezes for ~10s and users see nothing. Add caching, a timeout, and a friendly error surface. -29. **Medium – pip installs run on the UI thread** +26. **Medium – pip installs run on the UI thread** `src/REvoDesign/tools/package_manager.py:861-925` calls `pip install/uninstall` directly from menu actions, blocking PyMOL for minutes during heavy installs. Dispatch these calls to `run_worker_thread_in_pool` and stream progress to the dashboard. -30. **Medium – Archive browsing keeps every extraction forever** - `src/REvoDesign/driver/file_dialog.py:240-269` always expands compressed inputs into `/expanded_compressed_files/` with no cleanup or quota. Repeated browsing bloats the session dir and a malicious 20 GB tar can fill the disk. Extract into a temp dir and delete it once the dialog closes. diff --git a/dev/tools/generate_ui_typing.py b/dev/tools/generate_ui_typing.py index 7a42e4cbf..95efed645 100644 --- a/dev/tools/generate_ui_typing.py +++ b/dev/tools/generate_ui_typing.py @@ -8,9 +8,10 @@ import argparse import sys -import xml.etree.ElementTree as ET from pathlib import Path +from REvoDesign.Qt import QtCore + REPO_ROOT = Path(__file__).resolve().parents[2] UI_PATH = REPO_ROOT / "src/REvoDesign/UI/REvoDesign.ui" TYPES_PATH = REPO_ROOT / "src/REvoDesign/UI/types.py" @@ -70,20 +71,42 @@ class UiBinding: + __slots__ = ("name", "type_name", "warning") + def __init__(self, name: str, type_name: str, warning: str | None = None) -> None: self.name = name self.type_name = type_name self.warning = warning -def _extract_custom_widget_map(root: ET.Element) -> dict[str, str]: +def _read_ui_elements(ui_path: Path) -> tuple[dict[str, str], list[tuple[str, str, str]]]: custom_types: dict[str, str] = {} - for custom_widget in root.findall("./customwidgets/customwidget"): - class_name = custom_widget.findtext("class") - if not class_name: - continue - custom_types[class_name] = FALLBACK_WIDGET_TYPE - return custom_types + elements: list[tuple[str, str, str]] = [] + reader = QtCore.QXmlStreamReader(ui_path.read_text(encoding="utf-8")) + in_custom_widget = False + + while not reader.atEnd(): + reader.readNext() + if reader.isStartElement(): + tag = reader.name() + if tag == "customwidget": + in_custom_widget = True + continue + if in_custom_widget and tag == "class": + class_name = reader.readElementText().strip() + if class_name: + custom_types[class_name] = FALLBACK_WIDGET_TYPE + continue + if tag in {"widget", "layout", "action", "buttongroup"}: + attributes = reader.attributes() + elements.append((tag, attributes.value("name"), attributes.value("class"))) + elif reader.isEndElement() and reader.name() == "customwidget": + in_custom_widget = False + + if reader.hasError(): + raise ValueError(f"Failed to parse {ui_path}: {reader.errorString()}") + + return custom_types, elements def _resolve_type_name(tag: str, class_name: str, custom_types: dict[str, str]) -> tuple[str, str | None]: @@ -103,18 +126,13 @@ def _resolve_type_name(tag: str, class_name: str, custom_types: dict[str, str]) def collect_ui_bindings(ui_path: Path = UI_PATH) -> list[UiBinding]: - root = ET.parse(ui_path).getroot() - custom_types = _extract_custom_widget_map(root) + custom_types, elements = _read_ui_elements(ui_path) bindings: dict[str, UiBinding] = {} - for element in root.iter(): - if element.tag not in {"widget", "layout", "action", "buttongroup"}: - continue - name = element.attrib.get("name", "") - class_name = element.attrib.get("class", "") + for tag, name, class_name in elements: if not name or not name.isidentifier(): continue - type_name, warning = _resolve_type_name(element.tag, class_name, custom_types) + type_name, warning = _resolve_type_name(tag, class_name, custom_types) bindings.setdefault(name, UiBinding(name=name, type_name=type_name, warning=warning)) return [bindings[name] for name in sorted(bindings)] diff --git a/dev/tools/validate_package_data.py b/dev/tools/validate_package_data.py index 1b110ac0b..f63368f22 100644 --- a/dev/tools/validate_package_data.py +++ b/dev/tools/validate_package_data.py @@ -14,6 +14,8 @@ REPO_ROOT / "src/REvoDesign/UI/REvoDesign.ui", REPO_ROOT / "src/REvoDesign/UI/language/eng-chs.qm", REPO_ROOT / "src/REvoDesign/UI/language/eng-cht.qm", + REPO_ROOT / "src/REvoDesign/UI/REvoDesign-PyMOL-entry.ui", + REPO_ROOT / "jsons/REvoDesignExtrasTableRich.json", ) diff --git a/dev/tools/validate_ui_i18n.py b/dev/tools/validate_ui_i18n.py index 06207c42d..82926429e 100644 --- a/dev/tools/validate_ui_i18n.py +++ b/dev/tools/validate_ui_i18n.py @@ -10,7 +10,7 @@ import sys from pathlib import Path -from REvoDesign.Qt import QtCore, QtWidgets +from REvoDesign.Qt import QtWidgets from REvoDesign.Qt.ui_runtime_loader import load_runtime_ui REPO_ROOT = Path(__file__).resolve().parents[2] @@ -18,6 +18,11 @@ LINGUIST_PROJECT = REPO_ROOT / "src/REvoDesign/UI/liguist.pro" +def require(condition: bool, message: str) -> None: + if not condition: + raise RuntimeError(message) + + def main(argv: list[str] | None = None) -> int: parser = argparse.ArgumentParser(description=__doc__) parser.parse_args(argv) @@ -36,26 +41,26 @@ def main(argv: list[str] | None = None) -> int: window, ui = load_runtime_ui(UI_PATH) try: # --- RuntimeUiProxy i18n contract --- - assert hasattr(ui, "trans"), "RuntimeUiProxy must expose a `trans` attribute for legacy i18n code" + require(hasattr(ui, "trans"), "RuntimeUiProxy must expose a `trans` attribute for legacy i18n code") trans = ui.trans - assert hasattr(trans, "load"), "translator must expose `load`" - assert hasattr(trans, "isEmpty"), "translator must expose `isEmpty`" - assert hasattr(trans, "translate"), "translator must expose `translate`" + require(hasattr(trans, "load"), "translator must expose `load`") + require(hasattr(trans, "isEmpty"), "translator must expose `isEmpty`") + require(hasattr(trans, "translate"), "translator must expose `translate`") # retranslateUi must work ui.retranslateUi(window) # `trans` must survive retranslateUi's internal refresh_bindings call - assert hasattr(ui, "trans"), "RuntimeUiProxy.trans must survive retranslateUi → refresh_bindings" + require(hasattr(ui, "trans"), "RuntimeUiProxy.trans must survive retranslateUi -> refresh_bindings") # --- Dynamic language action registry --- - assert hasattr(ui, "menuLanguage"), "menuLanguage menu must be available for dynamic actions" + require(hasattr(ui, "menuLanguage"), "menuLanguage menu must be available for dynamic actions") new_action = QtWidgets.QAction() new_action.setObjectName("actionTestLang") - setattr(ui, "actionTestLang", new_action) + ui.actionTestLang = new_action ui.menuLanguage.addAction(new_action) - assert hasattr(ui, "actionTestLang"), "RuntimeUiProxy must accept dynamic action attributes" - assert getattr(ui, "actionTestLang") is new_action + require(hasattr(ui, "actionTestLang"), "RuntimeUiProxy must accept dynamic action attributes") + require(ui.actionTestLang is new_action, "RuntimeUiProxy must retain dynamic action attributes") # --- LanguageSwitch instantiation / exercise (smoke) --- # LanguageSwitch depends on a ConfigBus singleton that has `ui` already @@ -68,9 +73,9 @@ def main(argv: list[str] | None = None) -> int: config_bus.ui = ui lang_switch = LanguageSwitch(window) - assert hasattr(lang_switch, "trans"), "LanguageSwitch must own a `trans` reference" - assert lang_switch.trans is not None, "LanguageSwitch.trans must not be None" - assert hasattr(lang_switch.trans, "load"), "LanguageSwitch.trans must be a translator-like object" + require(hasattr(lang_switch, "trans"), "LanguageSwitch must own a `trans` reference") + require(lang_switch.trans is not None, "LanguageSwitch.trans must not be None") + require(hasattr(lang_switch.trans, "load"), "LanguageSwitch.trans must be a translator-like object") # English / source-language switch must not raise english = lang_switch.language_items[0] diff --git a/docs/api/download-registry.md b/docs/api/download-registry.md index 0a5f712e3..cd5be13e4 100644 --- a/docs/api/download-registry.md +++ b/docs/api/download-registry.md @@ -26,7 +26,7 @@ registry = FileDownloadRegistry( base_url="https://example.com/data/", registry={ "file1.pdb": "md5:abc123def456", - "file2.pkl": None, # no hash verification + "file2.pkl": "sha256:0123456789abcdef...", }, alternative_base_urls=["https://mirror.example.com/data/"], retry_count=3, @@ -65,6 +65,10 @@ local filesystem path. The archive extraction feature (`flatten_archive`) can automatically unpack `.tar.gz` and `.zip` archives into a `_flatten/` subdirectory. +Every registry entry must include an `md5:`, `sha256:`, or other Pooch-supported +hash string. Entries with `None` or an empty hash are rejected before any +download starts. + ## Creating Registries from MD5 Checksums ```python diff --git a/docs/dev-guide/package-manager.md b/docs/dev-guide/package-manager.md index 2d7697aae..2ec6f8e31 100644 --- a/docs/dev-guide/package-manager.md +++ b/docs/dev-guide/package-manager.md @@ -85,8 +85,8 @@ killed. ### REvoDesignPackageManager (`:::` REvoDesign.tools.package_manager.REvoDesignPackageManager) The top-level orchestrator. It: -1. Fetches the UI file (`REvoDesign-PyMOL-entry.ui`) from Gist and loads it via `loadUi` -2. Fetches the extras manifest JSON (`REvoDesignExtrasTableRich.json`) from Gist +1. Self-bootstraps the UI file (`REvoDesign_installer.ui`) via retrying Gist fetches +2. Self-bootstraps the extras manifest JSON (`REvoDesignExtrasTableRich.json`) via retrying Gist fetches 3. Fetches GitHub release tags for version selection 4. Registers PyMOL menu items (Install, Reinstall, Uninstall, Manager) 5. Configures proxy, mirror, and extras in the UI @@ -94,16 +94,31 @@ The top-level orchestrator. It: ## The UI File -The manager's graphical interface is defined in `REvoDesign-PyMOL-entry.ui`, -a Qt Designer file stored at `src/REvoDesign/UI/REvoDesign-PyMOL-entry.ui`. -This file is uploaded to a GitHub Gist via `make upload-gists` and fetched at -runtime by the manager plugin. +The manager is expected to work as a single PyMOL startup file. It therefore +self-bootstraps its UI from the canonical Gist source generated from +`src/REvoDesign/UI/REvoDesign-PyMOL-entry.ui` into a writable runtime bootstrap +directory under the system temp directory by default. Set +`REVODESIGN_PM_BOOTSTRAP_DIR` to override that bootstrap location in tests or +special deployments. Startup re-fetches bootstrap assets with retry/backoff +instead of treating a stale local copy as the source of truth, and it never +writes runtime assets into `src/` or an installed package directory. The +downloaded UI must match the HMAC in `manifest.json` before it replaces the +runtime bootstrap copy. ## Extras Registry -The file `REvoDesignExtrasTableRich.json` (version-tracked alongside the UI -file on Gist) defines available extras. It is fetched at manager startup and -displayed in a `CheckableListView` widget. Each extras entry can specify: +The canonical source file `jsons/REvoDesignExtrasTableRich.json` defines +available extras and is uploaded to Gist with `make upload-gists`. The manager +self-bootstraps the runtime copy into the same writable bootstrap directory +with retry/backoff. The Refresh button is the explicit user-triggered path for +fetching and writing a newer runtime copy from Gist. Each extras entry can +specify: + +The bootstrap manifest uses HMAC-SHA256 to catch corruption, truncation, and +cross-project asset mismatches before runtime files are written. The HMAC key +is distributed with the standalone manager, so this is an integrity check for +the managed bootstrap flow rather than a private release signature against a +compromised Gist account. - **`name`** -- Display name - **`extras_id`** -- The pip extras identifier (e.g., `[scatter]`) diff --git a/docs/dev-guide/server.md b/docs/dev-guide/server.md index 23fd0df86..12f1175f0 100644 --- a/docs/dev-guide/server.md +++ b/docs/dev-guide/server.md @@ -383,6 +383,44 @@ Flask test clients with temporary SQLite databases and environment variables. A dedicated CI workflow (``.github/workflows/server-test.yml``) runs non-Docker tests on every push touching ``server/**``. +### Current Security Test Coverage + +The current automated security coverage is concentrated in: + +- `server/tests/test_security.py`: login throttling, enumeration resistance, + cookie-only write rejection, token tampering, password policy, admin + self-protection, batch-user authorization, upload/path handling, SQL/XSS + payload handling, plus-alias registration policy, task-id validation, and + endpoint method/content-type rejection. +- `server/tests/test_security_advanced.py`: upload/download path traversal, + CRLF/header-injection probes, SSTI/RCE-style payloads, command-injection + filenames, eval/prototype-pollution payloads, YAML/pickle body rejection, API + key brute-force rejection, and binary upload rejection. +- `server/tests/test_auth.py`: auth/session behavior that security checks rely + on, including Bearer-vs-cookie gates, API key lifecycle and privilege limits, + guest restrictions, CAPTCHA/token validation, forgot-password throttling, + deleted/pending/unverified-user rejection, registration/upload rate limits, + upload validation, task ownership separation, and schema validation. + +Current local evaluation snapshot, after `make clean`: + +```bash +conda run -n REvoDesignTestFlight python -m pytest \ + server/tests/test_security.py \ + server/tests/test_security_advanced.py \ + server/tests/test_auth.py \ + -k "security or attack or api_key or rate_limit or deleted_user or pending_user or unverified_email or admin" \ + --no-cov +``` + +Result on 2026-07-17: `66 passed, 43 deselected in 22.93s`. + +This is a non-Docker Flask test-client evaluation. It does not prove live Docker +socket behavior, Docker runtime group mapping, reverse-proxy header handling, or +staging deployment isolation. Use the manual security validation checks below +after Docker Compose, user/group, runner-launch, auth, or account-status +changes. + ## Security Validation Run these checks after auth, account-status, Docker Compose, user/group, or @@ -411,6 +449,283 @@ runner-launch changes: same IP, the login endpoint returns HTTP 429 with `retry_after_seconds`; the login page should disable the submit button and count down until retry. +### Docker Socket A/B Attack Test + +Run this test after Docker, Compose, user/group, or runner-launch changes. It +checks whether an authenticated low-privilege web user can turn task submission +into Docker daemon control. The expected secure result is not just "the app +works"; HTTP users must only submit normal FASTA tasks, must not reach Docker +API routes, must not choose Docker image/command/mounts, and must not use a +cookie-only CSRF request for task writes. + +Use an isolated staging instance. Do not run destructive Docker escape payloads +on a host that stores real data. + +Prepare: + +```bash +export BASE_URL=http://127.0.0.1:8080 +export ADMIN_USER=admin +export ADMIN_PASS='' + +docker ps --format '{{.ID}} {{.Names}} {{.Image}} {{.Ports}}' +docker inspect server-web-1 server-worker-1 \ + --format '{{.Name}} User={{.Config.User}} Groups={{json .HostConfig.GroupAdd}} Mounts={{json .Mounts}}' +docker exec server-web-1 id +docker exec server-worker-1 id +docker exec server-web-1 ls -ln /var/run/docker.sock +docker exec server-worker-1 ls -ln /var/run/docker.sock +``` + +A result: the socket is not accessible from web/worker. This proves the live +instance is not Docker-root-escape-capable, but Docker-backed task execution +will fail until permissions are fixed. Expected evidence: + +```text +/var/run/docker.sock -> srw-rw---- 1 0 0 ... +docker.from_env() from web/worker -> PermissionError(13, 'Permission denied') +submitted task -> failed with "Docker daemon unavailable" or PermissionError +``` + +B result: the socket is accessible from web/worker. This is operationally +required for runner containers, but it is host-root-equivalent if arbitrary code +ever runs in web/worker. In this case the test must prove the HTTP contract does +not expose Docker control. Expected evidence: + +```text +limited user /api/auth/admin/users -> 403 +cookie-only POST /PSSM_GREMLIN/api/post -> 403 Bearer token required +POST /containers/create and GET /containers/json on the Flask port -> 404 +path-like FASTA filename upload -> normal task only, sanitized filename, no host-path mount selection +server code fixes image, command, user, environment, and mounts from server config +``` + +The following probe creates a temporary non-admin user, checks the CSRF gate, +submits a harmless FASTA with a path-like filename, probes Docker-like HTTP +routes, and verifies API-key privilege limits: + +```bash +python - <<'PY' +import json +import os +import re +import uuid + +import requests + +base = os.environ["BASE_URL"] +admin_user = os.environ["ADMIN_USER"] +admin_pass = os.environ["ADMIN_PASS"] + + +def show(label, resp): + ctype = resp.headers.get("content-type", "") + body = resp.text[:220].replace("\n", " ") + if "application/json" in ctype: + try: + data = resp.json() + body = { + k: ("[redacted]" if k in {"token", "api_key"} else v) + for k, v in data.items() + } + except Exception: + pass + print( + json.dumps( + { + "label": label, + "status": resp.status_code, + "location": resp.headers.get("Location"), + "body": body, + }, + sort_keys=True, + ) + ) + + +admin = requests.Session() +r = admin.post( + f"{base}/PSSM_GREMLIN/api/auth/login", + json={"username": admin_user, "password": admin_pass}, + timeout=10, +) +show("admin_login", r) +r.raise_for_status() +admin_h = {"Authorization": "Bearer " + r.json()["token"]} + +name = "limited_" + uuid.uuid4().hex[:8] +pw = "Pw_" + uuid.uuid4().hex + "!aA1" +r = admin.post( + f"{base}/PSSM_GREMLIN/api/auth/admin/users", + headers=admin_h, + json={ + "username": name, + "email": name + "@audit.local", + "password": pw, + "affiliation": "audit", + "is_admin": False, + }, + timeout=10, +) +show("create_limited", r) + +limited = requests.Session() +r = limited.post( + f"{base}/PSSM_GREMLIN/api/auth/login", + json={"username": name, "password": pw}, + timeout=10, +) +show("limited_login", r) +r.raise_for_status() +limited_h = {"Authorization": "Bearer " + r.json()["token"]} + +r = limited.get( + f"{base}/PSSM_GREMLIN/api/auth/admin/users", + headers=limited_h, + timeout=10, +) +show("limited_admin_list", r) + +files = { + "file": ( + "../../../../var/run/docker.sock.fasta", + b">audit\nACDEFGHIKLMNPQRSTVWY\n", + ) +} +r = limited.post( + f"{base}/PSSM_GREMLIN/api/post", + files=files, + allow_redirects=False, + timeout=10, +) +show("cookie_only_upload", r) + +files = { + "file": ( + "../../../../--privileged--var-run-docker-sock.fasta", + b">audit\nACDEFGHIKLMNPQRSTVWY\n", + ) +} +r = limited.post( + f"{base}/PSSM_GREMLIN/api/post", + files=files, + headers=limited_h, + allow_redirects=False, + timeout=15, +) +show("bearer_upload_pathlike_name", r) +loc = r.headers.get("Location") or "" +match = re.search(r"/running/([0-9a-f]{32})", loc) +if match: + md5 = match.group(1) + show( + "running_after_upload", + limited.get( + f"{base}/PSSM_GREMLIN/api/running/{md5}", + headers=limited_h, + timeout=10, + ), + ) + show( + "delete_uploaded_task", + limited.delete( + f"{base}/PSSM_GREMLIN/api/delete/{md5}", + headers=limited_h, + timeout=10, + ), + ) + +for method, path in [ + ("get", "/containers/json"), + ("get", "/version"), + ("post", "/containers/create"), + ("get", "/PSSM_GREMLIN/api/docker"), + ("post", "/PSSM_GREMLIN/api/docker/run"), +]: + resp = getattr(limited, method)(base + path, headers=limited_h, timeout=10) + show(f"probe_{method}_{path}", resp) + +r = limited.post( + f"{base}/PSSM_GREMLIN/api/auth/me/api-key", + headers=limited_h, + timeout=10, +) +show("limited_generate_api_key", r) +if r.status_code == 201: + api_h = {"X-API-Key": r.json()["api_key"]} + clean = requests.Session() + show( + "apikey_admin_list", + clean.get( + f"{base}/PSSM_GREMLIN/api/auth/admin/users", + headers=api_h, + timeout=10, + ), + ) + files = {"file": ("api-key-task.fasta", b">audit\nACDEFGHIKLMNPQRSTVWY\n")} + show( + "apikey_upload_task", + clean.post( + f"{base}/PSSM_GREMLIN/api/post", + files=files, + headers=api_h, + allow_redirects=False, + timeout=15, + ), + ) +PY +``` + +Treat a failure in the expected status codes above as a security regression. If +B can access Docker and an HTTP user can influence image, command, privileged +mode, bind source paths, or socket mounts, this is a real host-root escape path. + +### Admin Self-Lockout Check + +Run this after admin user-management changes. The expected result is that an +admin cannot remove their own access, either through a single-user action or a +batch action. + +Expected results: + +```text +self PUT user_status=banned -> 400 Administrators cannot ban their own account +self DELETE -> 400 Administrators cannot delete their own account +batch disable [self, other] -> count 1, self stays active, other becomes banned +batch delete [self, other] -> count 1, self remains undeleted, other is deleted +``` + +### Banned-User Authentication Check + +Run this after account-status or token-auth changes. The check creates a +temporary user, confirms it can log in before the ban, bans it, and verifies +that new login, old Bearer token, and pre-existing API key are all rejected. + +Expected results: + +```text +pre_ban_login -> 200 +pre_ban_api_key -> 201 +ban_user -> 200 +post_ban_login -> 403 Account has been suspended +post_ban_old_bearer_me -> 401 Authentication required +post_ban_api_key_me -> 401 Authentication required +``` + +### Login Throttling Check + +Run this after changes to login, client IP extraction, CDN headers, or rate +limits. + +Expected results: + +```text +first 5 failed login attempts in one minute -> 401 +sixth failed login attempt from the same client identity -> 429 +429 JSON body includes retry_after_seconds +login page disables submit and counts down until retry +``` + ## REvoDesign Plugin Integration The REvoDesign plugin does **not** connect to the PSSM_GREMLIN server diff --git a/docs/dev-guide/testing.md b/docs/dev-guide/testing.md index 52885e05f..46a526b1d 100644 --- a/docs/dev-guide/testing.md +++ b/docs/dev-guide/testing.md @@ -91,6 +91,42 @@ tests/ └── test_openkinetics_scorer.py # example test module ``` +### Mini UniRef Test Databases + +Server GREMLIN tests use tiny UniRef30/UniRef90 mocks under +`tests/data/msa/miniuc`. The raw inputs stay in `tests/data/msa` because +non-server evolution tests also read `2KL8.i90c75_aln.fas` directly. + +The shared `miniuc_databases` fixture in `tests/conftest.py` rebuilds the mock +databases on demand when all required binaries are available: + +- `makeblastdb` +- `psiblast` +- `hhblits` +- `cstranslate` +- `ffindex_from_fasta` + +Make sure the REvoDesign Conda environment is active and `CONDA_PREFIX` points +to it before running server tests that need these databases: + +```bash +conda activate REvoDesignTestFlight +echo "$CONDA_PREFIX" +``` + +To rebuild the mini databases manually, follow +`tests/data/msa/README.md`. The expected generated locations are: + +``` +tests/data/msa/miniuc/uc30/miniuc30* +tests/data/msa/miniuc/uc90/uniref90* +tests/data/msa/testminiuc/uc30/ +tests/data/msa/testminiuc/uc90/ +``` + +If those files are missing and any required binary is unavailable, the relevant +server tests skip with a message pointing back to the MSA README. + ## Running Tests Tests **must** run from a temporary directory. The Makefile targets handle this diff --git a/docs/getting-started.md b/docs/getting-started.md index 761afb73a..95f58642d 100644 --- a/docs/getting-started.md +++ b/docs/getting-started.md @@ -92,11 +92,16 @@ Example `environ.yaml` (keys are commented out by default; uncomment and set your values): ```yaml +override_existing: false variables: # OPENKINETICS_API_KEY: your-openkinetics-api-key # auto-registers the key # ROSETTA_BIN: /path/to/rosetta/main/source/bin # Rosetta executable directory ``` +`environ.yaml` preserves variables already provided by the shell. Set +`override_existing: true` only when the config file should replace existing +process values. + The OpenKinetics API key is auto-registered by default: `OpenKineticsScorerAbstract` will automatically generate and persist a key on first use, so manual setup in `environ.yaml` is optional and takes diff --git a/pyproject.toml b/pyproject.toml index e240cec9d..2ca09bab5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -226,7 +226,7 @@ include = [ ] [tool.bandit] -exclude_dirs = ["build","dist","tests","scripts"] +exclude_dirs = ["build","dist","tests","server/tests","scripts"] number = 4 recursive = true targets = "src" diff --git a/server/.env.example b/server/.env.example index a3ef11558..eb1fe45a7 100644 --- a/server/.env.example +++ b/server/.env.example @@ -15,19 +15,22 @@ RUNNER_IMAGE=revodesign-pssm-gremlin-non-root # Host paths — must exist on the Docker host # ============================================================================ # Root directory for uploads, results, and the SQLite task database. -SERVER_DIR=/Users/yyy/Documents/protein_design/REvoDesign/playground/server_test +SERVER_DIR=/srv/revodesign/server +# Host directory that the Docker runner may bind from. Defaults to the parent +# of SERVER_DIR; set this when SERVER_DIR is under a symlink or alternate mount. +# RUNNER_HOST_ROOT=/srv/revodesign # Directory for Gunicorn and Celery log files. -LOG_DIR=/Users/yyy/Documents/protein_design/REvoDesign/playground/server_test/revodesign/server/logs +LOG_DIR=/srv/revodesign/logs # ============================================================================ # MSA databases — large read-only datasets mounted into the runner container # ============================================================================ # Prefix for the UniRef30 HHblits database (the directory containing the .cs219 # and .a3m files; omit extensions). -DB_UNIREF30=/Users/yyy/Documents/protein_design/REvoDesign/playground/miniuc/uc30/miniuc30 +DB_UNIREF30=/srv/revodesign/databases/uniref30/UniRef30_2023_02 # Prefix for the UniRef90 BLAST database (the directory containing the .pin/.phr # etc. files; omit extensions). -DB_UNIREF90=/Users/yyy/Documents/protein_design/REvoDesign/playground/miniuc/uc90/uniref90 +DB_UNIREF90=/srv/revodesign/databases/uniref90/uniref90 # ============================================================================ # Authentication — token signing and user database diff --git a/server/README.md b/server/README.md index 2ddbc5643..ff9d1ef91 100644 --- a/server/README.md +++ b/server/README.md @@ -1,22 +1,7 @@ # PSSM GREMLIN Server -The server is a pip-installable package (`pssm_gremlin_server`) providing a -Flask + Celery + Docker-runner web service for GREMLIN co-evolution analysis. - -## Development - -```bash -# Install in editable mode with test dependencies -pip install -e "server/[test]" - -# Run non-Docker tests from the repo root -pytest server/tests/ -v -k "not Docker and not docker" - -# Run the server directly (no Docker) -python -m pssm_gremlin_server.pssm_gremlin -``` - ---- +The server is a Docker-deployed Flask + Celery + Docker-runner web service for +GREMLIN co-evolution analysis. ## Production Deployment @@ -146,10 +131,11 @@ Fallback when `REVODESIGN_SERVER_ENV` is unset: | Variable | Purpose | | --- | --- | -| `SERVER_DIR` | Host root for uploads, sqlite, and result folders. | +| `SERVER_DIR` | Host root for uploads, sqlite, and result folders (default: `./pssm_gremlin_data`). | +| `RUNNER_HOST_ROOT` | Host root allowed for Docker runner bind mounts (default: parent of `SERVER_DIR`). | | `LOG_DIR` | Host directory for Gunicorn/Celery logs. | -| `DB_UNIREF30` | UniRef30 prefix path. | -| `DB_UNIREF90` | UniRef90 BLAST prefix path. | +| `DB_UNIREF30` | UniRef30 prefix path (default: `{SERVER_DIR}/db/uniref30/UniRef30_2022_02`). | +| `DB_UNIREF90` | UniRef90 BLAST prefix path (default: `{SERVER_DIR}/db/uniref90/uniref90`). | | `AUTH_SECRET_KEY` | Fixed secret for signing auth tokens. Set in production so tokens survive restarts. | | `AUTH_TOKEN_MAX_AGE` | Token lifetime in seconds (default: 604800 = 7 days). | | `USER_DB_PATH` | Path to the user database (default: `{SERVER_DIR}/users.sqlite3`). | @@ -166,7 +152,6 @@ Fallback when `REVODESIGN_SERVER_ENV` is unset: | `WORKER_CONCURRENCY` | Celery worker concurrency. | | `GUNICORN_WORKERS` | Gunicorn worker count. | | `PORT` | Public HTTP port. | -| `ALLOWED_EMAIL_DOMAINS` | Comma-separated allowed email domains for registration (empty = all allowed). Also normalises plus-aliased addresses. | | `PUBLIC_DASHBOARD` | `false` by default; scopes task visibility to owner unless admin. | | `ADMIN_USERS` | Comma-separated admin usernames for cross-user management. | | `ADMIN_NOTIFY_EMAIL` | Comma-separated admin email addresses for new-user registration digests (default: empty = no notification). | @@ -440,254 +425,20 @@ The web and worker containers mount `/var/run/docker.sock` to spawn runner conta - Consider using a Docker socket proxy (e.g. `docker-socket-proxy`) to restrict API access in untrusted environments. - Never expose the Docker socket to a public network. -#### Docker socket A/B attack test - -Run this test after Docker, Compose, user/group, or runner-launch changes. It -checks whether an authenticated low-privilege web user can turn task submission -into Docker daemon control. The expected secure result is not just "the app -works"; it is that HTTP users can only submit normal FASTA tasks, cannot reach -Docker API routes, cannot choose Docker image/command/mounts, and cannot use a -cookie-only CSRF request for task writes. - -Use an isolated staging instance. Do not run destructive Docker escape payloads -on a host that stores real data. - -Prepare: - -```bash -export BASE_URL=http://127.0.0.1:8080 -export ADMIN_USER=admin -export ADMIN_PASS='' - -docker ps --format '{{.ID}} {{.Names}} {{.Image}} {{.Ports}}' -docker inspect server-web-1 server-worker-1 \ - --format '{{.Name}} User={{.Config.User}} Groups={{json .HostConfig.GroupAdd}} Mounts={{json .Mounts}}' -docker exec server-web-1 id -docker exec server-worker-1 id -docker exec server-web-1 ls -ln /var/run/docker.sock -docker exec server-worker-1 ls -ln /var/run/docker.sock -``` - -A result: the socket is not accessible from web/worker. This proves the live -instance is not Docker-root-escape-capable, but Docker-backed task execution -will fail until permissions are fixed. Expected evidence: - -```text -/var/run/docker.sock -> srw-rw---- 1 0 0 ... -docker.from_env() from web/worker -> PermissionError(13, 'Permission denied') -submitted task -> failed with "Docker daemon unavailable" or PermissionError -``` - -B result: the socket is accessible from web/worker. This is operationally -required for runner containers, but it is host-root-equivalent if arbitrary code -ever runs in web/worker. In this case the test must prove the HTTP contract does -not expose Docker control. Expected evidence: - -```text -limited user /api/auth/admin/users -> 403 -cookie-only POST /PSSM_GREMLIN/api/post -> 403 Bearer token required -POST /containers/create and GET /containers/json on the Flask port -> 404 -path-like FASTA filename upload -> normal task only, sanitized filename, no host-path mount selection -server code fixes image, command, user, environment, and mounts from server config -``` - -The following probe creates a temporary non-admin user, checks the CSRF gate, -submits a harmless FASTA with a path-like filename, probes Docker-like HTTP -routes, and verifies API-key privilege limits: - -```bash -python - <<'PY' -import json, re, uuid -import os -import requests - -base = os.environ["BASE_URL"] -admin_user = os.environ["ADMIN_USER"] -admin_pass = os.environ["ADMIN_PASS"] - -def show(label, resp): - ctype = resp.headers.get("content-type", "") - body = resp.text[:220].replace("\n", " ") - if "application/json" in ctype: - try: - data = resp.json() - body = {k: ("[redacted]" if k in {"token", "api_key"} else v) for k, v in data.items()} - except Exception: - pass - print(json.dumps({"label": label, "status": resp.status_code, "location": resp.headers.get("Location"), "body": body}, sort_keys=True)) - -admin = requests.Session() -r = admin.post(f"{base}/PSSM_GREMLIN/api/auth/login", json={"username": admin_user, "password": admin_pass}, timeout=10) -show("admin_login", r) -r.raise_for_status() -admin_h = {"Authorization": "Bearer " + r.json()["token"]} - -name = "limited_" + uuid.uuid4().hex[:8] -pw = "Pw_" + uuid.uuid4().hex + "!aA1" -r = admin.post( - f"{base}/PSSM_GREMLIN/api/auth/admin/users", - headers=admin_h, - json={"username": name, "email": name + "@audit.local", "password": pw, "affiliation": "audit", "is_admin": False}, - timeout=10, -) -show("create_limited", r) - -limited = requests.Session() -r = limited.post(f"{base}/PSSM_GREMLIN/api/auth/login", json={"username": name, "password": pw}, timeout=10) -show("limited_login", r) -r.raise_for_status() -limited_h = {"Authorization": "Bearer " + r.json()["token"]} - -r = limited.get(f"{base}/PSSM_GREMLIN/api/auth/admin/users", headers=limited_h, timeout=10) -show("limited_admin_list", r) - -files = {"file": ("../../../../var/run/docker.sock.fasta", b">audit\nACDEFGHIKLMNPQRSTVWY\n")} -r = limited.post(f"{base}/PSSM_GREMLIN/api/post", files=files, allow_redirects=False, timeout=10) -show("cookie_only_upload", r) - -files = {"file": ("../../../../--privileged--var-run-docker-sock.fasta", b">audit\nACDEFGHIKLMNPQRSTVWY\n")} -r = limited.post(f"{base}/PSSM_GREMLIN/api/post", files=files, headers=limited_h, allow_redirects=False, timeout=15) -show("bearer_upload_pathlike_name", r) -loc = r.headers.get("Location") or "" -match = re.search(r"/running/([0-9a-f]{32})", loc) -if match: - md5 = match.group(1) - show("running_after_upload", limited.get(f"{base}/PSSM_GREMLIN/api/running/{md5}", headers=limited_h, timeout=10)) - show("delete_uploaded_task", limited.delete(f"{base}/PSSM_GREMLIN/api/delete/{md5}", headers=limited_h, timeout=10)) - -for method, path in [ - ("get", "/containers/json"), - ("get", "/version"), - ("post", "/containers/create"), - ("get", "/PSSM_GREMLIN/api/docker"), - ("post", "/PSSM_GREMLIN/api/docker/run"), -]: - resp = getattr(limited, method)(base + path, headers=limited_h, timeout=10) - show(f"probe_{method}_{path}", resp) - -r = limited.post(f"{base}/PSSM_GREMLIN/api/auth/me/api-key", headers=limited_h, timeout=10) -show("limited_generate_api_key", r) -if r.status_code == 201: - api_h = {"X-API-Key": r.json()["api_key"]} - clean = requests.Session() - show("apikey_admin_list", clean.get(f"{base}/PSSM_GREMLIN/api/auth/admin/users", headers=api_h, timeout=10)) - files = {"file": ("api-key-task.fasta", b">audit\nACDEFGHIKLMNPQRSTVWY\n")} - show("apikey_upload_task", clean.post(f"{base}/PSSM_GREMLIN/api/post", files=files, headers=api_h, allow_redirects=False, timeout=15)) -PY -``` - -Treat a failure in the expected status codes above as a security regression. -If B can access Docker and an HTTP user can influence image, command, privileged -mode, bind source paths, or socket mounts, this is a real host-root escape path. - -#### Admin self-lockout check - -Run this after admin user-management changes. The expected result is that an -admin cannot remove their own access, either through a single-user action or a -batch action. - -Expected results: - -```text -self PUT user_status=banned -> 400 Administrators cannot ban their own account -self DELETE -> 400 Administrators cannot delete their own account -batch disable [self, other] -> count 1, self stays active, other becomes banned -batch delete [self, other] -> count 1, self remains undeleted, other is deleted -``` - -#### Banned-user authentication check - -Run this after account-status or token-auth changes. The check creates a -temporary user, confirms it can log in before the ban, bans it, and verifies -that new login, old Bearer token, and pre-existing API key are all rejected. - -Expected results: - -```text -pre_ban_login -> 200 -pre_ban_api_key -> 201 -ban_user -> 200 -post_ban_login -> 403 Account has been suspended -post_ban_old_bearer_me -> 401 Authentication required -post_ban_api_key_me -> 401 Authentication required -``` - -Command: - -```bash -python - <<'PY' -import json, os, uuid -import requests - -base = os.environ["BASE_URL"] -admin_user = os.environ["ADMIN_USER"] -admin_pass = os.environ["ADMIN_PASS"] - -def show(label, resp): - ctype = resp.headers.get("content-type", "") - body = resp.text[:180].replace("\n", " ") - if "application/json" in ctype: - try: - data = resp.json() - body = {k: ("[redacted]" if k in {"token", "api_key"} else v) for k, v in data.items()} - except Exception: - pass - print(json.dumps({"label": label, "status": resp.status_code, "body": body}, sort_keys=True)) - -admin = requests.Session() -r = admin.post(f"{base}/PSSM_GREMLIN/api/auth/login", json={"username": admin_user, "password": admin_pass}, timeout=10) -show("admin_login", r) -r.raise_for_status() -admin_h = {"Authorization": "Bearer " + r.json()["token"]} - -name = "bancheck_" + uuid.uuid4().hex[:8] -pw = "Pw_" + uuid.uuid4().hex + "!aA1" -r = admin.post( - f"{base}/PSSM_GREMLIN/api/auth/admin/users", - headers=admin_h, - json={"username": name, "email": name + "@audit.local", "password": pw, "affiliation": "audit", "is_admin": False}, - timeout=10, -) -show("create_user", r) - -r = admin.get(f"{base}/PSSM_GREMLIN/api/auth/admin/users", headers=admin_h, timeout=10) -r.raise_for_status() -user_id = next(u["id"] for u in r.json()["users"] if u["username"] == name) - -limited = requests.Session() -r = limited.post(f"{base}/PSSM_GREMLIN/api/auth/login", json={"username": name, "password": pw}, timeout=10) -show("pre_ban_login", r) -r.raise_for_status() -limited_h = {"Authorization": "Bearer " + r.json()["token"]} - -r = limited.post(f"{base}/PSSM_GREMLIN/api/auth/me/api-key", headers=limited_h, timeout=10) -show("pre_ban_api_key", r) -api_key = r.json().get("api_key") if r.status_code == 201 else None - -r = admin.put(f"{base}/PSSM_GREMLIN/api/auth/admin/users/{user_id}", headers=admin_h, json={"user_status": "banned"}, timeout=10) -show("ban_user", r) -show("post_ban_login", requests.post(f"{base}/PSSM_GREMLIN/api/auth/login", json={"username": name, "password": pw}, timeout=10)) -show("post_ban_old_bearer_me", requests.get(f"{base}/PSSM_GREMLIN/api/auth/me", headers=limited_h, timeout=10)) -if api_key: - show("post_ban_api_key_me", requests.get(f"{base}/PSSM_GREMLIN/api/auth/me", headers={"X-API-Key": api_key}, timeout=10)) -PY -``` +Security regression checks for Docker socket exposure, admin self-lockout, +banned users, and login throttling are maintained in +`docs/dev-guide/server.md#security-validation`. ### Authentication -- Tokens are signed with `itsdangerous.URLSafeTimedSerializer` (HMAC-SHA1). -- Gunicorn uses `--preload` so the auth secret key is generated once in the - arbiter; all workers share the same key and tokens validate consistently. - Set `AUTH_SECRET_KEY` to a fixed, high-entropy value in production; otherwise tokens are lost on restart. - Browser page navigations use an `HttpOnly`/`SameSite=Lax` cookie; JavaScript cannot read it, so logout requires the server endpoint (`POST /api/auth/logout`). - Rate limiting: 5 login attempts/minute/IP, 3 registrations/hour/IP. - All state-changing endpoints require a valid Bearer token or API key. - API keys have restricted privileges (task operations only) — Bearer tokens are required for profile changes and admin actions. -- CSRF is mitigated: state-changing endpoints require a Bearer token in the - `Authorization` header (browser same-origin policy prevents cross-origin - requests from setting custom headers). The `HttpOnly` cookie is used for - read-only page navigations only. +- Cookie-only writes are rejected; state-changing API calls require a Bearer + token or API key. ### Redis @@ -715,7 +466,23 @@ PY - Regularly back up sqlite and result archives. - If a task is deleted, result artifacts are removed, but the sqlite record remains for audit. -## 11. Troubleshooting +## 11. Local Development + +```bash +# Install in editable mode with test dependencies +pip install -e "server/[test]" + +# Run non-Docker tests from the repo root +pytest server/tests/ -v -k "not Docker and not docker" + +# Run the server directly without Docker +python -m pssm_gremlin_server.pssm_gremlin +``` + +Full test and security validation guidance is maintained in +`docs/dev-guide/server.md`. + +## 12. Troubleshooting ### Network issues diff --git a/server/pssm_gremlin_server/auth.py b/server/pssm_gremlin_server/auth.py index 684418179..0a1d6dafe 100644 --- a/server/pssm_gremlin_server/auth.py +++ b/server/pssm_gremlin_server/auth.py @@ -23,7 +23,7 @@ from typing import Any import sqlalchemy as sa -from flask import current_app, g, jsonify, redirect, render_template, request, url_for +from flask import current_app, g, jsonify, redirect, request, url_for from itsdangerous import BadSignature, SignatureExpired, URLSafeTimedSerializer from werkzeug.security import check_password_hash, generate_password_hash @@ -471,10 +471,13 @@ def load_current_user() -> dict[str, Any] | None: payload = validate_token(token) if payload is not None: user = db.get_user(payload["uid"]) - if user is not None and _is_account_blocked(user) is None: - if payload.get("ver", 0) == user.get("token_version", 0): - g.auth_method = "bearer" - return user + if ( + user is not None + and _is_account_blocked(user) is None + and payload.get("ver", 0) == user.get("token_version", 0) + ): + g.auth_method = "bearer" + return user # 2. Cookie — browser page navigations after login (read-only; CSRF-prone) # Guest accounts ARE permitted to use cookie-based web access. token = request.cookies.get("auth_token") @@ -482,10 +485,13 @@ def load_current_user() -> dict[str, Any] | None: payload = validate_token(token) if payload is not None: user = db.get_user(payload["uid"]) - if user is not None and _is_account_blocked(user) is None: - if payload.get("ver", 0) == user.get("token_version", 0): - g.auth_method = "cookie" - return user + if ( + user is not None + and _is_account_blocked(user) is None + and payload.get("ver", 0) == user.get("token_version", 0) + ): + g.auth_method = "cookie" + return user # 3. API key (programmatic access — restricted privileges) # Guest accounts are not permitted to use API keys. diff --git a/server/pssm_gremlin_server/pssm_gremlin.py b/server/pssm_gremlin_server/pssm_gremlin.py index e58086692..e8912b417 100644 --- a/server/pssm_gremlin_server/pssm_gremlin.py +++ b/server/pssm_gremlin_server/pssm_gremlin.py @@ -2,13 +2,11 @@ # Distributed under the terms of the GNU General Public License v3.0. # SPDX-License-Identifier: GPL-3.0-only - -#! /mnt/data/envs/conda_env/envs/REvoDesign/bin/python - from __future__ import annotations import grp import hashlib +import importlib import json import logging import os @@ -16,17 +14,17 @@ import re import shutil import signal -import tempfile import time from dataclasses import dataclass from datetime import datetime +from pathlib import Path from typing import Any import docker from celery import Celery from celery.result import AsyncResult from docker import types -from flask import Flask, current_app, g, jsonify, request +from flask import Flask, g, jsonify, request from pssm_gremlin_server.db import TaskDatabase from sqlalchemy.exc import IntegrityError from werkzeug.utils import secure_filename @@ -39,13 +37,15 @@ from pssm_gremlin_server.auth import UserDatabase # noqa: E402 from pssm_gremlin_server.auth import _env_bool # noqa: E402 -from pssm_gremlin_server.auth import _env_int # noqa: E402 from pssm_gremlin_server.auth import _env_str # noqa: E402 from pssm_gremlin_server.auth import send_admin_digest # noqa: E402 THIS_FILE = os.path.abspath(__file__) THIS_DIR = os.path.dirname(THIS_FILE) TEMPLATE_IMAGE_DIR = os.path.join(THIS_DIR, "templates", "images") +DEFAULT_SERVER_DIR = os.path.abspath(os.path.join(os.getcwd(), "pssm_gremlin_data")) +DEFAULT_UNIREF30_DB = os.path.join(DEFAULT_SERVER_DIR, "db", "uniref30", "UniRef30_2022_02") +DEFAULT_UNIREF90_DB = os.path.join(DEFAULT_SERVER_DIR, "db", "uniref90", "uniref90") app = Flask(__name__, template_folder="./templates") app.config["MAX_CONTENT_LENGTH"] = 16 * 1024 * 1024 # 16 MiB upload limit @@ -172,7 +172,7 @@ class GremlinConfig: @classmethod def from_env(cls) -> GremlinConfig: - server_dir = _env_path("SERVER_DIR", "/mnt/data/yinying/server/") + server_dir = _env_path("SERVER_DIR", DEFAULT_SERVER_DIR) upload_folder = os.path.join(server_dir, "upload") results_folder = os.path.join(server_dir, "results") return cls( @@ -182,11 +182,8 @@ def from_env(cls) -> GremlinConfig: db_path=_env_path("DB_PATH", os.path.join(server_dir, "pssm_gremlin.sqlite3")), docker_image=os.environ.get("RUNNER_IMAGE", "revodesign-pssm-gremlin"), docker_user=_resolve_docker_user(), - uniref30_db=_env_path( - "DB_UNIREF30", - "/mnt/db/uniref30_uc30/UniRef30_2022_02/UniRef30_2022_02", - ), - uniref90_db=_env_path("DB_UNIREF90", "/mnt/db/uniref90/uniref90"), + uniref30_db=_env_path("DB_UNIREF30", DEFAULT_UNIREF30_DB), + uniref90_db=_env_path("DB_UNIREF90", DEFAULT_UNIREF90_DB), nproc=_env_int("NPROC", 16), maxmem=_env_int("MAXMEM", 64), port=_env_int("PORT", 8080), @@ -233,15 +230,14 @@ def from_env(cls) -> GremlinConfig: # notified so each user appears only once. _admin_digest_minutes = _env_int("ADMIN_NEW_USER_INFORM", 0) if _admin_digest_minutes > 0 and _env_str("ADMIN_NOTIFY_EMAIL", ""): + import secrets import threading _digest_lock = os.path.join(_env_str("SERVER_DIR", os.getcwd()), ".admin_digest.lock") def _digest_loop() -> None: - import random - while True: - _jitter = random.uniform(-5, 5) # spread worker wake-ups + _jitter = secrets.randbelow(10001) / 1000 - 5 # spread worker wake-ups time.sleep(_admin_digest_minutes * 60 + _jitter) try: # ponytail: file lock so only one worker sends the digest @@ -253,7 +249,7 @@ def _digest_loop() -> None: with open(_digest_lock, "w") as _lock_fh: fcntl.flock(_lock_fh.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB) send_admin_digest() - except (BlockingIOError, OSError): + except OSError: pass # another worker has the lock — try next cycle except Exception: logging.exception("Admin digest thread failed") @@ -290,7 +286,7 @@ def _digest_loop() -> None: def _is_binary_file(path: str) -> bool: try: - with open(path, "rb") as f: + with Path(path).open("rb") as f: chunk = f.read(4096) except OSError: return True @@ -308,7 +304,7 @@ def _is_binary_file(path: str) -> bool: def _is_fasta_content(path: str) -> bool: """Return True if *path* looks like a FASTA file (first non-blank line starts with '>').""" try: - with open(path, encoding="utf-8", errors="replace") as f: + with Path(path).open(encoding="utf-8", errors="replace") as f: for line in f: stripped = line.strip() if not stripped: @@ -542,11 +538,8 @@ def _build_running_trace(task: dict[str, Any]) -> str: return "\n".join(traced_lines) -try: - _ROOT_MOUNT_DIRECTORY = f"/home/{os.getlogin()}" -except BaseException: - _ROOT_MOUNT_DIRECTORY = os.path.abspath(tempfile.gettempdir()) - os.makedirs(_ROOT_MOUNT_DIRECTORY, exist_ok=True) +_ROOT_MOUNT_DIRECTORY = _env_path("RUNNER_HOST_ROOT", os.path.dirname(CONFIG.server_dir)) +os.makedirs(_ROOT_MOUNT_DIRECTORY, exist_ok=True) def _is_admin_user(username: str | None = None) -> bool: @@ -660,7 +653,7 @@ def _create_mount(mount_name: str, path: str, read_only=True) -> tuple[types.Mou target_path = os.path.join(_ROOT_MOUNT_DIRECTORY, mount_name) if not read_only: - logging.warning(f"{mount_name} is not read-only!") + logging.warning("%s is not read-only!", mount_name) if os.path.isdir(path): source_path = path @@ -670,7 +663,7 @@ def _create_mount(mount_name: str, path: str, read_only=True) -> tuple[types.Mou mounted_path = os.path.join(target_path, os.path.basename(path)) if not os.path.exists(source_path): os.makedirs(source_path) - logging.info(f"Mounting {source_path} -> {target_path}") + logging.info("Mounting %s -> %s", source_path, target_path) mount = types.Mount( target=str(target_path), source=str(source_path), @@ -782,8 +775,6 @@ def run_pssm_gremlin_in_docker(fasta_path, output_dir, docker_client=None, stage except docker.errors.DockerException: pass - return - def _pack_results_archive(task: dict) -> None: zip_filename = _task_zip_path(task) @@ -818,8 +809,7 @@ def _pack_failed_results_archive(task: dict, error: Any) -> None: def format_times(timestamp): if timestamp: return datetime.fromtimestamp(timestamp).strftime("%Y-%m-%d %H:%M:%S") - else: - return None + return None def format_walltime(seconds: Any) -> str: @@ -961,7 +951,8 @@ def _on_stage_change(stage: str) -> None: # Register HTTP routes (imported late to avoid circular imports — routes.py # needs ``app`` and helpers that are only available after this module loads). # --------------------------------------------------------------------------- -from pssm_gremlin_server import routes # noqa: E402, F401 +importlib.import_module("pssm_gremlin_server.routes") if __name__ == "__main__": - app.run(host="0.0.0.0", port=CONFIG.port) # nosec B104: containerized server, binding to all interfaces by design + # Containerized server binds to all interfaces by design. + app.run(host="0.0.0.0", port=CONFIG.port) # nosec B104 diff --git a/server/pssm_gremlin_server/routes.py b/server/pssm_gremlin_server/routes.py index 8bf6f2f8b..b564170d0 100644 --- a/server/pssm_gremlin_server/routes.py +++ b/server/pssm_gremlin_server/routes.py @@ -17,6 +17,7 @@ import os import shutil import time +from typing import Any from celery.result import AsyncResult from flask import current_app, g, jsonify, redirect, render_template, request, send_from_directory, url_for @@ -735,23 +736,21 @@ def auth_forgot_password(): return jsonify({"message": "If that email is registered, a reset link has been sent."}), 200 -@app.route("/PSSM_GREMLIN/reset_password", methods=["GET", "POST"]) -def auth_reset_password(): - """Password-reset page. +@app.route("/PSSM_GREMLIN/reset_password", methods=["GET"]) +def auth_reset_password_page(): + """Render the password-reset page for a valid reset token.""" + token = request.args.get("c", "").strip() + if not token: + return render_template("error.html", code=400, message="Missing reset token."), 400 + user_id = validate_reset_token(token) + if user_id is None: + return render_template("error.html", code=400, message="Invalid or expired reset token."), 400 + return render_template("reset-password.html", token=token), 200 - ``GET`` — renders the new-password form (requires ``?c=`` token). - ``POST`` — sets the new password. - """ - if request.method == "GET": - token = request.args.get("c", "").strip() - if not token: - return render_template("error.html", code=400, message="Missing reset token."), 400 - user_id = validate_reset_token(token) - if user_id is None: - return render_template("error.html", code=400, message="Invalid or expired reset token."), 400 - return render_template("reset-password.html", token=token), 200 - - # POST + +@app.route("/PSSM_GREMLIN/reset_password", methods=["POST"]) +def auth_reset_password(): + """Set a new password using a password-reset token.""" req = _parse_body(ResetPasswordRequest) if isinstance(req, tuple): return req @@ -1069,27 +1068,29 @@ def auth_revoke_api_key(): return jsonify({"message": "API key revoked"}), 200 -@app.route("/PSSM_GREMLIN/api/auth/admin/users", methods=["GET", "POST"]) +@app.route("/PSSM_GREMLIN/api/auth/admin/users", methods=["GET"]) @login_required def admin_users(): - """Admin-only user management. - - ``GET`` — list all users (safe fields only). - ``POST`` — create a new user (auto-verified, immediately active). - """ + """Admin-only user listing with safe fields only.""" if _blocked := require_admin(): return _blocked db = _get_user_db() + users = db.list_users() + safe = [UserResponse.model_validate(u).model_dump() for u in users] + return jsonify({"users": safe}), 200 - if request.method == "GET": - users = db.list_users() - safe = [UserResponse.model_validate(u).model_dump() for u in users] - return jsonify({"users": safe}), 200 - # POST — state-changing, require Bearer token for CSRF protection +@app.route("/PSSM_GREMLIN/api/auth/admin/users", methods=["POST"]) +@login_required +def admin_create_user(): + """Admin-only user creation for pre-approved accounts.""" + if _blocked := require_admin(): + return _blocked if _blocked := require_bearer_auth(): return _blocked + + db = _get_user_db() req = _parse_body(AdminCreateUserRequest) if isinstance(req, tuple): return req diff --git a/server/scripts/GREMLIN_TFv1.py b/server/scripts/GREMLIN_TFv1.py index 1b220eece..7360d8ff2 100644 --- a/server/scripts/GREMLIN_TFv1.py +++ b/server/scripts/GREMLIN_TFv1.py @@ -3,7 +3,7 @@ # SPDX-License-Identifier: GPL-3.0-only -#!/usr/bin/env python +# Original script marker: /usr/bin/env python # # GREMLIN_TF # @@ -12,7 +12,9 @@ import os import pathlib -import pickle # nosec B403: MRF model serialization, trusted data + +# MRF model serialization uses trusted data. +import pickle # nosec B403 import sys import matplotlib @@ -82,17 +84,16 @@ def parse_fasta(filename, limit=-1): """function to parse fasta""" header = [] sequence = [] - lines = open(filename) - for line in lines: - line = line.rstrip() - if line[0] == ">": - if len(header) == limit: - break - header.append(line[1:]) - sequence.append([]) - else: - sequence[-1].append(line) - lines.close() + with open(filename) as lines: + for line in lines: + line = line.rstrip() + if line[0] == ">": + if len(header) == limit: + break + header.append(line[1:]) + sequence.append([]) + else: + sequence[-1].append(line) sequence = ["".join(seq) for seq in sequence] return np.array(header), np.array(sequence) @@ -326,7 +327,14 @@ def get_loss(): # ## Explore the contact map # ### Contact prediction: # -# For contact prediction, the W matrix is reduced from LxLx21x21 to LxL matrix (by taking the L2norm for each of the 20x20). In the code below, you can access this as mtx["raw"]. Further correction (average product correction) is then performed to the mtx["raw"] to remove the effects of entropy, mtx["apc"]. The relative ranking of mtx["apc"] is used to assess importance. When there are enough effective sequences (>1000), we find that the top 1.0L contacts are ~90% accurate! When the number of effective sequences is lower, NN can help clean noise and fill in missing contacts. +# For contact prediction, the W matrix is reduced from LxLx21x21 to LxL matrix +# by taking the L2 norm for each of the 20x20 slices. In the code below, you can +# access this as mtx["raw"]. Further correction (average product correction) is +# then performed to remove entropy effects, producing mtx["apc"]. The relative +# ranking of mtx["apc"] is used to assess importance. When there are enough +# effective sequences (>1000), we find that the top 1.0L contacts are ~90% +# accurate. When the number of effective sequences is lower, NN can help clean +# noise and fill in missing contacts. # # ## Functions for extracting contacts from MRF @@ -365,7 +373,8 @@ def plot_mtx(mtx, key="zscore", vmin=1, vmax=3): # save mtx file -pickle.dump(mrf, open(f"{pth}/{instance}.GREMLIN.mrf.pkl", "wb")) +with open(f"{pth}/{instance}.GREMLIN.mrf.pkl", "wb") as mrf_file: + pickle.dump(mrf, mrf_file) mtx = get_mtx(mrf) plot_mtx(mtx) diff --git a/server/scripts/fasta_lower_char_rm.py b/server/scripts/fasta_lower_char_rm.py index c60bb8485..2750ab0c5 100644 --- a/server/scripts/fasta_lower_char_rm.py +++ b/server/scripts/fasta_lower_char_rm.py @@ -3,7 +3,7 @@ # SPDX-License-Identifier: GPL-3.0-only -#! python +# Original script marker: python # created by Yao Yin Ying # for hhblits a3m alignment file treatment diff --git a/server/scripts/make_paired_MSA_simple.py b/server/scripts/make_paired_MSA_simple.py index 6292a5a67..5093474d7 100644 --- a/server/scripts/make_paired_MSA_simple.py +++ b/server/scripts/make_paired_MSA_simple.py @@ -34,13 +34,38 @@ def calc_seqID(query, cand): def read_a3m(fn): # read sequences in a3m file # only take one (having the highest seqID to query) per each taxID + if fn.split(".")[-1] == "gz": + with gzip.open(fn, "rt") as fp: + query, tmp = _read_a3m_taxa(fp) + else: + with open(fn) as fp: + query, tmp = _read_a3m_taxa(fp) + + query_in_num = seq2number(query) + a3m = {} + for TaxID in tmp: + if len(tmp[TaxID]) < 1: + continue + if len(tmp[TaxID]) < 2: + a3m[TaxID] = tmp[TaxID][0] + continue + # Get the best sequence only + score_s = list() + for _seqID, seq in tmp[TaxID]: + seq_in_num = seq2number(seq) + score = calc_seqID(query_in_num, seq_in_num) + score_s.append(score) + # + idx = np.argmax(score_s) + a3m[TaxID] = tmp[TaxID][idx] + + return query, a3m + + +def _read_a3m_taxa(fp): is_first = True is_ignore = False tmp = {} - if fn.split(".")[-1] == "gz": - fp = gzip.open(fn, "rt") - else: - fp = open(fn) for line in fp: if line[0] == ">": @@ -51,7 +76,7 @@ def read_a3m(fn): try: idx = line.index("OX") is_ignore = False - except BaseException: + except ValueError: is_ignore = True continue TaxID = line[idx:].split()[0].split("=")[-1] @@ -66,25 +91,7 @@ def read_a3m(fn): else: tmp[TaxID].append((seqID, line.strip())) - query_in_num = seq2number(query) - a3m = {} - for TaxID in tmp: - if len(tmp[TaxID]) < 1: - continue - if len(tmp[TaxID]) < 2: - a3m[TaxID] = tmp[TaxID][0] - continue - # Get the best sequence only - score_s = list() - for seqID, seq in tmp[TaxID]: - seq_in_num = seq2number(seq) - score = calc_seqID(query_in_num, seq_in_num) - score_s.append(score) - # - idx = np.argmax(score_s) - a3m[TaxID] = tmp[TaxID][idx] - - return query, a3m + return query, tmp if len(sys.argv) == 1: diff --git a/server/tests/test_config.py b/server/tests/test_config.py index 0923a4c86..96e11f4bd 100644 --- a/server/tests/test_config.py +++ b/server/tests/test_config.py @@ -5,14 +5,7 @@ from __future__ import annotations import pytest -from conftest import ( - _admin_client_auth, - _extract_md5, - _insert_pending_task, - _load_pssm_module, - _test_client_auth, - _upsert_task_for_user, -) +from conftest import _load_pssm_module # config tests # ================================================================== @@ -42,6 +35,44 @@ def test_pssm_config_uses_named_runner_identity(monkeypatch, tmp_path): assert module.CONFIG.docker_user == "revodesign:revodesign_appgroup" +def test_pssm_config_defaults_are_not_cluster_paths(monkeypatch, tmp_path): + monkeypatch.chdir(tmp_path) + module = _load_pssm_module( + monkeypatch, + tmp_path, + extra_env={ + "RUNNER_UID": "1234", + "RUNNER_GID": "5678", + "SERVER_DIR": None, + "DB_PATH": None, + "DB_UNIREF30": None, + "DB_UNIREF90": None, + "RUNNER_HOST_ROOT": None, + }, + ) + + assert module.CONFIG.server_dir.endswith("pssm_gremlin_data") + assert not module.CONFIG.server_dir.startswith("/mnt/") + assert not module.CONFIG.uniref30_db.startswith("/mnt/") + assert not module.CONFIG.uniref90_db.startswith("/mnt/") + assert module._ROOT_MOUNT_DIRECTORY == str(module.Path(module.CONFIG.server_dir).parent) + + +def test_pssm_config_uses_runner_host_root_override(monkeypatch, tmp_path): + host_root = tmp_path / "runner_host_root" + module = _load_pssm_module( + monkeypatch, + tmp_path, + extra_env={ + "RUNNER_UID": "1234", + "RUNNER_GID": "5678", + "RUNNER_HOST_ROOT": str(host_root), + }, + ) + + assert module._ROOT_MOUNT_DIRECTORY == str(host_root) + + def test_pssm_config_requires_runner_identity(monkeypatch, tmp_path): with pytest.raises(RuntimeError): _load_pssm_module(monkeypatch, tmp_path) diff --git a/server/tests/test_runner_script_static.py b/server/tests/test_runner_script_static.py new file mode 100644 index 000000000..2089f6926 --- /dev/null +++ b/server/tests/test_runner_script_static.py @@ -0,0 +1,27 @@ +# Copyright (c) 2026 The REvoDesign Developers. +# Distributed under the terms of the GNU General Public License v3.0. +# SPDX-License-Identifier: GPL-3.0-only + +from __future__ import annotations + +from pathlib import Path + +RUNNER_SCRIPT = Path(__file__).resolve().parents[1] / "REvoDesign_PSSM_GREMLIN.sh" + + +def test_runner_script_does_not_eval_user_controlled_commands(): + script = RUNNER_SCRIPT.read_text() + + assert 'eval "$cmd"' not in script + assert "eval $cmd" not in script + assert "bash -c" not in script + assert "sh -c" not in script + + +def test_runner_script_executes_pipeline_commands_as_arrays(): + script = RUNNER_SCRIPT.read_text() + + for command in ("hhblits", "hhfilter", "GREMLIN_TFv1.py", "psiblast"): + assert command in script + + assert '"${cmd[@]}"' in script diff --git a/server/tests/test_tasks.py b/server/tests/test_tasks.py index f7e0d8cdb..7922c9fc0 100644 --- a/server/tests/test_tasks.py +++ b/server/tests/test_tasks.py @@ -13,9 +13,8 @@ from pathlib import Path import docker -import pytest +import requests from conftest import ( - _admin_client_auth, _extract_md5, _insert_pending_task, _load_pssm_module, @@ -358,7 +357,10 @@ def test_upload_records_headers_and_local_user(monkeypatch, tmp_path): "RUNNER_GID": "5678", }, ) - monkeypatch.setattr(module.routes, "_local_user_identity", lambda: "pytest:staff-1000:20") + upload_file = module.app.view_functions["upload_file"] + while "_local_user_identity" not in upload_file.__globals__: + upload_file = upload_file.__wrapped__ + monkeypatch.setitem(upload_file.__globals__, "_local_user_identity", lambda: "pytest:staff-1000:20") class _DummyAsyncResult: id = "celery-test-id" @@ -482,7 +484,7 @@ def test_dashboard_masks_host_file_paths_on_read_errors(monkeypatch, tmp_path): md5sum = uuid.uuid4().hex result_dir = tmp_path / "result" result_dir.mkdir(parents=True, exist_ok=True) - leaked_host_path = "/Users/yyy/Documents/protein_design/REvoDesign/playground/server_test/upload/2KL8.fasta" + leaked_host_path = "/home/server-user/REvoDesign/playground/server_test/upload/2KL8.fasta" _upsert_task_for_user( module, @@ -517,7 +519,7 @@ def test_failed_status_masks_host_paths_in_api_error(monkeypatch, tmp_path): md5sum = uuid.uuid4().hex result_dir = tmp_path / "result" result_dir.mkdir(parents=True, exist_ok=True) - leaked_host_path = "/Users/yyy/Documents/protein_design/REvoDesign/playground/server_test/upload/2KL8.fasta" + leaked_host_path = "/home/server-user/REvoDesign/playground/server_test/upload/2KL8.fasta" _upsert_task_for_user( module, @@ -538,7 +540,7 @@ def test_failed_status_masks_host_paths_in_api_error(monkeypatch, tmp_path): payload = response.get_json() assert payload["status"] == "failed" assert "/srv/REvoDesign/PSSM_GREMLIN/upload/2KL8.fasta" in payload["error"] - assert "/Users/yyy/Documents/protein_design/REvoDesign" not in payload["error"] + assert "/home/server-user/REvoDesign" not in payload["error"] def test_private_dashboard_blocks_non_owner_access(monkeypatch, tmp_path): diff --git a/src/REvoDesign/Qt/ui_runtime_loader.py b/src/REvoDesign/Qt/ui_runtime_loader.py index 367f6a546..537381ba5 100644 --- a/src/REvoDesign/Qt/ui_runtime_loader.py +++ b/src/REvoDesign/Qt/ui_runtime_loader.py @@ -29,7 +29,7 @@ def _load_pyqt_uic_module() -> Any: def _language_change_event_type(): event_type_owner = getattr(QtCore.QEvent, "Type", QtCore.QEvent) - return getattr(event_type_owner, "LanguageChange") + return event_type_owner.LanguageChange def is_language_change_event(event: QtCore.QEvent) -> bool: diff --git a/src/REvoDesign/REvoDesign.py b/src/REvoDesign/REvoDesign.py index 324ae5a8b..75425fe4d 100644 --- a/src/REvoDesign/REvoDesign.py +++ b/src/REvoDesign/REvoDesign.py @@ -11,7 +11,6 @@ import gc import os import tempfile -import traceback import warnings # using partial module to reduce duplicate code. @@ -104,7 +103,8 @@ def __init__( warnings.warn( issues.DisabledFunctionWarning( "Teamwork is disabled because the active PyMOL Qt backend does not provide QtWebSockets." - ) + ), + stacklevel=2, ) self.teamwork_enabled = False @@ -147,8 +147,6 @@ def set_working_directory(self, new_dir: str | None = None): return self.PWD = PWD - os.chdir(self.PWD) - self.bus.set_value("work_dir", os.path.abspath(self.PWD)) def run_plugin_gui(self): @@ -210,7 +208,8 @@ def reinitialize(self, delete: bool = False): warnings.warn( issues.ConflictWarning( "Reinitialized with default configuration. " "Restart REvoDesign to take effort." - ) + ), + stacklevel=2, ) def __del__(self): @@ -661,7 +660,9 @@ def reload_molecule_info(self): cmd.reinitialize() cmd.load(self.temperal_session) else: - warnings.warn(issues.EmptySessionWarning("Current session is empty! Please load one PDB/PSE/PZE!")) + warnings.warn( + issues.EmptySessionWarning("Current session is empty! Please load one PDB/PSE/PZE!"), stacklevel=2 + ) new_session_file = self.file_dialog.browse_filename( mode="r", exts=( @@ -671,10 +672,10 @@ def reload_molecule_info(self): ), ) if not new_session_file: - warnings.warn(issues.NoInputWarning("Abored recognizing sessions from input.")) + warnings.warn(issues.NoInputWarning("Abored recognizing sessions from input."), stacklevel=2) return if not os.path.exists(new_session_file): - warnings.warn(issues.NoInputWarning(f"File does not exist: {new_session_file}.")) + warnings.warn(issues.NoInputWarning(f"File does not exist: {new_session_file}."), stacklevel=2) return cmd.reinitialize() @@ -701,13 +702,13 @@ def save_as_a_session(self, cfg_to_pse: str): self.logging.info(f"Output file is set as {output_pse_fn}") self.bus.set_widget_value(cfg_to_pse, output_pse_fn) else: - warnings.warn(issues.NoInputWarning(f"Invalid output path: {output_pse_fn}.")) + warnings.warn(issues.NoInputWarning(f"Invalid output path: {output_pse_fn}."), stacklevel=2) def update_chain_id(self): """Update chain ids for picked molecule.""" molecule = self.bus.get_widget_value("ui.header_panel.input.molecule", str) if not molecule: - warnings.warn(issues.NoInputWarning("No available designable molecule!")) + warnings.warn(issues.NoInputWarning("No available designable molecule!"), stacklevel=2) return chain_ids = find_all_protein_chain_ids_in_protein(molecule) self.designable_sequences = RosettaPyProteinSequence.from_dict( @@ -735,18 +736,23 @@ def find_session_path(self) -> str | None: session_path: str = cmd.get("session_file") if not session_path: - warnings.warn(issues.EmptySessionWarning("Session not found, please use a new session path to save.")) + warnings.warn( + issues.EmptySessionWarning("Session not found, please use a new session path to save."), stacklevel=2 + ) return self.file_dialog.browse_filename(mode="w", exts=(file_extensions.Session,)) if not os.path.exists(session_path): - warnings.warn(issues.NoInputWarning("Invalid session file path, please use a new session path to save.")) + warnings.warn( + issues.NoInputWarning("Invalid session file path, please use a new session path to save."), stacklevel=2 + ) return self.file_dialog.browse_filename(mode="w", exts=(file_extensions.Session,)) if os.path.basename(session_path).startswith("tmp") and session_path.endswith(".pse"): warnings.warn( issues.InvalidSessionWarning( f"Found temperal session path: {session_path}, " "please use a new session path to save." - ) + ), + stacklevel=2, ) return self.file_dialog.browse_filename(mode="w", exts=(file_extensions.Session,)) @@ -761,7 +767,7 @@ def reload_determine_tab_setup(self): """Setup pocket determination""" molecule = self.bus.get_widget_value("ui.header_panel.input.molecule", str) if not molecule: - warnings.warn(issues.NoResultsWarning("No design molecule found.")) + warnings.warn(issues.NoResultsWarning("No design molecule found."), stacklevel=2) return small_molecules = [""] @@ -769,7 +775,7 @@ def reload_determine_tab_setup(self): small_molecules.extend(more_hetatm) self.logging.info(f"Small molecules found: {more_hetatm}") else: - warnings.warn(issues.NoResultsWarning("No small molecule found.")) + warnings.warn(issues.NoResultsWarning("No small molecule found."), stacklevel=2) self.bus.set_widget_value("ui.prepare.input.pocket.substrate", small_molecules) self.bus.set_widget_value("ui.prepare.input.pocket.cofactor", small_molecules) @@ -945,7 +951,7 @@ def update_mutant_table_columns(self): comboBox_group_name = self.bus.get_widget_from_cfg_item("ui.visualize.input.group_name") if not os.path.exists(mut_table_fp): - warnings.warn(issues.NoInputWarning(f"Mutant Table path is not valid: {mut_table_fp}")) + warnings.warn(issues.NoInputWarning(f"Mutant Table path is not valid: {mut_table_fp}"), stacklevel=2) # reset cols to combo boxes to empty for comboBox in [comboBox_best_leaf, comboBox_totalscore]: set_widget_value(comboBox, [""]) @@ -956,7 +962,9 @@ def update_mutant_table_columns(self): mut_table_cols = get_mutant_table_columns(mutfile=mut_table_fp) if not mut_table_cols: - warnings.warn(issues.BadDataWarning(f"Mutant Table column names is not valid: {mut_table_cols}")) + warnings.warn( + issues.BadDataWarning(f"Mutant Table column names is not valid: {mut_table_cols}"), stacklevel=2 + ) return # set cols to combo boxes @@ -1179,7 +1187,7 @@ def multi_mutagenesis_design_auto(self): self.multi_mutagenesis_design_stop_design() self.multi_mutagenesis_design_save_design() except Exception: - traceback.print_exc() + self.logging.exception("Failed to generate multi-mutagenesis designs.") # Tab Interact via GREMLIN def load_gremlin_mrf(self): @@ -1306,9 +1314,8 @@ def toggle_ws_server_mode(self): self.logging.warning("Server is already stopped. Do nothing.") return self.ws_server.stop_server() - except Exception as e: - self.logging.warning(e) - traceback.print_exc() + except Exception: + self.logging.exception("Failed to toggle WebSocket server connection.") self.logging.warning(f'Server status: {"ON" if self.ws_server.is_running else "OFF"}') @@ -1352,9 +1359,8 @@ def toggle_ws_client_connection(self, connect: bool = True): self.ws_client_connect_to_server() else: self.ws_client_disconnect_from_server() - except Exception as e: - self.logging.warning(e) - traceback.print_exc() + except Exception: + self.logging.exception("Failed to toggle WebSocket client connection.") self.logging.warning(f'Client status: {"ON" if self.ws_client.connected else "OFF"}') diff --git a/src/REvoDesign/UI/REvoDesign-PyMOL-entry.ui b/src/REvoDesign/UI/REvoDesign-PyMOL-entry.ui index 5e4741708..bedd70f1f 100644 --- a/src/REvoDesign/UI/REvoDesign-PyMOL-entry.ui +++ b/src/REvoDesign/UI/REvoDesign-PyMOL-entry.ui @@ -151,7 +151,7 @@ - /Users/yyy/Documents/protein_design/REvoDesign + diff --git a/src/REvoDesign/UI/UI.pyproject.user b/src/REvoDesign/UI/UI.pyproject.user deleted file mode 100644 index 82ab870d8..000000000 --- a/src/REvoDesign/UI/UI.pyproject.user +++ /dev/null @@ -1,375 +0,0 @@ - - - - - - EnvironmentId - {0214636d-c5aa-4424-985a-d088db583ed6} - - - ProjectExplorer.Project.ActiveTarget - 0 - - - ProjectExplorer.Project.EditorSettings - - true - false - true - - Cpp - - CppGlobal - - - - QmlJS - - QmlJSGlobal - - - 2 - UTF-8 - false - 4 - false - 80 - true - true - 1 - false - true - false - 0 - true - true - 0 - 8 - true - false - 1 - true - true - true - *.md, *.MD, Makefile - false - true - true - - - - ProjectExplorer.Project.PluginSettings - - - true - false - true - true - true - true - - - 0 - true - - true - true - Builtin.DefaultTidyAndClazy - 5 - true - - - - true - - - - - ProjectExplorer.Project.Target.0 - - Desktop - Desktop (x86-darwin-generic-mach_o-64bit) - Desktop (x86-darwin-generic-mach_o-64bit) - {1e4bf9bd-468b-4836-ba8c-71a0373db4bd} - 0 - 0 - 0 - - /Users/yyy/Documents/protein_design/REvoDesign/REvoDesign/UI - - - true - Python.PysideBuildStep - - 1 - Build - Build - ProjectExplorer.BuildSteps.Build - - - 0 - Clean - Clean - ProjectExplorer.BuildSteps.Clean - - 2 - false - - false - - build - Python.PySideBuildConfiguration - - 1 - - - 0 - Deploy - Deploy - ProjectExplorer.BuildSteps.Deploy - - 1 - - false - ProjectExplorer.DefaultDeployConfiguration - - 1 - - true - true - true - - 2 - - false - Ui_REvoDesign.py - PythonEditor.RunConfiguration./Users/yyy/Documents/protein_design/REvoDesign/REvoDesign/UI/Ui_REvoDesign.py - /Users/yyy/Documents/protein_design/REvoDesign/REvoDesign/UI/Ui_REvoDesign.py - /Users/yyy/Documents/protein_design/REvoDesign/REvoDesign/UI/Ui_REvoDesign.py - true - true - /Users/yyy/Documents/protein_design/REvoDesign/REvoDesign/UI - - 1 - - - - ProjectExplorer.Project.Target.1 - - Desktop - Desktop (arm-darwin-generic-mach_o-64bit) - Desktop (arm-darwin-generic-mach_o-64bit) - {ce734272-b7d7-4c06-994a-976506a526c1} - 0 - 0 - 0 - - /Users/yyy/Documents/protein_design/REvoDesign/REvoDesign/UI - - - true - Python.PysideBuildStep - - 1 - Build - Build - ProjectExplorer.BuildSteps.Build - - - 0 - Clean - Clean - ProjectExplorer.BuildSteps.Clean - - 2 - false - - false - - build - Python.PySideBuildConfiguration - - 1 - - - 0 - Deploy - Deploy - ProjectExplorer.BuildSteps.Deploy - - 1 - - false - ProjectExplorer.DefaultDeployConfiguration - - 1 - - true - true - true - - 2 - - false - Ui_REvoDesign.py - PythonEditor.RunConfiguration./Users/yyy/Documents/protein_design/REvoDesign/REvoDesign/UI/Ui_REvoDesign.py - /Users/yyy/Documents/protein_design/REvoDesign/REvoDesign/UI/Ui_REvoDesign.py - /Users/yyy/Documents/protein_design/REvoDesign/REvoDesign/UI/Ui_REvoDesign.py - true - true - /Users/yyy/Documents/protein_design/REvoDesign/REvoDesign/UI - - 1 - - - - ProjectExplorer.Project.Target.2 - - Desktop - Desktop (arm-darwin-generic-mach_o-32bit) - Desktop (arm-darwin-generic-mach_o-32bit) - {1c5021eb-a3b5-4932-90c6-41451ea69647} - 0 - 0 - 0 - - /Users/yyy/Documents/protein_design/REvoDesign/REvoDesign/UI - - - true - Python.PysideBuildStep - - 1 - Build - Build - ProjectExplorer.BuildSteps.Build - - - 0 - Clean - Clean - ProjectExplorer.BuildSteps.Clean - - 2 - false - - false - - build - Python.PySideBuildConfiguration - - 1 - - - 0 - Deploy - Deploy - ProjectExplorer.BuildSteps.Deploy - - 1 - - false - ProjectExplorer.DefaultDeployConfiguration - - 1 - - true - true - true - - 2 - - false - Ui_REvoDesign.py - PythonEditor.RunConfiguration./Users/yyy/Documents/protein_design/REvoDesign/REvoDesign/UI/Ui_REvoDesign.py - /Users/yyy/Documents/protein_design/REvoDesign/REvoDesign/UI/Ui_REvoDesign.py - /Users/yyy/Documents/protein_design/REvoDesign/REvoDesign/UI/Ui_REvoDesign.py - true - true - /Users/yyy/Documents/protein_design/REvoDesign/REvoDesign/UI - - 1 - - - - ProjectExplorer.Project.Target.3 - - Desktop - Desktop (x86-darwin-generic-mach_o-32bit) - Desktop (x86-darwin-generic-mach_o-32bit) - {3ad391aa-4d57-4fd8-94ae-5f8e44f65cc1} - 0 - 0 - 0 - - /Users/yyy/Documents/protein_design/REvoDesign/REvoDesign/UI - - - true - Python.PysideBuildStep - - 1 - Build - Build - ProjectExplorer.BuildSteps.Build - - - 0 - Clean - Clean - ProjectExplorer.BuildSteps.Clean - - 2 - false - - false - - build - Python.PySideBuildConfiguration - - 1 - - - 0 - Deploy - Deploy - ProjectExplorer.BuildSteps.Deploy - - 1 - - false - ProjectExplorer.DefaultDeployConfiguration - - 1 - - true - true - true - - 2 - - false - Ui_REvoDesign.py - PythonEditor.RunConfiguration./Users/yyy/Documents/protein_design/REvoDesign/REvoDesign/UI/Ui_REvoDesign.py - /Users/yyy/Documents/protein_design/REvoDesign/REvoDesign/UI/Ui_REvoDesign.py - /Users/yyy/Documents/protein_design/REvoDesign/REvoDesign/UI/Ui_REvoDesign.py - true - true - /Users/yyy/Documents/protein_design/REvoDesign/REvoDesign/UI - - 1 - - - - ProjectExplorer.Project.TargetCount - 4 - - - ProjectExplorer.Project.Updater.FileVersion - 22 - - - Version - 22 - - diff --git a/src/REvoDesign/application/cluster_tab.py b/src/REvoDesign/application/cluster_tab.py index 301d46624..c9b52d45d 100644 --- a/src/REvoDesign/application/cluster_tab.py +++ b/src/REvoDesign/application/cluster_tab.py @@ -6,9 +6,12 @@ from __future__ import annotations +from REvoDesign.logger import ROOT_LOGGER from REvoDesign.Qt import QtCore, QtWidgets from REvoDesign.tools.customized_widgets import decide +logging = ROOT_LOGGER.getChild(__name__) + FALLBACK_CLUSTER_METHODS: tuple[str, ...] = ( "AgglomerativeCluster", "EvoCluster", @@ -77,8 +80,8 @@ def _configured_method_name(self) -> str: if self.bus is not None: try: return str(self.bus.get_value("ui.cluster.method.use", default_value="AgglomerativeCluster")) - except Exception: - pass + except Exception as exc: + logging.debug("Could not read configured cluster method; using current UI value: %s", exc) current = self.ui.comboBox_cluster_method.currentText() return current or "AgglomerativeCluster" @@ -142,19 +145,24 @@ def _install_tooltips(self) -> None: "AgglomerativeCluster uses the sequence distance matrix with average linkage and a precomputed metric." ), self.ui.label_cluster_agglomerative_representative_value: ( - "Representatives are selected as the nearest centroid among clustered variants. This is not a medoid policy." + "Representatives are selected as the nearest centroid among clustered variants. " + "This is not a medoid policy." ), self.ui.page_cluster_evo: ( - "Optional Evo inputs with positive weights are used when available; missing optional inputs are skipped and remaining weights are renormalized." + "Optional Evo inputs with positive weights are used when available; " + "missing optional inputs are skipped and remaining weights are renormalized." ), self.ui.lineEdit_cluster_evo_pssm_profile: ( - "Optional PSSM input for EvoCluster. If the corresponding weight is positive, the file path must exist." + "Optional PSSM input for EvoCluster. If the corresponding weight is positive, " + "the file path must exist." ), self.ui.lineEdit_cluster_evo_esm1v_table: ( - "Optional ESM-1v input for EvoCluster. If the corresponding weight is positive, the file path must exist." + "Optional ESM-1v input for EvoCluster. If the corresponding weight is positive, " + "the file path must exist." ), self.ui.lineEdit_cluster_evo_structure_pdb: ( - "Optional structure input for EvoCluster spatial distance. A temporary input PDB may be generated when needed." + "Optional structure input for EvoCluster spatial distance. " + "A temporary input PDB may be generated when needed." ), self.ui.lineEdit_cluster_evo_esm_mutation_col: "Column name used to read ESM mutation identifiers.", self.ui.page_cluster_kmeans: ( @@ -167,13 +175,16 @@ def _install_tooltips(self) -> None: "LegacyCluster is a deprecated compatibility mode that uses Ward linkage on the score matrix." ), self.ui.label_cluster_legacy_warning: ( - "Use this only to reproduce historical behavior. Prefer AgglomerativeCluster or EvoCluster for new analyses." + "Use this only to reproduce historical behavior. " + "Prefer AgglomerativeCluster or EvoCluster for new analyses." ), self.ui.checkBox_cluster_mutate_and_relax: ( - "Run Rosetta Mutate/Relax scoring after clustering without rewriting representative FASTA files by default." + "Run Rosetta Mutate/Relax scoring after clustering without rewriting representative " + "FASTA files by default." ), self.ui.checkBox_cluster_rosetta_override_representatives: ( - "Rosetta scoring normally exports score tables only. Representative FASTA files are rewritten only when this explicit override option is enabled." + "Rosetta scoring normally exports score tables only. Representative FASTA files " + "are rewritten only when this explicit override option is enabled." ), } diff --git a/src/REvoDesign/application/font/font_manager.py b/src/REvoDesign/application/font/font_manager.py index aca89ed95..dbda9377b 100644 --- a/src/REvoDesign/application/font/font_manager.py +++ b/src/REvoDesign/application/font/font_manager.py @@ -98,10 +98,12 @@ def set_window_font(self, custom_font: QtGui.QFont | str | None = None): """ Set the window font based on the operating system type. - This method retrieves the current client operating system type and checks if it exists in the predefined font dictionary. - If the OS type is not found, the method returns without making any changes. If the OS type is found, it iterates through - the list of fonts associated with that OS type and sets the window font to the first available font family found in the - `self.font_families` list. + This method retrieves the current client operating system type and + checks if it exists in the predefined font dictionary. If the OS type is + not found, the method returns without making any changes. If the OS type + is found, it iterates through the list of fonts associated with that OS + type and sets the window font to the first available font family found + in the `self.font_families` list. Parameters: - None @@ -155,7 +157,7 @@ def set_font(font: QtGui.QFont | str | None = None): # Iterate through and set font for all open windows if hasattr(ui, "open_windows"): - open_windows: list[QtWidgets.QWidget] = getattr(ui, "open_windows") + open_windows: list[QtWidgets.QWidget] = ui.open_windows for window in open_windows: window.setFont(font) diff --git a/src/REvoDesign/application/i18n/language_settings.py b/src/REvoDesign/application/i18n/language_settings.py index 0c40d6245..654e7ccf1 100644 --- a/src/REvoDesign/application/i18n/language_settings.py +++ b/src/REvoDesign/application/i18n/language_settings.py @@ -224,7 +224,7 @@ def _ensure_translator(self) -> QtCore.QTranslator: if app is not None: for child in app.children(): if isinstance(child, QtCore.QTranslator): - setattr(self.bus.ui, "trans", child) + self.bus.ui.trans = child self._translator_installed = True return child @@ -238,7 +238,7 @@ def _ensure_translator(self) -> QtCore.QTranslator: return existing translator = QtCore.QTranslator(self.window) - setattr(self.bus.ui, "trans", translator) + self.bus.ui.trans = translator return translator def _bind_to_action(self, language: LanguageItem): diff --git a/src/REvoDesign/application/launching.py b/src/REvoDesign/application/launching.py index 3d3045027..499ba0d60 100644 --- a/src/REvoDesign/application/launching.py +++ b/src/REvoDesign/application/launching.py @@ -23,7 +23,8 @@ QLabel#labelStatus { color: #8b95a8; font-size: 13pt; } QLabel#labelDot { color: #4ecdc4; font-size: 10pt; } QLabel#labelInfoLeft, QLabel#labelInfoRight { color: #5a6578; font-size: 9pt; } - QProgressBar#progressBar { background-color: #232a3a; border: none; border-radius: 2px; min-height: 3px; max-height: 3px; } + QProgressBar#progressBar { background-color: #232a3a; border: none; border-radius: 2px; + min-height: 3px; max-height: 3px; } QProgressBar#progressBar::chunk { background-color: #4ecdc4; border-radius: 2px; } """ @@ -34,7 +35,8 @@ QLabel#labelStatus { color: #5a6578; font-size: 13pt; } QLabel#labelDot { color: #3db8b0; font-size: 10pt; } QLabel#labelInfoLeft, QLabel#labelInfoRight { color: #8b95a8; font-size: 9pt; } - QProgressBar#progressBar { background-color: #dfdcd6; border: none; border-radius: 2px; min-height: 3px; max-height: 3px; } + QProgressBar#progressBar { background-color: #dfdcd6; border: none; border-radius: 2px; + min-height: 3px; max-height: 3px; } QProgressBar#progressBar::chunk { background-color: #3db8b0; border-radius: 2px; } """ diff --git a/src/REvoDesign/basic/abc_singleton.py b/src/REvoDesign/basic/abc_singleton.py index 4066436b4..c57407f1b 100644 --- a/src/REvoDesign/basic/abc_singleton.py +++ b/src/REvoDesign/basic/abc_singleton.py @@ -260,8 +260,10 @@ def reset_singletons(): """ Reset singleton classes. - This function is used to gracefully reset all singleton classes. It iterates through all subclasses of SingletonAbstract, - and if the subclass has an '_instance' attribute and it is not None, it calls the reset_instance method of that class. + This function is used to gracefully reset all singleton classes. It iterates + through all subclasses of SingletonAbstract, and if the subclass has an + '_instance' attribute and it is not None, it calls the reset_instance method + of that class. Note: This function does not accept parameters and does not return any value. """ diff --git a/src/REvoDesign/basic/extensions.py b/src/REvoDesign/basic/extensions.py index 03455c7c8..34a14763a 100644 --- a/src/REvoDesign/basic/extensions.py +++ b/src/REvoDesign/basic/extensions.py @@ -72,7 +72,8 @@ def list_all(self) -> list[str]: """ Returns a list of all file extensions. - This property method iterates over the `self.extensions` list, extracting the 'ext' attribute from each extension object. + This property method iterates over the `self.extensions` list, + extracting the 'ext' attribute from each extension object. """ return [e.ext for e in self.extensions] @@ -189,10 +190,12 @@ def resolve_extension(extension: str) -> FileExtensionCollection: user-friendly names with a prefix `'Customized - '`. Args: - extension (str): The extension string to be resolved. It can be a predefined name or a custom list like `"pdb;csv"`. + extension (str): The extension string to be resolved. It can be a + predefined name or a custom list like `"pdb;csv"`. Returns: - Fext.ExtColl: An object representing the file extension collection, either from predefined values or custom input. + Fext.ExtColl: An object representing the file extension collection, + either from predefined values or custom input. Example: Given input `"pdb;csv"`, this will generate: diff --git a/src/REvoDesign/basic/param_toggle.py b/src/REvoDesign/basic/param_toggle.py index 082f0993f..524cf7297 100644 --- a/src/REvoDesign/basic/param_toggle.py +++ b/src/REvoDesign/basic/param_toggle.py @@ -22,8 +22,10 @@ class ParamChangeRegistryItem: """ A class that registers a mapping between two config items. - This class is used to define how a change in one configuration item (source_cfg_item) should be mapped and registered - to another configuration item (target_cfg_item) through a UI widget signal. It uses a dictionary to map parameter values. + This class is used to define how a change in one configuration item + (source_cfg_item) should be mapped and registered to another + configuration item (target_cfg_item) through a UI widget signal. It uses + a dictionary to map parameter values. Attributes: widget_name (str): The name of the UI widget. @@ -64,10 +66,12 @@ def widget_signal(self, ui) -> QtCore.pyqtBoundSignal: def register(self, register_func: Callable[[str, str, dict[str, tuple]], None], ui: Any): """ - Registers the mapping between two configuration items by connecting the widget signal to the registration function. + Registers the mapping between two configuration items by connecting the + widget signal to the registration function. Args: - register_func (Callable[[str, str, Dict[str, Tuple]], None]): The function to be called when the signal is emitted. + register_func (Callable[[str, str, Dict[str, Tuple]], None]): The function to be + called when the signal is emitted. ui (Any): The UI instance containing the widgets. """ event = self.widget_signal(ui=ui) diff --git a/src/REvoDesign/bootstrap/set_config.py b/src/REvoDesign/bootstrap/set_config.py index 6ef50084f..a901a9784 100644 --- a/src/REvoDesign/bootstrap/set_config.py +++ b/src/REvoDesign/bootstrap/set_config.py @@ -17,13 +17,14 @@ from omegaconf import DictConfig, OmegaConf from platformdirs import user_cache_dir, user_data_dir -from REvoDesign.Qt import QtCompat, QtWidgets +CONFIG_RESET_ENV = "REVODESIGN_RESET_CONFIG" def decide(title="", description="", rich: bool = False, details: str | None = None): """ A copy of decide function from package_manager.py """ + from REvoDesign.Qt import QtCompat, QtWidgets # A confirmation message. msg = QtWidgets.QMessageBox() @@ -40,7 +41,27 @@ def decide(title="", description="", rich: bool = False, details: str | None = N return result == QtCompat.Yes -def set_REvoDesign_config_file(delete_user_config_tree: bool = False): +def _env_flag_enabled(name: str) -> bool: + return os.environ.get(name, "").strip().lower() in {"1", "true", "yes", "on"} + + +def _has_qapplication() -> bool: + try: + from REvoDesign.Qt import QtWidgets + except Exception: + return False + return QtWidgets.QApplication.instance() is not None + + +def _should_prompt_for_config_reset(prompt_user: bool | None) -> bool: + if prompt_user is not None: + return prompt_user + if os.environ.get("PYMOL_GIT_TEST"): + return False + return _has_qapplication() + + +def set_REvoDesign_config_file(delete_user_config_tree: bool = False, prompt_user: bool | None = None): """ Sets the REvoDesign configuration directory. If the main configuration file does not exist, it will be copied from the template directory. If the configuration directory exists, @@ -48,6 +69,8 @@ def set_REvoDesign_config_file(delete_user_config_tree: bool = False): Arguments: delete_user_config_tree (bool): Whether to delete the user's configuration tree. Defaults to False. + prompt_user (bool | None): Whether to ask before resetting an upgrade-time config tree. + ``None`` prompts only when a QApplication exists and test mode is not active. Returns: str: The path to the REvoDesign configuration directory. @@ -63,8 +86,9 @@ def set_REvoDesign_config_file(delete_user_config_tree: bool = False): config_dir = os.path.join(default_storage_path, "config") main_config_file = os.path.join(config_dir, "main.yaml") + reset_requested = delete_user_config_tree or _env_flag_enabled(CONFIG_RESET_ENV) - if delete_user_config_tree and os.path.exists(config_dir): + if reset_requested and os.path.exists(config_dir): print("WARNING: The configuration directory will be deleted as required") shutil.rmtree(config_dir) @@ -82,7 +106,7 @@ def set_REvoDesign_config_file(delete_user_config_tree: bool = False): ) print(reset_warning) - if decide( + if _should_prompt_for_config_reset(prompt_user) and decide( title="Reset REvoDesign Configuration?", description=reset_warning + "Do you want to continue? \n" "You can still choose to cancel the reset to proceed it manually.", @@ -90,7 +114,10 @@ def set_REvoDesign_config_file(delete_user_config_tree: bool = False): shutil.rmtree(config_dir) else: - print("Please manually delete the configuration directory and restart REvoDesign.") + print( + "Configuration reset was not approved; keeping existing files. " + f"Set {CONFIG_RESET_ENV}=1 or pass delete_user_config_tree=True to force a reset." + ) print(f"Copied configurations from {template_config_dir} to {config_dir}") shutil.copytree(src=template_config_dir, dst=config_dir, dirs_exist_ok=True) @@ -117,10 +144,12 @@ def reload_config_file( config_name: str = "main", overrides: list[str] | None = None, return_hydra_config: bool = False ) -> DictConfig: """ - Reload a configuration yaml file in a Hydra manner. As we initialize hydra w/ initialize_config_dir, - which is the REVODESIGN_CONFIG_DIR at user's data directory, the config_name is supposed as a relative path of the yaml file. + Reload a configuration yaml file in a Hydra manner. As we initialize hydra + w/ initialize_config_dir, which is the REVODESIGN_CONFIG_DIR at user's data + directory, the config_name is supposed as a relative path of the yaml file. e.g. config_name="experiments/my_experiment" refers to experiments/my_experiment.yaml - the DictConfig object can be accessed like `reload_config_file(config_name="experiments/my_experiment")["experiments"]` + the DictConfig object can be accessed like + `reload_config_file(config_name="experiments/my_experiment")["experiments"]` Arguments: config_name (str): The name of the configuration file. Defaults to "main". @@ -276,13 +305,15 @@ def list_all_config_files(config_dir: str, tree: bool = False) -> list[str]: def _iter_yaml_rel_paths(base_dir: str) -> list[str]: """ - Iterates through the specified directory and its subdirectories to collect relative paths of all .yaml files + Iterates through the specified directory and its subdirectories to collect + relative paths of all .yaml files Args: base_dir (str): The base directory path to traverse Returns: - list[str]: A list containing relative paths of all found .yaml files with respect to base_dir, sorted in lexicographical order + list[str]: A list containing relative paths of all found .yaml files + with respect to base_dir, sorted in lexicographical order """ base_dir = os.path.abspath(base_dir) rel_paths: list[str] = [] diff --git a/src/REvoDesign/citations/citation_manager.py b/src/REvoDesign/citations/citation_manager.py index c6c6f62ae..3bafa319d 100644 --- a/src/REvoDesign/citations/citation_manager.py +++ b/src/REvoDesign/citations/citation_manager.py @@ -69,7 +69,7 @@ def update(self, new_citations: dict): new_citations (dict): A dictionary containing new citations. """ if not (new_citations and isinstance(new_citations, dict)): - warnings.warn(issues.NoInputWarning(f"{new_citations=} is not a valid dictionary.")) + warnings.warn(issues.NoInputWarning(f"{new_citations=} is not a valid dictionary."), stacklevel=2) return self.called_citations.update(new_citations) @@ -91,7 +91,7 @@ def output(self, cwd: str = "."): library = bibtexparser.parse_string("\n".join(self.collected_citations)) if library.failed_blocks: - warnings.warn(issues.REvoDesignWarning(f"Could not parse {library.failed_blocks=}")) + warnings.warn(issues.REvoDesignWarning(f"Could not parse {library.failed_blocks=}"), stacklevel=2) citation_output = os.path.join( cwd, @@ -99,7 +99,8 @@ def output(self, cwd: str = "."): f'{time.strftime("%Y%m%d", time.localtime())}.bib', ) os.makedirs(os.path.dirname(citation_output), exist_ok=True) - bibtexparser.write_file(file=open(citation_output, "w", encoding="utf8"), library=library) + with open(citation_output, "w", encoding="utf8") as citation_file: + bibtexparser.write_file(file=citation_file, library=library) logging.info(f"Citation is created at {citation_output}") def dismiss(self, modulename: str): @@ -130,7 +131,9 @@ def notice(cls): """ Display the citation notice. - This method checks if there are citations to be displayed, and if the current module's citations have not been silenced, it logs the citation information. + This method checks if there are citations to be displayed, and if the + current module's citations have not been silenced, it logs the citation + information. """ # If the __bibtex__ dictionary is empty, log a debug message and return. if not cls.__bibtex__: @@ -171,7 +174,8 @@ def cite(cls): """ Add citation to the citation manager. - This method adds the current module's citation information to the citation manager and then displays the citation notice. + This method adds the current module's citation information to the + citation manager and then displays the citation notice. """ citations = cls.__bibtex__ # Check if the citation information is a dictionary. @@ -200,7 +204,7 @@ def get_citable_class(cls, func: Callable) -> type["CitableModuleAbstract"]: class AnonymousCitableModule(cls): __name__ = f"{func.__name__}_CitableModule" - __bibtex__ = getattr(func, "__bibtex__") + __bibtex__ = func.__bibtex__ AnonymousCitableModule.__doc__ = f"""Citable module for function { func.__name__}, auto generated by CitableModuleAbstract.""" # type: ignore[assignment] diff --git a/src/REvoDesign/clients/QtSocketConnector.py b/src/REvoDesign/clients/QtSocketConnector.py index da07b03e4..86ba66211 100644 --- a/src/REvoDesign/clients/QtSocketConnector.py +++ b/src/REvoDesign/clients/QtSocketConnector.py @@ -15,10 +15,11 @@ import base64 import copy -import pickle # nosec B403: inter-process Qt socket serialization, trusted local context + +# Inter-process Qt socket serialization uses a trusted local context. +import pickle # nosec B403 import socket import time -import traceback import warnings from collections.abc import Iterable, Mapping from dataclasses import dataclass, field @@ -81,7 +82,7 @@ def does_user_exist(self, uuid: str) -> bool: def add_user(self, uuid: str, user: Client): if uuid in self.clients: - warnings.warn(issues.SocketUserAlreadyExists("user already joined.")) + warnings.warn(issues.SocketUserAlreadyExists("user already joined."), stacklevel=2) return self.clients[uuid] = user @@ -202,13 +203,15 @@ def unpack(package: bytes) -> dict[str, str]: def digest_dict(self, unpacked_msg_dict: dict[str, str]) -> tuple[supported_datatypes | None, Any]: datatype: self.supported_datatypes = unpacked_msg_dict.get("datatype") if not datatype: - warnings.warn(issues.BadDataWarning(f"bad data {datatype=} has been discarded")) + warnings.warn(issues.BadDataWarning(f"bad data {datatype=} has been discarded"), stacklevel=2) return None, None obj_encoded = unpacked_msg_dict.get("obj") obj = self.decode(obj_encoded) expected_obj_type = self.supported_datatypes_mapping.get(datatype) if not isinstance(obj, expected_obj_type): - warnings.warn(issues.BadDataWarning(f"Bad data type mismatch: {datatype=}: {expected_obj_type=}")) + warnings.warn( + issues.BadDataWarning(f"Bad data type mismatch: {datatype=}: {expected_obj_type=}"), stacklevel=2 + ) return None, None return datatype, obj @@ -227,7 +230,7 @@ def broadcast( all_clients = [meetingroom.client] elif isinstance(meetingroom, MeetingRoom): if not meetingroom or meetingroom.empty: - warnings.warn(issues.SocketMeetingRoomIsEmpty("Empty meeting room.")) + warnings.warn(issues.SocketMeetingRoomIsEmpty("Empty meeting room."), stacklevel=2) return all_clients = meetingroom.current_clients elif isinstance(meetingroom, list): @@ -276,7 +279,7 @@ def received(self, packed_msg: bytes) -> tuple[None, None] | tuple[str, Any] | t unpacked_msg = self.unpack(packed_msg) (datatype, obj) = self.digest_dict(unpacked_msg_dict=unpacked_msg) if not datatype: - warnings.warn(issues.BadDataWarning("Bad data received with unexpected none datatype.")) + warnings.warn(issues.BadDataWarning("Bad data received with unexpected none datatype."), stacklevel=2) return None, None if datatype != "MessageStack": return datatype, obj @@ -476,11 +479,15 @@ def processBinaryMessage(self, client, message: QtCore.QByteArray): try: msg_type, msg_content = self.bc_worker.received(packed_msg=message) except issues.FobbidenDataTypeError: - warnings.warn(issues.BadDataWarning(f"Bad data from {username}@{node}({user_id}) is discarded.")) + warnings.warn( + issues.BadDataWarning(f"Bad data from {username}@{node}({user_id}) is discarded."), stacklevel=2 + ) return if not self.bc_worker.typing_is_valid(msg_type=msg_type, msg_content=msg_content): - warnings.warn(issues.BadDataWarning(f"Bad data from socket discarded: {msg_type=},{msg_content=}")) + warnings.warn( + issues.BadDataWarning(f"Bad data from socket discarded: {msg_type=},{msg_content=}"), stacklevel=2 + ) return # printout readable msgs @@ -631,7 +638,7 @@ async def broadcast_object(self, obj, data_type): None """ if not self.is_running: - warnings.warn(issues.SocketWarning("Host is not running")) + warnings.warn(issues.SocketWarning("Host is not running"), stacklevel=2) return self.bc_worker.broadcast(meetingroom=self.meetingroom, obj=obj, obj_type=data_type) @@ -802,8 +809,7 @@ def connect_to_server(self): logging.info("Connection established.") except Exception: - logging.error("Unexpected error during connection:") - traceback.print_exc() + logging.exception("Unexpected error during connection.") @property def connected(self): @@ -821,8 +827,7 @@ def close_connection(self): refresh_tree_widget(user_tree={}, treeWidget_ws_peers=self.treeWidget_ws_peers) except Exception: - traceback.print_exc() - logging.error("Client disconnection failed.") + logging.exception("Client disconnection failed.") @property def server_is_reachable(self): @@ -915,7 +920,8 @@ def messageDispatcher(self, msg_type: str, msg_content: Any): warnings.warn( issues.FobbidenDataTypeError( f"The content of a MessageStack is expected to be an Iterable instead of a {type(msg_content)}" - ) + ), + stacklevel=2, ) return @@ -925,12 +931,14 @@ def messageDispatcher(self, msg_type: str, msg_content: Any): issues.FobbidenDataTypeError( f"The content of a MessageStack item {msg_data_item=} is expected to be " f"a tuple with 2 item instead of a {type(msg_content)}" - ) + ), + stacklevel=2, ) continue if (len_msg_content := len(msg_data_item)) != 2: warnings.warn( - issues.FobbidenDataTypeError(f"The MessageStack item has invalid length {len_msg_content}") + issues.FobbidenDataTypeError(f"The MessageStack item has invalid length {len_msg_content}"), + stacklevel=2, ) continue @@ -950,11 +958,13 @@ def handleBinaryMessage(self, message: QtCore.QByteArray): try: msg_type, msg_content = self.bc_worker.received(packed_msg=msg_bytes) except issues.FobbidenDataTypeError: - warnings.warn(issues.BadDataWarning("Bad data from Host is discarded.")) + warnings.warn(issues.BadDataWarning("Bad data from Host is discarded."), stacklevel=2) return if not self.bc_worker.typing_is_valid(msg_type=msg_type, msg_content=msg_content): - warnings.warn(issues.BadDataWarning(f"Bad data from socket discarded: {msg_type=},{msg_content=}")) + warnings.warn( + issues.BadDataWarning(f"Bad data from socket discarded: {msg_type=},{msg_content=}"), stacklevel=2 + ) return if msg_type in self.bc_worker.readble_type: diff --git a/src/REvoDesign/clusters/cluster_runner.py b/src/REvoDesign/clusters/cluster_runner.py index 50cac28a3..38fe4240e 100644 --- a/src/REvoDesign/clusters/cluster_runner.py +++ b/src/REvoDesign/clusters/cluster_runner.py @@ -76,8 +76,8 @@ def _representative_policy(clustering) -> str: if hasattr(clustering, "get_method_spec"): try: return clustering.get_method_spec().representative_policy - except Exception: - pass + except Exception as exc: + logging.debug("Could not read representative policy from clustering backend: %s", exc) return "Representative selection policy is not provided by this clustering backend." @staticmethod diff --git a/src/REvoDesign/clusters/cluster_sequence.py b/src/REvoDesign/clusters/cluster_sequence.py index bfd242f95..f584fba46 100644 --- a/src/REvoDesign/clusters/cluster_sequence.py +++ b/src/REvoDesign/clusters/cluster_sequence.py @@ -243,10 +243,11 @@ def override_cluster_centers_with_rosetta(self, rosetta_results): def _calculate_pairwise_scores(self, progressbar) -> np.ndarray: from REvoDesign.tools.customized_widgets import QtParallelExecutor - handle = open(self.fastafile) - self.records = list(SeqIO.parse(handle, "fasta")) + with open(self.fastafile) as handle: + self.records = list(SeqIO.parse(handle, "fasta")) if self.shuffle_variant: - rng = random.Random(self.random_seed) + # Deterministic clustering order is controlled by the UI random-seed setting. + rng = random.Random(self.random_seed) # nosec B311 self.records = rng.sample(self.records, len(self.records)) nm_seqs = len(self.records) @@ -405,7 +406,9 @@ def build_method_report(self) -> dict[str, object]: "deprecated": spec.deprecated, "representative_selection_policy": spec.representative_policy, "compatibility_outputs": { - "cluster_centers_stochastic.fasta": "Deprecated compatibility alias of nearest-centroid representatives." + "cluster_centers_stochastic.fasta": ( + "Deprecated compatibility alias of nearest-centroid representatives." + ) }, } diff --git a/src/REvoDesign/clusters/methods/evo.py b/src/REvoDesign/clusters/methods/evo.py index 56192e0b0..51797f646 100644 --- a/src/REvoDesign/clusters/methods/evo.py +++ b/src/REvoDesign/clusters/methods/evo.py @@ -186,7 +186,8 @@ def _build_pssm_distance_matrix(self, mutations_per_record: list[list]) -> np.nd continue try: values.append(float(df.loc[mutation.mut_res, col])) - except Exception: + except Exception as exc: + logging.debug("Skipping invalid PSSM score for %s at %s: %s", mutation.mut_res, col, exc) continue if values: scalar_scores[idx] = float(np.mean(values)) diff --git a/src/REvoDesign/common/multi_mutant_designer.py b/src/REvoDesign/common/multi_mutant_designer.py index 0f376d0b0..f79a17149 100644 --- a/src/REvoDesign/common/multi_mutant_designer.py +++ b/src/REvoDesign/common/multi_mutant_designer.py @@ -10,7 +10,7 @@ import copy import os -import random +import secrets import warnings from itertools import pairwise @@ -99,7 +99,9 @@ def get_input_and_initialize(self): self.design_case = cmd.get_unused_name("multi_design") logging.info( - f"Mutant Tree for multi-design is initialized. {len(self.design_pool_tree.all_mutant_branch_ids)} groups with {len(self.design_pool_tree.all_mutant_ids)} mutants." + "Mutant Tree for multi-design is initialized. " + f"{len(self.design_pool_tree.all_mutant_branch_ids)} groups with " + f"{len(self.design_pool_tree.all_mutant_ids)} mutants." ) @staticmethod @@ -133,7 +135,7 @@ def refresh_design_color(self): mut_obj.mutant_score for mut_obj in self.all_design_multi_design_mutant_object if mut_obj.mutant_score ] - for (i_obj, obj), (j_des, des) in zip( + for (i_obj, obj), (_j_des, des) in zip( enumerate(self.design_case_variant_objects), enumerate(self.all_design_multi_design_mutant_object), ): @@ -182,7 +184,9 @@ def evaluate_design(self) -> Mutant: tmp_mutant_obj.wt_protein_sequence = self.designable_sequences if not self.magician.gimmick: - warnings.warn(issues.ConflictWarning("Abord design evaluation because no external scorer is defined.")) + warnings.warn( + issues.ConflictWarning("Abord design evaluation because no external scorer is defined."), stacklevel=2 + ) else: tmp_mutant_obj.mutant_score = self.magician.gimmick.scorer(mutant=tmp_mutant_obj) @@ -226,7 +230,7 @@ def _auto_pick_tryout(self, tryout=30): This method tries to automatically pick mutants for the design process. """ - for i in range(tryout): + for _i in range(tryout): try: branch, (mutant_id, mutant_obj) = self._select_random_mutant() except IndexError: @@ -268,7 +272,8 @@ def pick_next_mutant(self): warnings.warn( issues.NoInputWarning( "Temperal mutant tree for multi-design is empty after designing! This design is ended." - ) + ), + stacklevel=2, ) self.start_new_design() return @@ -278,7 +283,7 @@ def pick_next_mutant(self): self._auto_pick_tryout() if num_mut_before_picking == self.in_design_multi_design_case.mutant_num: # failed picking - warnings.warn(issues.NoResultsWarning("Failed auto picking. Please take anther try.")) + warnings.warn(issues.NoResultsWarning("Failed auto picking. Please take anther try."), stacklevel=2) return # last mutant @@ -461,8 +466,8 @@ def _select_random_mutant(self) -> tuple[str, tuple[str, Mutant]]: This method randomly selects a mutant for the design process. """ - branch = random.choice(self.design_pool_tree_copy.all_mutant_branch_ids) - mut = random.choice(list(self.design_pool_tree_copy.get_a_branch(branch).items())) + branch = secrets.choice(self.design_pool_tree_copy.all_mutant_branch_ids) + mut = secrets.choice(list(self.design_pool_tree_copy.get_a_branch(branch).items())) return branch, mut def _is_compatible_mutant( diff --git a/src/REvoDesign/common/mutant_tree.py b/src/REvoDesign/common/mutant_tree.py index eb956813f..ce000ed4d 100644 --- a/src/REvoDesign/common/mutant_tree.py +++ b/src/REvoDesign/common/mutant_tree.py @@ -24,11 +24,16 @@ class MutantDict(TypedDict): class MutateRunner(Protocol): - def run_mutate_parallel(self, mutants: list[Mutant], nproc: int = 2) -> list[str]: ... + # fmt: off + def run_mutate_parallel(self, mutants: list[Mutant], nproc: int = 2) -> list[str]: + ... - def mutated_pdb_mapping(self, mutant_tree: "MutantTree", pdb_fps: list[str]) -> "MutantTree": ... + def mutated_pdb_mapping(self, mutant_tree: "MutantTree", pdb_fps: list[str]) -> "MutantTree": + ... - def cite(self) -> None: ... + def cite(self) -> None: + ... + # fmt: on class MutantTree: @@ -153,7 +158,8 @@ def has(self, obj: str | Mutant) -> bool: """ Check if the given object exists in the instance. - This function accepts an object of type string or Mutant and determines whether the object is present in the instance's mutant collection. + This function accepts an object of type string or Mutant and determines + whether the object is present in the instance's mutant collection. Parameters: - obj (Union[str, Mutant]): The object to check. It can be a string representing a mutant ID or a Mutant object. diff --git a/src/REvoDesign/common/mutant_visualise.py b/src/REvoDesign/common/mutant_visualise.py index 99e834183..09be02c1d 100644 --- a/src/REvoDesign/common/mutant_visualise.py +++ b/src/REvoDesign/common/mutant_visualise.py @@ -218,14 +218,15 @@ def _get_mutation_data(self): return pd.read_fwf(self.mutfile) if self.mutfile.lower().endswith((".fasta", ".fas", ".fa")): # Read mutant data from fasta file. - _mutation_objs = [ - extract_mutant_from_sequences( - mutant_sequence=str(mut_record.seq), - wt_sequences=self.designable_sequences, - chain_id=self.chain_id, - ) - for mut_record in SeqIO.parse(open(self.mutfile), format="fasta") - ] + with open(self.mutfile) as mutant_file: + _mutation_objs = [ + extract_mutant_from_sequences( + mutant_sequence=str(mut_record.seq), + wt_sequences=self.designable_sequences, + chain_id=self.chain_id, + ) + for mut_record in SeqIO.parse(mutant_file, format="fasta") + ] return pd.DataFrame.from_dict( {self.key_col: [mut_obj.short_mutant_id for mut_obj in _mutation_objs if mut_obj is not None]} @@ -418,6 +419,7 @@ def merge_sessions_via_commandline(self): f"Temperal merged result is failed to create. Try again with a clean PyMOL session. \n" f"STDOUT:\n{merge_results.stdout}\n" f"STDERR:\n{merge_results.stderr}" - ) + ), + stacklevel=2, ) return diff --git a/src/REvoDesign/common/profile_parsers.py b/src/REvoDesign/common/profile_parsers.py index 6b43fb9b6..768c984fa 100644 --- a/src/REvoDesign/common/profile_parsers.py +++ b/src/REvoDesign/common/profile_parsers.py @@ -153,11 +153,12 @@ def convert_PSSM_file_to_df(input_pssm_file): PSSM_Alphabet = "ARNDCQEGHILKMFPSTWYV" # Fetch table header of PSSM c = 0 - for line in open(input_pssm_file): - pssm_header = line - c += 1 - if c == 3: - break + with open(input_pssm_file) as pssm_file: + for line in pssm_file: + pssm_header = line + c += 1 + if c == 3: + break logging.info(pssm_header) diff --git a/src/REvoDesign/config/environ.yaml b/src/REvoDesign/config/environ.yaml index a3b781f18..20875e50f 100644 --- a/src/REvoDesign/config/environ.yaml +++ b/src/REvoDesign/config/environ.yaml @@ -1,5 +1,6 @@ # register all variables here +override_existing: false variables: - # ROSETTA_BIN: /Users/yyy/rosetta/rosetta_binary/rosetta.source.release-371/main/source/bin/ + # ROSETTA_BIN: /path/to/rosetta/source/bin/ # OPENKINETICS_API_KEY: your-openkinetics-api-key diff --git a/src/REvoDesign/driver/environ_register.py b/src/REvoDesign/driver/environ_register.py index bbc07113c..d93ed209c 100644 --- a/src/REvoDesign/driver/environ_register.py +++ b/src/REvoDesign/driver/environ_register.py @@ -26,17 +26,21 @@ def register_environment_variables(): bus.cfg_group["environ"].reload() # retrieve the environment variables from config object - EnvironBindItemCollection: Mapping[str, Any] | None = bus.get_value("variables", dict, cfg="environ") - if not EnvironBindItemCollection: + environ_bind_items: Mapping[str, Any] | None = bus.get_value("variables", dict, cfg="environ") + if not environ_bind_items: logging.debug("No environment variables to register") return + override_existing = bool(bus.get_value("override_existing", bool, cfg="environ", default_value=False)) # register environment variables logging.debug("Registering environment variables") - if isinstance(EnvironBindItemCollection, Mapping): - for k, v in EnvironBindItemCollection.items(): + if isinstance(environ_bind_items, Mapping): + for k, v in environ_bind_items.items(): if v is None: continue + if k in os.environ and not override_existing: + logging.debug("Skipping %s because it already exists in the process environment", k) + continue logging.debug(f"Adding {k}: {v}") os.environ[k] = v logging.debug("Environment variables registered") diff --git a/src/REvoDesign/driver/file_dialog.py b/src/REvoDesign/driver/file_dialog.py index af8e51dc3..21773fd3d 100644 --- a/src/REvoDesign/driver/file_dialog.py +++ b/src/REvoDesign/driver/file_dialog.py @@ -8,6 +8,10 @@ """ import os +import shutil +import tarfile +import tempfile +import zipfile from functools import partial from typing import Any, Literal @@ -24,6 +28,8 @@ logging = ROOT_LOGGER.getChild(__name__) +MAX_ARCHIVE_BROWSE_BYTES = 2 * 1024 * 1024 * 1024 + class FileDialog(SingletonAbstract): """ @@ -45,6 +51,7 @@ def singleton_init(self, window: Any | None, pwd: str | None): """ self.window = window self.PWD = pwd or os.getcwd() + self._archive_temp_dirs: set[str] = set() self.register_file_dialog_buttons() # Mark the instance as initialized to prevent reinitialization self.initialize() @@ -58,7 +65,10 @@ def browse_multiple_files( return getMultipleFiles(self.window, exts) def browse_filename( - self, mode: IO_MODE = "r", exts: tuple[FileExtensionCollection, ...] = (file_extensions.Any,) + self, + mode: IO_MODE = "r", + exts: tuple[FileExtensionCollection, ...] = (file_extensions.Any,), + directory: str | None = None, ) -> str | None: """Open Finder/Explorer to browse from a filename @@ -83,7 +93,12 @@ def browse_filename( # otherwise, open a file to read browse_title = "Open ..." - filename = getOpenFileNameWithExt(self.window, browse_title, filter=filter_strings) + filename = getOpenFileNameWithExt( + self.window, + browse_title, + directory or "", + filter=filter_strings, + ) # no file selected if not filename: @@ -110,8 +125,9 @@ def browse_filename( return filename # otherwise, extract the archive and browse the extracted file - flatten_compressed_files(filename, self.PWD) - return self.browse_filename(mode, exts=exts) + extracted_dir = flatten_compressed_files(filename) + self._archive_temp_dirs.add(extracted_dir) + return self.browse_filename(mode, exts=exts, directory=extracted_dir) # A universal and versatile function for input file path browsing. @@ -275,36 +291,70 @@ def register_file_dialog_buttons(self): partial(self.open_mutant_table, "ui.evaluate.input.to_mutant_txt", "w") ) + def cleanup_archive_dirs(self) -> None: + temp_dirs = getattr(self, "_archive_temp_dirs", set()) + for temp_dir in list(temp_dirs): + shutil.rmtree(temp_dir, ignore_errors=True) + temp_dirs.clear() + + def __del__(self): + self.cleanup_archive_dirs() -def flatten_compressed_files(compressed_file: str, target_dir: str | None = None) -> str: + +def flatten_compressed_files( + compressed_file: str, + target_dir: str | None = None, + max_archive_bytes: int = MAX_ARCHIVE_BROWSE_BYTES, +) -> str: """ Flattens and extracts the contents of a compressed file. Parameters: - compressed_file (str): The path to the compressed file to be extracted. - - target_dir (Optional[str]): The directory where the extracted files will be placed. - If not provided, the current working directory is used. + - target_dir (Optional[str]): Parent directory for the temporary extraction directory. + If not provided, the system temporary directory is used. Returns: - str: The path to the directory where the files have been extracted. """ - # Set the target directory to the current working directory if not specified - if target_dir is None: - target_dir = os.getcwd() + archive_size = os.path.getsize(compressed_file) + if archive_size > max_archive_bytes: + raise ValueError( + f"Archive is too large to browse safely: {archive_size} bytes " f"(limit: {max_archive_bytes} bytes)" + ) + payload_size = _compressed_archive_payload_size(compressed_file) + if payload_size > max_archive_bytes: + raise ValueError( + f"Archive expands too large to browse safely: {payload_size} bytes " f"(limit: {max_archive_bytes} bytes)" + ) # Create a path for the extracted files - flatten_path = os.path.join( - target_dir, - "expanded_compressed_files", - os.path.basename(compressed_file), + temp_parent = target_dir or tempfile.gettempdir() + flatten_path = tempfile.mkdtemp( + prefix=f"revodesign-{os.path.basename(compressed_file)}-", + dir=temp_parent, ) - # Create the directory if it does not exist - os.makedirs(flatten_path, exist_ok=True) - # Extract the archive to the specified directory - extract_archive(archive_file=compressed_file, extract_to=flatten_path) + try: + extract_archive(archive_file=compressed_file, extract_to=flatten_path) + except Exception: + shutil.rmtree(flatten_path, ignore_errors=True) + raise # Return the path to the extracted directory return flatten_path + + +def _compressed_archive_payload_size(compressed_file: str) -> int: + if compressed_file.endswith(".zip"): + with zipfile.ZipFile(compressed_file, "r") as zip_ref: + return sum(info.file_size for info in zip_ref.infolist() if not info.is_dir()) + + tar_suffixes = (".tar", ".tar.gz", ".tgz", ".tar.bz2", ".tbz", ".tar.xz") + if compressed_file.endswith(tar_suffixes): + with tarfile.open(compressed_file, "r:*") as tar_ref: + return sum(member.size for member in tar_ref.getmembers() if member.isfile()) + + return os.path.getsize(compressed_file) diff --git a/src/REvoDesign/driver/param_toggle_register.py b/src/REvoDesign/driver/param_toggle_register.py index e83cb9cb4..7986ac966 100644 --- a/src/REvoDesign/driver/param_toggle_register.py +++ b/src/REvoDesign/driver/param_toggle_register.py @@ -37,7 +37,8 @@ class ParserOrDesigner(Protocol): """ Protocol class to define the structure for a parser or designer type. - This class inherits from Protocol and specifies the attributes that any instance conforming to this protocol must have. + This class inherits from Protocol and specifies the attributes that any + instance conforming to this protocol must have. Attributes: name (str): The name of the parser or designer. diff --git a/src/REvoDesign/driver/ui_driver.py b/src/REvoDesign/driver/ui_driver.py index 60ac1bb71..110252bca 100644 --- a/src/REvoDesign/driver/ui_driver.py +++ b/src/REvoDesign/driver/ui_driver.py @@ -63,7 +63,7 @@ def reset_instance(cls): for k, s in attr_dict.items(): if hasattr(s, "controller"): - controller = getattr(s, "controller") + controller = s.controller try: if issubclass(s.controller.__class__, SingletonAbstract): print(f"Resetting {k}: {controller.__class__.__name__}", end=" ") @@ -97,7 +97,8 @@ class HeadlessProtocol(Protocol): @dataclass class Config: """ - A dataclass to represent a configuration file. It contains the name, path, and configuration data of a configuration file. + A dataclass to represent a configuration file. It contains the name, path, + and configuration data of a configuration file. Attributes: name: str -- The name of the configuration file. @@ -344,22 +345,29 @@ class ConfigBus(SingletonAbstract, CitableModuleAbstract): Methods: Non-headless Methods: - initialize_widget_with_cfg_group(): Initializes UI widgets with their corresponding configuration settings. - update_cfg_item_from_widget(widget_id: str): Updates a configuration setting based on the value of a UI widget. - register_widget_changes_to_cfg(): Registers UI widget changes to update the configuration settings. + initialize_widget_with_cfg_group(): Initializes UI widgets with their corresponding + configuration settings. + update_cfg_item_from_widget(widget_id: str): Updates a configuration setting based on + the value of a UI widget. + register_widget_changes_to_cfg(): Registers UI widget changes to update the + configuration settings. get_widget_from_id(widget_id: str): Retrieves a UI widget based on its ID. - get_widget_from_cfg_item(cfg_item: str): Retrieves a UI widget based on its corresponding configuration item. + get_widget_from_cfg_item(cfg_item: str): Retrieves a UI widget based on its + corresponding configuration item. get_widget_value(cfg_item: str): Retrieves the value of a UI widget based on its corresponding configuration item. set_widget_value(cfg_item: str, value): Sets the value of a UI widget based on its corresponding configuration item. - restore_widget_value(cfg_item: str): Restores the value of a UI widget to its default configuration setting. + restore_widget_value(cfg_item: str): Restores the value of a UI widget to its default + configuration setting. get_cfg_item(widget_id: str): Retrieves the configuration item corresponding to a UI widget ID. button(id: str): Retrieves a button widget based on its ID. - toggle_buttons(buttons: Iterable, set_enabled: bool = False): Toggles the enabled state of a list of buttons. + toggle_buttons(buttons: Iterable, set_enabled: bool = False): Toggles the enabled + state of a list of buttons. Headless Only Methods: - get_value(cfg_item: str, typing=None): Retrieves the value of a configuration item, with optional type casting. + get_value(cfg_item: str, typing=None): Retrieves the value of a configuration item, + with optional type casting. set_value(cfg_item: str, value): Sets the value of a configuration item. fp_lock(cfg_fps: Union[list, tuple, str], buttons_id_to_release: Union[list, tuple, str]): Locks or unlocks @@ -391,7 +399,7 @@ def singleton_init(self, ui=None): def initialize_widget_with_group(self): # Initializes UI widgets with their corresponding configuration settings. - for i, gr in enumerate(GroupRegistryCollection): + for gr in GroupRegistryCollection: group_values = [] widget = self.get_widget_from_id(widget_id=gr.cfg_item) if isinstance(widget, str): @@ -503,6 +511,7 @@ def get_cfg_item(self, widget_id: str) -> str: raise ValueError(f"{widget_id} is not a valid widget ID.") return cfg_item + # fmt: off @overload def get_value( self, @@ -511,7 +520,8 @@ def get_value( reject_none: bool, default_value: None = ..., cfg: Config | str = "main", - ) -> ValueFromConfigT: ... + ) -> ValueFromConfigT: + ... @overload def get_value( @@ -521,7 +531,8 @@ def get_value( reject_none: bool, default_value: bool = ..., cfg: Config | str = "main", - ) -> bool: ... + ) -> bool: + ... @overload def get_value( @@ -531,10 +542,13 @@ def get_value( reject_none: bool = True, default_value: ValueFromConfigT | None = ..., cfg: Config | str = "main", - ) -> ValueFromConfigT: ... + ) -> ValueFromConfigT: + ... @overload - def get_value(self, cfg_item: str, converter: None) -> Any: ... + def get_value(self, cfg_item: str, converter: None) -> Any: + ... + # fmt: on def get_value( self, diff --git a/src/REvoDesign/editor/monaco/monaco.py b/src/REvoDesign/editor/monaco/monaco.py index 9b433190d..faa35c2c8 100644 --- a/src/REvoDesign/editor/monaco/monaco.py +++ b/src/REvoDesign/editor/monaco/monaco.py @@ -3,10 +3,12 @@ # SPDX-License-Identifier: GPL-3.0-only +import base64 import os import shutil from urllib.parse import urlencode +import requests from platformdirs import user_data_dir from REvoDesign import issues @@ -21,6 +23,29 @@ logging = ROOT_LOGGER.getChild(__name__) +def _npm_integrity_to_pooch_hash(integrity: str) -> str: + hash_name, hash_value = integrity.split("-", 1) + return f"{hash_name}:{base64.b64decode(hash_value).hex()}" + + +def _get_monaco_tarball_hash(version: str) -> str: + metadata_url = f"https://registry.npmjs.org/monaco-editor/{version}" + response = requests.get(metadata_url, timeout=30) + response.raise_for_status() + package_metadata = response.json() + + tarball_dist = package_metadata.get("dist", {}) + integrity = tarball_dist.get("integrity") + if integrity: + return _npm_integrity_to_pooch_hash(integrity) + + shasum = tarball_dist.get("shasum") + if shasum: + return f"sha1:{shasum}" + + raise RuntimeError(f"Could not find npm integrity metadata for monaco-editor {version}") + + class MonacoEditorManager: def __init__(self, app_name="YourAppName", app_author="YourCompany", version="latest"): """ @@ -100,7 +125,7 @@ def download_monaco_editor(self, version="latest"): down_registry = FileDownloadRegistry( name="monaco-editor", base_url=cdn_base_url, - registry={f"monaco-editor-{version}.tgz": None}, + registry={f"monaco-editor-{version}.tgz": _get_monaco_tarball_hash(version)}, version=version, customized_directory=self.editor_path, # https://github.com/amio/npm-mirrors/blob/master/index.js diff --git a/src/REvoDesign/editor/monaco/server.py b/src/REvoDesign/editor/monaco/server.py index 4ad65c62e..820c1b67b 100644 --- a/src/REvoDesign/editor/monaco/server.py +++ b/src/REvoDesign/editor/monaco/server.py @@ -12,6 +12,7 @@ from contextlib import asynccontextmanager from html import escape from pathlib import Path +from typing import Annotated import uvicorn from fastapi import FastAPI, HTTPException, Query, Request @@ -107,7 +108,7 @@ def initialize_token() -> str: SECRET_TOKEN = secrets.token_urlsafe(32) config_store.set("editor.token", SECRET_TOKEN) - logging.info(f"Generated SECRET_TOKEN: {SECRET_TOKEN}") + logging.info("Generated editor authentication token.") return SECRET_TOKEN @@ -149,7 +150,7 @@ async def lifespan(app: FastAPI): config_store = ConfigStore() html_dir = config_store.get("editor.backend.html_dir") if not html_dir: - warnings.warn(issues.MissingExternalTool("Monaco Editor is not ready.")) + warnings.warn(issues.MissingExternalTool("Monaco Editor is not ready."), stacklevel=2) # ensure monaco editor is ready ensure_monaco() # re-fetch the updated value from config store @@ -225,7 +226,7 @@ async def serve_editor(file_path: str, token: str = None, request: Request = Non @app.get("/load_file", response_class=JSONResponse) -async def load_file(file_path: str, token: str = Query(None), request: Request = None): +async def load_file(file_path: str, token: Annotated[str | None, Query()] = None, request: Request = None): verify_token(token, request) # Validate file existence @@ -253,7 +254,7 @@ class SaveFileRequest(BaseModel): @app.post("/save_file", response_class=JSONResponse) -async def save_file(data: SaveFileRequest, token: str = Query(None), request: Request = None): +async def save_file(data: SaveFileRequest, token: Annotated[str | None, Query()] = None, request: Request = None): verify_token(token, request) # Validate the file path diff --git a/src/REvoDesign/evaluate/evaluator.py b/src/REvoDesign/evaluate/evaluator.py index c6f2aec3a..4ca9030af 100644 --- a/src/REvoDesign/evaluate/evaluator.py +++ b/src/REvoDesign/evaluate/evaluator.py @@ -99,7 +99,8 @@ def mutant_decision(self, decision_to_accept: bool): return logging.debug( - f'{"Accepting" if decision_to_accept else "Rejecting"} mutant {self.mutant_tree_candidates.current_mutant_id}' + f'{"Accepting" if decision_to_accept else "Rejecting"} mutant ' + f"{self.mutant_tree_candidates.current_mutant_id}" ) if decision_to_accept: @@ -271,8 +272,10 @@ def find_all_best_mutants(self): # Ask whether to overide confirmed = decide( title="Override existed mutant table choices?", - description="You currently have existed mutant table choices, which shall be overriden by using `I'm lucky`. \n \ - Are you really sure? ", + description=( + "You currently have existed mutant table choices, which shall be overriden " + "by using `I'm lucky`.\nAre you really sure? " + ), ) if not confirmed: diff --git a/src/REvoDesign/logger/logger.py b/src/REvoDesign/logger/logger.py index e526c188c..f7a160623 100644 --- a/src/REvoDesign/logger/logger.py +++ b/src/REvoDesign/logger/logger.py @@ -330,11 +330,11 @@ def get_current_logger_level(channel: str = "root"): key = f"handlers.{channel}.level" try: return ConfigBus().get_value(key, str, reject_none=True, cfg="logger") - except Exception as e: + except Exception: logging.error(f"Logger channel {channel} does not exist like key: {key}") logging.error(f"Available channels: {list_all_logger_channels()}") logging.error(f'All config: {ConfigBus().cfg_group["logger"]}') - raise e + raise def list_all_logger_levels(): diff --git a/src/REvoDesign/magician/designers/cart_ddg.py b/src/REvoDesign/magician/designers/cart_ddg.py index 5eb695d8c..a3709ce67 100644 --- a/src/REvoDesign/magician/designers/cart_ddg.py +++ b/src/REvoDesign/magician/designers/cart_ddg.py @@ -149,8 +149,11 @@ def scorer(self, mutant: Mutant | RosettaPyProteinSequence, **kwargs) -> float: __bibtex__ = rosetta_utils.copy_rosetta_citation( { "Cartesian-ddG": """@article{doi:10.1021/acs.jctc.6b00819, -author = {Park, Hahnbeom and Bradley, Philip and Greisen, Per Jr. and Liu, Yuan and Mulligan, Vikram Khipple and Kim, David E. and Baker, David and DiMaio, Frank}, -title = {Simultaneous Optimization of Biomolecular Energy Functions on Features from Small Molecules and Macromolecules}, +author = {Park, Hahnbeom and Bradley, Philip and Greisen, Per Jr. +and Liu, Yuan and Mulligan, Vikram Khipple and Kim, David E. and Baker, David +and DiMaio, Frank}, +title = {Simultaneous Optimization of Biomolecular Energy Functions on Features +from Small Molecules and Macromolecules}, journal = {Journal of Chemical Theory and Computation}, volume = {12}, number = {12}, diff --git a/src/REvoDesign/magician/designers/colabdesign.py b/src/REvoDesign/magician/designers/colabdesign.py index 89cae03e6..319e66d20 100644 --- a/src/REvoDesign/magician/designers/colabdesign.py +++ b/src/REvoDesign/magician/designers/colabdesign.py @@ -59,11 +59,8 @@ def initialize(self, *args, **kwargs): self.mpnn_model = mk_mpnn_model() if not os.path.exists(self.pdb_filename): raise FileNotFoundError(f"Input pdb file does not exist: {self.pdb_filename}") - self.mpnn_model.prep_inputs( - pdb_filename=self.pdb_filename, - *args, - **kwargs, - ) + prep_kwargs = {**kwargs, "pdb_filename": self.pdb_filename} + self.mpnn_model.prep_inputs(*args, **prep_kwargs) self.initialized = True self.cite() @@ -98,14 +95,14 @@ def scorer(self, mutant: Mutant | RosettaPyProteinSequence, **kwargs) -> float: if isinstance(mutant, Mutant): if len(mutant.wt_protein_sequence.all_chain_ids) > 1: - warnings.warn(issues.REvoDesignWarning("MPNN only supports single chain sequences")) + warnings.warn(issues.REvoDesignWarning("MPNN only supports single chain sequences"), stacklevel=2) sequence: str = mutant.get_mutant_sequence_single_chain( chain_id=mutant.wt_protein_sequence.all_chain_ids[0], ignore_missing=True, ).sequence else: if len(mutant.all_chain_ids) > 1: - warnings.warn(issues.REvoDesignWarning("MPNN only supports single chain sequences")) + warnings.warn(issues.REvoDesignWarning("MPNN only supports single chain sequences"), stacklevel=2) sequence: str = mutant.get_sequence_by_chain(chain_id=mutant.all_chain_ids[0]).replace("X", "") # scorer must return a float score value given a mutant sequence. return self.mpnn_model.score(seq=sequence)["score"] @@ -137,7 +134,12 @@ def designer(self, *args, **kwargs): __bibtex__ = { "ProteinMPNN": r"""@article{ doi:10.1126/science.add2187, -author = {J. Dauparas and I. Anishchenko and N. Bennett and H. Bai and R. J. Ragotte and L. F. Milles and B. I. M. Wicky and A. Courbet and R. J. de Haas and N. Bethel and P. J. Y. Leung and T. F. Huddy and S. Pellock and D. Tischer and F. Chan and B. Koepnick and H. Nguyen and A. Kang and B. Sankaran and A. K. Bera and N. P. King and D. Baker }, +author = {J. Dauparas and I. Anishchenko and N. Bennett and H. Bai +and R. J. Ragotte and L. F. Milles and B. I. M. Wicky and A. Courbet +and R. J. de Haas and N. Bethel and P. J. Y. Leung and T. F. Huddy +and S. Pellock and D. Tischer and F. Chan and B. Koepnick +and H. Nguyen and A. Kang and B. Sankaran and A. K. Bera +and N. P. King and D. Baker }, title = {Robust deep learning–based protein sequence design using ProteinMPNN}, journal = {Science}, volume = {378}, @@ -147,7 +149,34 @@ def designer(self, *args, **kwargs): doi = {10.1126/science.add2187}, URL = {https://www.science.org/doi/abs/10.1126/science.add2187}, eprint = {https://www.science.org/doi/pdf/10.1126/science.add2187}, -abstract = {Although deep learning has revolutionized protein structure prediction, almost all experimentally characterized de novo protein designs have been generated using physically based approaches such as Rosetta. Here, we describe a deep learning–based protein sequence design method, ProteinMPNN, that has outstanding performance in both in silico and experimental tests. On native protein backbones, ProteinMPNN has a sequence recovery of 52.4\% compared with 32.9\% for Rosetta. The amino acid sequence at different positions can be coupled between single or multiple chains, enabling application to a wide range of current protein design challenges. We demonstrate the broad utility and high accuracy of ProteinMPNN using x-ray crystallography, cryo–electron microscopy, and functional studies by rescuing previously failed designs, which were made using Rosetta or AlphaFold, of protein monomers, cyclic homo-oligomers, tetrahedral nanoparticles, and target-binding proteins. Deep learning approaches such as Alphafold and Rosettafold have made reliable protein structure prediction broadly accessible. For the inverse problem, finding a sequence that folds to a desired structure, most approaches remain based on energy optimization. In two papers, a range of protein design problems were addressed through deep learning methods. Dauparas et al. built on recent deep learning protein design approaches to develop a method called ProteinMPNN. They validated designs experimentally and showed that ProteinMPNN can rescue previously failed designs made using Rosetta or Alphafold. Wicky et al. started from a random sequence and used Monte Carlo sequence search coupled with structure prediction by Alphafold to design cyclic homo-oligomers. Although the designs were generated to achieve stable expression, the sequences had to be regenerated using ProteinMPNN. This approach allowed for the design of a range of experimentally validated cyclic oligomers and paves the way for the design of increasingly complex assemblies. —VV A network-based protein design enables the generation of cyclic homo-oligomers across the nanoscopic scale.}} +abstract = {Although deep learning has revolutionized protein structure prediction, +almost all experimentally characterized de novo protein designs have been +generated using physically based approaches such as Rosetta. Here, we describe a +deep learning–based protein sequence design method, ProteinMPNN, that has +outstanding performance in both in silico and experimental tests. On native +protein backbones, ProteinMPNN has a sequence recovery of 52.4\% compared with +32.9\% for Rosetta. The amino acid sequence at different positions can be +coupled between single or multiple chains, enabling application to a wide range +of current protein design challenges. We demonstrate the broad utility and high +accuracy of ProteinMPNN using x-ray crystallography, cryo–electron microscopy, +and functional studies by rescuing previously failed designs, which were made +using Rosetta or AlphaFold, of protein monomers, cyclic homo-oligomers, +tetrahedral nanoparticles, and target-binding proteins. Deep learning approaches +such as Alphafold and Rosettafold have made reliable protein structure +prediction broadly accessible. For the inverse problem, finding a sequence that +folds to a desired structure, most approaches remain based on energy +optimization. In two papers, a range of protein design problems were addressed +through deep learning methods. Dauparas et al. built on recent deep learning +protein design approaches to develop a method called ProteinMPNN. They validated +designs experimentally and showed that ProteinMPNN can rescue previously failed +designs made using Rosetta or Alphafold. Wicky et al. started from a random +sequence and used Monte Carlo sequence search coupled with structure prediction +by Alphafold to design cyclic homo-oligomers. Although the designs were +generated to achieve stable expression, the sequences had to be regenerated +using ProteinMPNN. This approach allowed for the design of a range of +experimentally validated cyclic oligomers and paves the way for the design of +increasingly complex assemblies. —VV A network-based protein design enables the +generation of cyclic homo-oligomers across the nanoscopic scale.}} """ } diff --git a/src/REvoDesign/magician/designers/openkinetics/_client.py b/src/REvoDesign/magician/designers/openkinetics/_client.py index 43a6340d3..f8952985e 100644 --- a/src/REvoDesign/magician/designers/openkinetics/_client.py +++ b/src/REvoDesign/magician/designers/openkinetics/_client.py @@ -737,7 +737,8 @@ def submit( Args: self: The instance of the class containing this method - rows: A list of dictionaries, each containing 'Protein Sequence' and 'Substrate' keys with corresponding values + rows: A list of dictionaries, each containing 'Protein Sequence' and + 'Substrate' keys with corresponding values method: The prediction method to be used for the analysis prediction_type: The type of prediction to perform handle_long_sequences: Strategy for handling long sequences, defaults to "truncate" @@ -749,7 +750,8 @@ def submit( A dictionary containing the response from the OpenKinetics service, including job identifier Raises: - OpenKineticsValidationError: If no rows are provided or if the submit response doesn't contain a job identifier + OpenKineticsValidationError: If no rows are provided or if the + submit response doesn't contain a job identifier """ if not rows: raise OpenKineticsValidationError("At least one row is required for submission") @@ -786,7 +788,10 @@ def submit( logging.info("OpenKinetics service health: %s", self.check_health()) submit_response = self._request( - "POST", OPENKINETICS_ENDPOINTS["submit"], json_payload=payload, timeout=self.timeout_seconds + "POST", + OPENKINETICS_ENDPOINTS["submit"], + json_payload=payload, + timeout=self.timeout_seconds, ) # Extract and validate the job ID from the response @@ -813,11 +818,14 @@ def get_result( Args: job_id: The unique identifier of the job to retrieve results for result_format: Format of the result, either 'json' or 'csv' (default is 'json') - poll_interval_seconds: Interval in seconds between polling attempts when result is not ready (default is 3) - timeout_seconds: Maximum time to wait for the result in seconds, None means default timeout (default is None) + poll_interval_seconds: Interval in seconds between polling attempts + when result is not ready (default is 3) + timeout_seconds: Maximum time to wait for the result in seconds, + None means default timeout (default is None) Returns: - A dictionary containing the result data if format is JSON, or a string containing CSV data if format is CSV + A dictionary containing the result data if format is JSON, or a + string containing CSV data if format is CSV Raises: OpenKineticsAPIError: If the API request fails or returns an error status diff --git a/src/REvoDesign/magician/designers/openkinetics/_scorers.py b/src/REvoDesign/magician/designers/openkinetics/_scorers.py index 6d31e327f..5199a09b2 100644 --- a/src/REvoDesign/magician/designers/openkinetics/_scorers.py +++ b/src/REvoDesign/magician/designers/openkinetics/_scorers.py @@ -61,7 +61,9 @@ _PREDICTOR_BIBTEX = { "CataPro": r"""@article{Wang2025CataPro, title = {Robust enzyme discovery and engineering with deep learning using CataPro}, - author = {Wang, Zechen and Xie, Dongqi and Wu, Dong and Luo, Xiaozhou and Wang, Sheng and Li, Yangyang and Yang, Yanmei and Li, Weifeng and Zheng, Liangzhen}, + author = {Wang, Zechen and Xie, Dongqi and Wu, Dong and Luo, Xiaozhou and + Wang, Sheng and Li, Yangyang and Yang, Yanmei and Li, Weifeng and Zheng, + Liangzhen}, journal = {Nature Communications}, volume = {16}, number = {1}, @@ -81,7 +83,8 @@ }""", "DLKcat": r"""@article{Li2022DLKcat, title = {Deep learning-based kcat prediction enables improved enzyme-constrained model reconstruction}, - author = {Li, Feiran and Yuan, Le and Lu, Hongzhong and Li, Gang and Chen, Yu and Engqvist, Martin K. M. and Kerkhoven, Eduard J. and Nielsen, Jens}, + author = {Li, Feiran and Yuan, Le and Lu, Hongzhong and Li, Gang and Chen, + Yu and Engqvist, Martin K. M. and Kerkhoven, Eduard J. and Nielsen, Jens}, journal = {Nature Catalysis}, volume = {5}, number = {8}, @@ -102,7 +105,8 @@ url = {https://doi.org/10.1016/j.checat.2024.101094} }""", "KinForm": r"""@article{Alwer2026KinForm, - title = {KinForm: kinetics-informed feature optimised representation models for enzyme kcat and KM prediction}, + title = {KinForm: kinetics-informed feature optimised representation models + for enzyme kcat and KM prediction}, author = {Alwer, Saleh and Fleming, Ronan M. T.}, journal = {npj Systems Biology and Applications}, volume = {12}, @@ -112,7 +116,8 @@ url = {https://doi.org/10.1038/s41540-026-00692-5} }""", "MMISA-KM": r"""@inproceedings{Song2025MMISAKM, - title = {MMISA-KM: a deep-learning method using multi-modal information and self-attention mechanisms for the prediction of Michaelis constants}, + title = {MMISA-KM: a deep-learning method using multi-modal information and + self-attention mechanisms for the prediction of Michaelis constants}, author = {Song, Aijie and Wang, Kai}, booktitle = {2025 IEEE 14th Data Driven Control and Learning Systems (DDCLS)}, pages = {2023--2028}, @@ -121,15 +126,22 @@ url = {https://doi.org/10.1109/DDCLS66240.2025.11064981} }""", "OmniESI": r"""@misc{Nie2025OmniESI, - title = {OmniESI: A unified framework for enzyme-substrate interaction prediction with progressive conditional deep learning}, - author = {Nie, Zhiwei and Zhang, Hongyu and Jiang, Hao and Liu, Yutian and Huang, Xiansong and Xu, Fan and Fu, Jie and Ren, Zhixiang and Tian, Yonghong and Zhang, Wen-Bin and Chen, Jie}, + title = {OmniESI: A unified framework for enzyme-substrate interaction + prediction with progressive conditional deep learning}, + author = {Nie, Zhiwei and Zhang, Hongyu and Jiang, Hao and Liu, Yutian and + Huang, Xiansong and Xu, Fan and Fu, Jie and Ren, Zhixiang and Tian, Yonghong + and Zhang, Wen-Bin and Chen, Jie}, year = {2025}, doi = {10.48550/arXiv.2506.17963}, url = {https://doi.org/10.48550/arXiv.2506.17963} }""", "RealKcat": r"""@article{Sajeevan2025RealKcat, title = {Robust Prediction of Enzyme Variant Kinetics with RealKcat}, - author = {Sajeevan, Karuna Anna and Osinuga, Abraham and B, Arunraj and Ferdous, Sakib and Shahreen, Nabia and Noor, Mohammed and Koneru, Shashank and Santa-Correa, Laura Mariana and Salehi, Rahil and Chowdhury, Niaz Bahar and Aryee, Randy and Calderon-Lopez, Brisa and Mali, Ankur and Saha, Rajib and Chowdhury, Ratul}, + author = {Sajeevan, Karuna Anna and Osinuga, Abraham and B, Arunraj and + Ferdous, Sakib and Shahreen, Nabia and Noor, Mohammed and Koneru, Shashank + and Santa-Correa, Laura Mariana and Salehi, Rahil and Chowdhury, Niaz Bahar + and Aryee, Randy and Calderon-Lopez, Brisa and Mali, Ankur and Saha, Rajib + and Chowdhury, Ratul}, year = {2025}, doi = {10.1101/2025.02.10.637555}, url = {https://doi.org/10.1101/2025.02.10.637555} @@ -145,8 +157,11 @@ url = {https://doi.org/10.1038/s41467-023-44113-1} }""", "IECata": r"""@article{Wang2025IECata, - title = {IECata: interpretable bilinear attention network and evidential deep learning improve the catalytic efficiency prediction of enzymes}, - author = {Wang, Jingjing and Zhao, Yanpeng and Yang, Zhijiang and Yao, Ge and Han, Penggang and Liu, Jiajia and Chen, Chang and Zan, Peng and Wan, Xiukun and Bo, Xiaochen and Jiang, Hui}, + title = {IECata: interpretable bilinear attention network and evidential deep + learning improve the catalytic efficiency prediction of enzymes}, + author = {Wang, Jingjing and Zhao, Yanpeng and Yang, Zhijiang and Yao, Ge and + Han, Penggang and Liu, Jiajia and Chen, Chang and Zan, Peng and Wan, Xiukun + and Bo, Xiaochen and Jiang, Hui}, journal = {Briefings in Bioinformatics}, volume = {26}, number = {3}, diff --git a/src/REvoDesign/phylogenetics/evo_mutator.py b/src/REvoDesign/phylogenetics/evo_mutator.py index 36152fce5..281ed6d77 100644 --- a/src/REvoDesign/phylogenetics/evo_mutator.py +++ b/src/REvoDesign/phylogenetics/evo_mutator.py @@ -15,7 +15,6 @@ import asyncio import itertools import os -import traceback import warnings from dataclasses import dataclass from functools import partial @@ -193,7 +192,7 @@ def run_mutant_loading_from_profile(self): ) if not os.path.isdir(os.path.dirname(self.design.output_pse)): - warnings.warn(issues.NoResultsWarning("No output PyMOL session is created.")) + warnings.warn(issues.NoResultsWarning("No output PyMOL session is created."), stacklevel=2) return cmd.reinitialize() @@ -212,7 +211,7 @@ def run_mutant_loading_from_profile(self): cmd.save(output_pse) except Exception: - traceback.print_exc() + logging.exception("Error while running the design workflow.") finally: CitationManager().output() @@ -324,8 +323,7 @@ def visualize_mutants(self): cmd.save(output_pse) except Exception: - logging.error("Error while running the visualization: ") - traceback.print_exc() + logging.exception("Error while running the visualization.") finally: CitationManager().output() @@ -409,7 +407,7 @@ def _get_dist( atom2 = self._get_ca_atom(chain_2, j_1) return atom1 - atom2 except ValueError as e: - warnings.warn(f"Error calculating distance: {e}", category=UserWarning) + warnings.warn(f"Error calculating distance: {e}", category=UserWarning, stacklevel=2) return -1 def bind_chains(self, coevolved_pairs: tuple[CoevolvedPair]) -> tuple[CoevolvedPair]: @@ -616,7 +614,7 @@ def run_gremlin_tool(self): ) if not coevolved_pairs: - warnings.warn(issues.NoResultsWarning("No Available co-evolutionary signal in global")) + warnings.warn(issues.NoResultsWarning("No Available co-evolutionary signal in global"), stacklevel=2) # early return if no data. return @@ -644,7 +642,7 @@ def run_gremlin_tool(self): del chain_binder if self.coevolved_pairs.empty: - warnings.warn(issues.NoResultsWarning("No coevolved_pairs passes filter.")) + warnings.warn(issues.NoResultsWarning("No coevolved_pairs passes filter."), stacklevel=2) return logging.info("Visualizing as bonds ...") @@ -658,7 +656,9 @@ def run_gremlin_tool(self): self.bus.button("previous").clicked.disconnect() self.bus.button("next").clicked.disconnect() except Exception as e: - warnings.warn(issues.AlreadyDisconnectedWarning(f"button is already disconnected. do nothing. {e=}")) + warnings.warn( + issues.AlreadyDisconnectedWarning(f"button is already disconnected. do nothing. {e=}"), stacklevel=2 + ) self.bus.button("previous").clicked.connect(partial(self.load_co_evolving_pairs, False)) diff --git a/src/REvoDesign/phylogenetics/gremlin_tools.py b/src/REvoDesign/phylogenetics/gremlin_tools.py index bd1480dd9..b0c94f03e 100644 --- a/src/REvoDesign/phylogenetics/gremlin_tools.py +++ b/src/REvoDesign/phylogenetics/gremlin_tools.py @@ -21,7 +21,6 @@ """ import os -import traceback import warnings from dataclasses import dataclass, field from typing import Literal @@ -83,7 +82,7 @@ def is_out_of_range(self, chain_pair: str) -> bool: @property def min_dist(self): if self.empty: - warnings.warn(issues.NoInputWarning(f"Pair {repr(self)} is empty! ")) + warnings.warn(issues.NoInputWarning(f"Pair {repr(self)} is empty! "), stacklevel=2) return -1 return min(d for d in self.homochains_dist.values() if d > 0) @@ -332,7 +331,8 @@ def plot_w(self, i: int, j: int, i_aa: str, j_aa: str, idx: int = 0) -> Coevolve if not matching_indices: # No matching pairs found, handle this case warnings.warn( - issues.NoResultsWarning(f"No matching co-evolutionary pairs found for positions {i} and {j}.") + issues.NoResultsWarning(f"No matching co-evolutionary pairs found for positions {i} and {j}."), + stacklevel=2, ) return None @@ -378,10 +378,14 @@ def plot_w(self, i: int, j: int, i_aa: str, j_aa: str, idx: int = 0) -> Coevolve wt_i_index = self.alphabet.index(wt_i_aa) wt_j_index = self.alphabet.index(wt_j_aa) except ValueError: - traceback.print_exc() - logging.error(f"Error occured while processing '{wt_i_aa=}' or '{wt_j_aa=}' from {self.alphabet=}") + logging.exception( + "Error occured while processing %r or %r from %s", + wt_i_aa, + wt_j_aa, + self.alphabet, + ) # early return to skip ploting - warnings.warn(issues.BadDataWarning(f"Bad pair: {str(a_pair)}")) + warnings.warn(issues.BadDataWarning(f"Bad pair: {str(a_pair)}"), stacklevel=2) return None plt.text( @@ -451,7 +455,7 @@ def plot_w_o2a(self, resi) -> tuple[CoevolvedPair]: top_N_pairs = coevolving_pairs[: self.topN] if not top_N_pairs: - warnings.warn(issues.NoResultsWarning("No coevolving pairs found!")) + warnings.warn(issues.NoResultsWarning("No coevolving pairs found!"), stacklevel=2) return {} logging.info(f"top {self.topN} items selected: {str(top_N_pairs)}") @@ -475,7 +479,8 @@ def plot_w_o2a(self, resi) -> tuple[CoevolvedPair]: "GREMLIN": r"""@article{ doi:10.1073/pnas.1314045110, author = {Hetunandan Kamisetty and Sergey Ovchinnikov and David Baker }, -title = {Assessing the utility of coevolution-based residue–residue contact predictions in a sequence- and structure-rich era}, +title = {Assessing the utility of coevolution-based residue–residue contact +predictions in a sequence- and structure-rich era}, journal = {Proceedings of the National Academy of Sciences}, volume = {110}, number = {39}, @@ -484,5 +489,22 @@ def plot_w_o2a(self, resi) -> tuple[CoevolvedPair]: doi = {10.1073/pnas.1314045110}, URL = {https://www.pnas.org/doi/abs/10.1073/pnas.1314045110}, eprint = {https://www.pnas.org/doi/pdf/10.1073/pnas.1314045110}, -abstract = {Recently developed methods have shown considerable promise in predicting residue–residue contacts in protein 3D structures using evolutionary covariance information. However, these methods require large numbers of evolutionarily related sequences to robustly assess the extent of residue covariation, and the larger the protein family, the more likely that contact information is unnecessary because a reasonable model can be built based on the structure of a homolog. Here we describe a method that integrates sequence coevolution and structural context information using a pseudolikelihood approach, allowing more accurate contact predictions from fewer homologous sequences. We rigorously assess the utility of predicted contacts for protein structure prediction using large and representative sequence and structure databases from recent structure prediction experiments. We find that contact predictions are likely to be accurate when the number of aligned sequences (with sequence redundancy reduced to 90\%) is greater than five times the length of the protein, and that accurate predictions are likely to be useful for structure modeling if the aligned sequences are more similar to the protein of interest than to the closest homolog of known structure. These conditions are currently met by 422 of the protein families collected in the Pfam database.}}""", +abstract = {Recently developed methods have shown considerable promise in predicting +residue–residue contacts in protein 3D structures using evolutionary covariance +information. However, these methods require large numbers of evolutionarily +related sequences to robustly assess the extent of residue covariation, and the +larger the protein family, the more likely that contact information is +unnecessary because a reasonable model can be built based on the structure of a +homolog. Here we describe a method that integrates sequence coevolution and +structural context information using a pseudolikelihood approach, allowing more +accurate contact predictions from fewer homologous sequences. We rigorously +assess the utility of predicted contacts for protein structure prediction using +large and representative sequence and structure databases from recent structure +prediction experiments. We find that contact predictions are likely to be +accurate when the number of aligned sequences (with sequence redundancy reduced +to 90\%) is greater than five times the length of the protein, and that accurate +predictions are likely to be useful for structure modeling if the aligned +sequences are more similar to the protein of interest than to the closest +homolog of known structure. These conditions are currently met by 422 of the +protein families collected in the Pfam database.}}""", } diff --git a/src/REvoDesign/phylogenetics/revo_designer.py b/src/REvoDesign/phylogenetics/revo_designer.py index 39371b03d..0ba6d6584 100644 --- a/src/REvoDesign/phylogenetics/revo_designer.py +++ b/src/REvoDesign/phylogenetics/revo_designer.py @@ -64,7 +64,8 @@ def __init__(self, input_profile): self.randomized_sample_num = 10 self.mutate_runner: MutateRunnerAbstract = None # type: ignore - self.pwd = "." # nosec B105: present working directory, not a password + # Present working directory, not a password. + self.pwd = "." # nosec B105 self.sequence = "" self.design_case = "default" @@ -182,7 +183,7 @@ def plot_custom_indices_segments( "mutations": {}, } mutations = [] - for idx, resid in enumerate(custom_indices): + for resid in custom_indices: # fetch wt aa from untruncated sequence wt_aa = sequence[resid - 1] profile_scores = df_trunc.loc[:, resid] @@ -217,7 +218,8 @@ def plot_custom_indices_segments( mutation_json_fp = f"{self.pwd}/mutations_design_profile/{file_name}.json" mutation_png_fp = f"{self.pwd}/mutations_design_profile/{file_name}.png" - json.dump(mutation_candidates, open(mutation_json_fp, "w"), indent=2) + with open(mutation_json_fp, "w") as mutation_json_file: + json.dump(mutation_candidates, mutation_json_file, indent=2) plt.savefig(mutation_png_fp) plt.close() diff --git a/src/REvoDesign/shortcuts/__init__.py b/src/REvoDesign/shortcuts/__init__.py index 9811de882..31acbed24 100644 --- a/src/REvoDesign/shortcuts/__init__.py +++ b/src/REvoDesign/shortcuts/__init__.py @@ -25,8 +25,11 @@ # How to set up autocompletion for a extend command with multiple arguments? # The general format is: -# `cmd.auto_arg[][""]=[, '', '']` where -# is the argument index (starting from 0) that we want to set autocompletion for in the extend command +# `cmd.auto_arg[][""]=[ +# , '', '' +# ]` where +# is the argument index (starting from 0) that we want to set +# autocompletion for in the extend command # is the name of the extend command # is a lambda function that returns a Shortcut object # is a string label for the argument that will be displayed in the autocomplete output, @@ -40,7 +43,8 @@ # 0 1 # '' is a string that will be appended to the argument value -# Normally, to support autocomplete dynamically during runtime, PyMOL uses a lambda function that returns a Shortcut object +# Normally, to support autocomplete dynamically during runtime, PyMOL uses a +# lambda function that returns a Shortcut object # A shortcut object defines the possible keywords for autocompletion of arguments. # One can either bollow the existing Shortcut lambda from `cmd.auto_arg` or create a new one # Here we borrow the existing ones for 'show' and 'select' diff --git a/src/REvoDesign/shortcuts/function_utils.py b/src/REvoDesign/shortcuts/function_utils.py index a5c4aaf48..28be7b135 100644 --- a/src/REvoDesign/shortcuts/function_utils.py +++ b/src/REvoDesign/shortcuts/function_utils.py @@ -7,8 +7,12 @@ Utils for shortcuts """ +import atexit import os -import subprocess +import shutil + +# Used for launching a user-requested PyMOL preview with an argv list. +import subprocess # nosec B404 from typing import Literal from pymol import cmd @@ -17,6 +21,40 @@ from REvoDesign import ROOT_LOGGER logging = ROOT_LOGGER.getChild(__name__) +_PREVIEW_PROCESSES: set[subprocess.Popen[bytes]] = set() + + +def _prune_finished_preview_processes() -> None: + for process in tuple(_PREVIEW_PROCESSES): + if process.poll() is not None: + _PREVIEW_PROCESSES.discard(process) + + +def cleanup_conformer_preview_processes(timeout: float = 2.0) -> None: + """Terminate PyMOL preview processes launched by this module.""" + for process in tuple(_PREVIEW_PROCESSES): + try: + if process.poll() is None: + process.terminate() + try: + process.wait(timeout=timeout) + except subprocess.TimeoutExpired: + process.kill() + process.wait(timeout=timeout) + except Exception as exc: + logging.debug("Could not stop PyMOL preview process %s: %s", process, exc) + finally: + _PREVIEW_PROCESSES.discard(process) + + +atexit.register(cleanup_conformer_preview_processes) + + +def _resolve_pymol_executable() -> str: + pymol_executable = shutil.which("pymol") + if pymol_executable is None: + raise RuntimeError("Unable to launch PyMOL preview: 'pymol' executable was not found on PATH.") + return pymol_executable def visualize_conformer_sdf(sdf_file_path: str, show_conformer: Literal["New Window", "Current Window"]): @@ -54,10 +92,18 @@ def visualize_conformer_sdf(sdf_file_path: str, show_conformer: Literal["New Win pmlh.write("bg_color white\n") # Explicitly call a new PyMOL instance in the background - pymol_command = ["pymol", "-xi", pml_file_path] - subprocess.Popen(pymol_command) + pymol_command = [_resolve_pymol_executable(), "-xi", pml_file_path] + # Launches a user-requested PyMOL preview with an argv list and a resolved executable. + _prune_finished_preview_processes() + process = subprocess.Popen( # nosec B603 + pymol_command, + stdin=subprocess.DEVNULL, + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + ) + _PREVIEW_PROCESSES.add(process) - print(f"PyMOL launched in the background with {sdf_file_path}.") + logging.info("PyMOL launched in the background with %s.", sdf_file_path) def smiles_conformer_batch(smi: dict[str, str], num_conformer: int, save_dir: str, n_jobs: int = 1): @@ -65,7 +111,8 @@ def smiles_conformer_batch(smi: dict[str, str], num_conformer: int, save_dir: st Generates 3D conformers for a SMILES string using RDKit. Args: - smi (Dict[str, str]): Dictionary containing the name of the molecule as the key and the SMILES string as the value. + smi (Dict[str, str]): Dictionary containing the name of the molecule as + the key and the SMILES string as the value. num_conformer (int): Number of conformers to generate for each molecule. save_dir (str): Directory to save the generated conformer files. n_jobs (int, optional): Number of parallel jobs to run. Defaults to 1. diff --git a/src/REvoDesign/shortcuts/tools/designs.py b/src/REvoDesign/shortcuts/tools/designs.py index 6e46893ce..42d88ced0 100644 --- a/src/REvoDesign/shortcuts/tools/designs.py +++ b/src/REvoDesign/shortcuts/tools/designs.py @@ -31,6 +31,6 @@ def shortcut_pssm2csv(pssm: str) -> None: expected_csv = f"{pssm}.csv" if not os.path.exists(expected_csv): - warnings.warn(issues.NoResultsWarning(f"Expected {expected_csv=}")) + warnings.warn(issues.NoResultsWarning(f"Expected {expected_csv=}"), stacklevel=2) logging.info(expected_csv) diff --git a/src/REvoDesign/shortcuts/tools/esm2.py b/src/REvoDesign/shortcuts/tools/esm2.py index 6363ad670..39983a530 100644 --- a/src/REvoDesign/shortcuts/tools/esm2.py +++ b/src/REvoDesign/shortcuts/tools/esm2.py @@ -38,6 +38,19 @@ ESM_MODEL_BASE_URL = "https://dl.fbaipublicfiles.com/fair-esm/models" +# run from yinying's download archive. +# md5sum calculated by using: +# `parallel -k md5sum {} ::: $(ls esm1v_t33* esm_msa1b_t12_100M_UR50S.pt esm2_t36_3B_UR50D.pt)` +ESM1V_MODEL_MD5 = """ +1f04c2d2636b02b544ecb5fbbef8fefd esm1v_t33_650M_UR90S_1.pt +e6c58fc8107d5e6629f9ca9b910319ec esm1v_t33_650M_UR90S_2.pt +6a1a631a87cf798f72f58c0a37aeebdd esm1v_t33_650M_UR90S_3.pt +fcba894a07ec87dfea4ca366e46c78b7 esm1v_t33_650M_UR90S_4.pt +b1649837610e741f8b766b01f18890a3 esm1v_t33_650M_UR90S_5.pt +d37a0d0dbe7431e48a72072b9180b16b esm2_t36_3B_UR50D.pt +f66445cc186f77dbbbe3184399523413 esm_msa1b_t12_100M_UR50S.pt +""" + ESM1V_MODEL_DICT: immutabledict[str, str] = immutabledict( { "esm-1v_1": "esm1v_t33_650M_UR90S_1", @@ -50,8 +63,11 @@ } ) +ESM1V_WEIGHT_FILES = tuple(f"{v}.pt" for v in ESM1V_MODEL_DICT.values()) ESM1V_WEIGHTS = FileDownloadRegistry( - name="ESM2", base_url=ESM_MODEL_BASE_URL, registry={f"{v}.pt": None for v in ESM1V_MODEL_DICT.values()} + name="ESM2", + base_url=ESM_MODEL_BASE_URL, + registry=FileDownloadRegistry.prepare_registry_from_md5(ESM1V_MODEL_MD5), ) @@ -207,7 +223,6 @@ def _resolve_model_weight(self, model_name: str): logging.info(f"Loading model from {expected_model_path=}") return expected_model_path - # Download model from remote server logging.info(f"Fetching model {model_name=} from {ESM1V_WEIGHTS.base_url}") return ESM1V_WEIGHTS.setup(model_name).downloaded @@ -336,7 +351,9 @@ def compute_pppl(self, row, sequence, model, alphabet, offset_idx): __bibtex__ = { "ESM": """@article{rives2019biological, - author={Rives, Alexander and Meier, Joshua and Sercu, Tom and Goyal, Siddharth and Lin, Zeming and Liu, Jason and Guo, Demi and Ott, Myle and Zitnick, C. Lawrence and Ma, Jerry and Fergus, Rob}, + author={Rives, Alexander and Meier, Joshua and Sercu, Tom and Goyal, + Siddharth and Lin, Zeming and Liu, Jason and Guo, Demi and Ott, Myle and + Zitnick, C. Lawrence and Ma, Jerry and Fergus, Rob}, title={Biological Structure and Function Emerge from Scaling Unsupervised Learning to 250 Million Protein Sequences}, year={2019}, doi={10.1101/622803}, @@ -345,14 +362,16 @@ def compute_pppl(self, row, sequence, model, alphabet, offset_idx): }""", "ESM-1v": """@article{meier2021language, author = {Meier, Joshua and Rao, Roshan and Verkuil, Robert and Liu, Jason and Sercu, Tom and Rives, Alexander}, - title = {Language models enable zero-shot prediction of the effects of mutations on protein function}, + title = {Language models enable zero-shot prediction of the effects of + mutations on protein function}, year={2021}, doi={10.1101/2021.07.09.450648}, url={https://www.biorxiv.org/content/10.1101/2021.07.09.450648v1}, journal={bioRxiv} }""", "MSA Transformer": """@article{rao2021msa, - author = {Rao, Roshan and Liu, Jason and Verkuil, Robert and Meier, Joshua and Canny, John F. and Abbeel, Pieter and Sercu, Tom and Rives, Alexander}, + author = {Rao, Roshan and Liu, Jason and Verkuil, Robert and Meier, Joshua + and Canny, John F. and Abbeel, Pieter and Sercu, Tom and Rives, Alexander}, title={MSA Transformer}, year={2021}, doi={10.1101/2021.02.12.430858}, @@ -361,7 +380,9 @@ def compute_pppl(self, row, sequence, model, alphabet, offset_idx): }""", "ESM-2": """@article{lin2022language, title={Language models of protein sequences at the scale of evolution enable accurate structure prediction}, - author={Lin, Zeming and Akin, Halil and Rao, Roshan and Hie, Brian and Zhu, Zhongkai and Lu, Wenting and Smetanin, Nikita and dos Santos Costa, Allan and Fazel-Zarandi, Maryam and Sercu, Tom and Candido, Sal and others}, + author={Lin, Zeming and Akin, Halil and Rao, Roshan and Hie, Brian and Zhu, + Zhongkai and Lu, Wenting and Smetanin, Nikita and dos Santos Costa, Allan and + Fazel-Zarandi, Maryam and Sercu, Tom and Candido, Sal and others}, journal={bioRxiv}, year={2022}, publisher={Cold Spring Harbor Laboratory} diff --git a/src/REvoDesign/shortcuts/tools/evolution.py b/src/REvoDesign/shortcuts/tools/evolution.py index eb61754ef..b1be293f8 100644 --- a/src/REvoDesign/shortcuts/tools/evolution.py +++ b/src/REvoDesign/shortcuts/tools/evolution.py @@ -8,7 +8,9 @@ """ import gc -import pickle # nosec B403: MRF model serialization, trusted data + +# MRF model serialization uses trusted data. +import pickle # nosec B403 from REvoDesign.basic.abc_third_party_module import ThirdPartyModuleAbstract, TorchModuleAbstract from REvoDesign.bootstrap.set_config import is_package_installed diff --git a/src/REvoDesign/shortcuts/tools/ligand_converters.py b/src/REvoDesign/shortcuts/tools/ligand_converters.py index b9a928f16..4f2451e2f 100644 --- a/src/REvoDesign/shortcuts/tools/ligand_converters.py +++ b/src/REvoDesign/shortcuts/tools/ligand_converters.py @@ -62,7 +62,8 @@ def shortcut_smiles_conformer_batch( Runs the smiles_conformer_batch function with parameters collected from the dialog. """ - smi = json.load(open(smiles)) + with open(smiles) as smiles_file: + smi = json.load(smiles_file) with timing("Get SMILES Conformers (Many)"): smiles_conformer_batch( smi=smi, diff --git a/src/REvoDesign/shortcuts/tools/mutation_effect_predictors.py b/src/REvoDesign/shortcuts/tools/mutation_effect_predictors.py index 49d4fbb10..8c5bcc91d 100644 --- a/src/REvoDesign/shortcuts/tools/mutation_effect_predictors.py +++ b/src/REvoDesign/shortcuts/tools/mutation_effect_predictors.py @@ -29,8 +29,9 @@ @require_installed class ThermoMpnnPredictor(ThirdPartyModuleAbstract, TorchModuleAbstract): """ - ThermoMpnnPredictor class for predicting the thermodynamic stability effects of protein mutations. - It utilizes the ThermoMPNN model to analyze protein structure and sequence to predict stability changes due to mutations. + ThermoMpnnPredictor class for predicting the thermodynamic stability effects + of protein mutations. It utilizes the ThermoMPNN model to analyze protein + structure and sequence to predict stability changes due to mutations. Attributes: name (str): Name of the predictor, set to 'ThermoMPNN'. diff --git a/src/REvoDesign/shortcuts/tools/openmm_utils.py b/src/REvoDesign/shortcuts/tools/openmm_utils.py index f76cd6ceb..e7db4b1a2 100644 --- a/src/REvoDesign/shortcuts/tools/openmm_utils.py +++ b/src/REvoDesign/shortcuts/tools/openmm_utils.py @@ -53,7 +53,9 @@ def start_server(self): """ Start the OpenMMSetup server. - This method extends the start_server method of the parent class, configures and starts the Uvicorn server in a separate thread, and attempts to open the server URL in the default browser. + This method extends the start_server method of the parent class, + configures and starts the Uvicorn server in a separate thread, and + attempts to open the server URL in the default browser. """ """ Behavior of the server start action. @@ -104,7 +106,13 @@ def start_server(self): __bibtex__ = { "OpenMM": """@article{doi:10.1021/acs.jpcb.3c06662, -author = {Eastman, Peter and Galvelis, Raimondas and Peláez, Raúl P. and Abreu, Charlles R. A. and Farr, Stephen E. and Gallicchio, Emilio and Gorenko, Anton and Henry, Michael M. and Hu, Frank and Huang, Jing and Krämer, Andreas and Michel, Julien and Mitchell, Joshua A. and Pande, Vijay S. and Rodrigues, João PGLM and Rodriguez-Guerra, Jaime and Simmonett, Andrew C. and Singh, Sukrit and Swails, Jason and Turner, Philip and Wang, Yuanqing and Zhang, Ivy and Chodera, John D. and De Fabritiis, Gianni and Markland, Thomas E.}, +author = {Eastman, Peter and Galvelis, Raimondas and Peláez, Raúl P. and +Abreu, Charlles R. A. and Farr, Stephen E. and Gallicchio, Emilio and Gorenko, +Anton and Henry, Michael M. and Hu, Frank and Huang, Jing and Krämer, Andreas +and Michel, Julien and Mitchell, Joshua A. and Pande, Vijay S. and Rodrigues, +João PGLM and Rodriguez-Guerra, Jaime and Simmonett, Andrew C. and Singh, +Sukrit and Swails, Jason and Turner, Philip and Wang, Yuanqing and Zhang, Ivy +and Chodera, John D. and De Fabritiis, Gianni and Markland, Thomas E.}, title = {OpenMM 8: Molecular Dynamics Simulation with Machine Learning Potentials}, journal = {The Journal of Physical Chemistry B}, volume = {128}, diff --git a/src/REvoDesign/shortcuts/tools/represents.py b/src/REvoDesign/shortcuts/tools/represents.py index 867013f9a..149d8b912 100644 --- a/src/REvoDesign/shortcuts/tools/represents.py +++ b/src/REvoDesign/shortcuts/tools/represents.py @@ -612,7 +612,8 @@ def assign_to_res_pymol(self, pos_one_idx: int, state: int = 0) -> float: warnings.warn( issues.BadDataWarning( f"Failed to set B-factor for position {pos_one_idx} (zero-indexed {pos_one_idx - 1})" - ) + ), + stacklevel=2, ) return bfact @@ -801,7 +802,9 @@ def _load_to_one_chain(chain_id: str, state: int = 0): bfact_ori = bf_chain.get(pos) except (IndexError, KeyError, issues.InvalidInputError): logging.warning( - f"Position {pos} (zero-indexed {pos - 1}) exceeds the length of new B-factor data ({len(bf_chain.bfactor_data)}); setting B-factor to -1.0" + f"Position {pos} (zero-indexed {pos - 1}) exceeds the length of " + f"new B-factor data ({len(bf_chain.bfactor_data)}); setting " + "B-factor to -1.0" ) continue bfacts_orignal.append(bfact_ori) @@ -855,5 +858,5 @@ def _load_to_one_chain(chain_id: str, state: int = 0): _load_to_one_chain(chain_id) -setattr(_load_b_factors, "__bibtex__", load_b_factors_citation) +_load_b_factors.__bibtex__ = load_b_factors_citation load_b_factors = get_cited(_load_b_factors) diff --git a/src/REvoDesign/shortcuts/tools/rfdiffusion_tasks.py b/src/REvoDesign/shortcuts/tools/rfdiffusion_tasks.py index b1e8a26e7..cf1a36c0c 100644 --- a/src/REvoDesign/shortcuts/tools/rfdiffusion_tasks.py +++ b/src/REvoDesign/shortcuts/tools/rfdiffusion_tasks.py @@ -5,7 +5,9 @@ import glob import os -import pickle # nosec B403: traceback serialization, trusted context + +# Traceback serialization uses a trusted context. +import pickle # nosec B403 import random import re import shutil @@ -17,13 +19,14 @@ import numpy as np from hydra import errors as hydra_errors from omegaconf import DictConfig, OmegaConf +from pymol.Qt import QtWidgets from REvoDesign import ROOT_LOGGER, issues from REvoDesign.basic.abc_third_party_module import ThirdPartyModuleAbstract, TorchModuleAbstract from REvoDesign.bootstrap import REVODESIGN_CONFIG_FILE from REvoDesign.bootstrap.set_config import is_package_installed, reload_config_file from REvoDesign.tools.download_registry import FileDownloadRegistry -from REvoDesign.tools.package_manager import run_command +from REvoDesign.tools.package_manager import decide, run_command from REvoDesign.tools.rfdiffusion_tools import SubstratePotentialVisualizer from REvoDesign.tools.utils import device_picker, get_cited, require_installed, timing @@ -51,6 +54,12 @@ 6f4d00394d34f6a9072d70976f6c8777 RF_structure_prediction_weights.pt """ +DGL_INSTALL_APPROVAL_ENV = "REVODESIGN_ALLOW_DGL_INSTALL" + + +def _has_qapplication() -> bool: + return QtWidgets.QApplication.instance() is not None + @dataclass class DglSolver: @@ -75,10 +84,30 @@ def fetch_cuda_version_before_install(self): self.cuda_version = "" warnings.warn( issues.PlatformNotSupportedWarning( - f"CUDA version {cuda_version} is not supported by DGL. Please install CUDA version >= 11.8 if you need to use DGL with CUDA support." - ) + f"CUDA version {cuda_version} is not supported by DGL. " + "Please install CUDA version >= 11.8 if you need to use " + "DGL with CUDA support." + ), + stacklevel=2, ) + def _install_approved(self, index_link: str) -> bool: + if os.environ.get(DGL_INSTALL_APPROVAL_ENV) == "1": + return True + if not _has_qapplication(): + logging.warning("DGL installation requires interactive approval; no QApplication is available.") + return False + + return decide( + title="Install DGL?", + description=( + "RFdiffusion requires DGL, but it is not installed. " + "Install dgl==2.2.1 into the current Python environment?\n\n" + f"Package index: {index_link}" + ), + rich=True, + ) + def install(self): self.fetch_cuda_version_before_install() if self.cuda_version: @@ -86,11 +115,16 @@ def install(self): else: index_link = "https://data.dgl.ai/wheels/repo.html" + if not self._install_approved(index_link): + logging.info("DGL installation was not approved.") + return False + c = run_command([sys.executable, "-m", "pip", "install", "dgl==2.2.1", "-f", index_link]) if c.returncode != 0: logging.error(f"Failed to install DGL: {c.stderr}") raise RuntimeError(f"Failed to install DGL: {c.stderr}") self.installed = True + return True RFD_WEIGHTS = FileDownloadRegistry( @@ -165,7 +199,8 @@ def pick_model(self, model_name: str | None = None) -> None: warnings.warn( issues.FallingBackWarning( f"Model {model_name} not found. Falling back to pick one according to the input." - ) + ), + stacklevel=2, ) # automatically pick the model based on the input @@ -208,10 +243,18 @@ def ensure_dgl(): if (dgl_solver := DglSolver()).installed: return - warnings.warn(issues.MissingExternalTool("DGL is not installed. Now try to install it.")) + warnings.warn( + issues.MissingExternalTool("DGL is not installed. User approval is required before installing it."), + stacklevel=2, + ) - dgl_solver.install() + install_started = dgl_solver.install() if not dgl_solver.installed: + if not install_started: + raise issues.MissingExternalToolError( + "DGL installation was cancelled. Install dgl==2.2.1 manually or approve installation " + f"by setting {DGL_INSTALL_APPROVAL_ENV}=1." + ) raise issues.MissingExternalToolError( "DGL may not be installed. Please install it manually or take a restart to take effect." ) @@ -447,7 +490,28 @@ def main(self) -> None: volume={620}, number={7976}, pages={1089-1100}, -abstract={There has been considerable recent progress in designing new proteins using deep-learning methods1--9. Despite this progress, a general deep-learning framework for protein design that enables solution of a wide range of design challenges, including de novo binder design and design of higher-order symmetric architectures, has yet to be described. Diffusion models10,11 have had considerable success in image and language generative modelling but limited success when applied to protein modelling, probably due to the complexity of protein backbone geometry and sequence--structure relationships. Here we show that by fine-tuning the RoseTTAFold structure prediction network on protein structure denoising tasks, we obtain a generative model of protein backbones that achieves outstanding performance on unconditional and topology-constrained protein monomer design, protein binder design, symmetric oligomer design, enzyme active site scaffolding and symmetric motif scaffolding for therapeutic and metal-binding protein design. We demonstrate the power and generality of the method, called RoseTTAFold diffusion (RFdiffusion), by experimentally characterizing the structures and functions of hundreds of designed symmetric assemblies, metal-binding proteins and protein binders. The accuracy of RFdiffusion is confirmed by the cryogenic electron microscopy structure of a designed binder in complex with influenza haemagglutinin that is nearly identical to the design model. In a manner analogous to networks that produce images from user-specified inputs, RFdiffusion enables the design of diverse functional proteins from simple molecular specifications.}, +abstract={There has been considerable recent progress in designing new proteins +using deep-learning methods1--9. Despite this progress, a general deep-learning +framework for protein design that enables solution of a wide range of design +challenges, including de novo binder design and design of higher-order symmetric +architectures, has yet to be described. Diffusion models10,11 have had +considerable success in image and language generative modelling but limited +success when applied to protein modelling, probably due to the complexity of +protein backbone geometry and sequence--structure relationships. Here we show +that by fine-tuning the RoseTTAFold structure prediction network on protein +structure denoising tasks, we obtain a generative model of protein backbones +that achieves outstanding performance on unconditional and topology-constrained +protein monomer design, protein binder design, symmetric oligomer design, enzyme +active site scaffolding and symmetric motif scaffolding for therapeutic and +metal-binding protein design. We demonstrate the power and generality of the +method, called RoseTTAFold diffusion (RFdiffusion), by experimentally +characterizing the structures and functions of hundreds of designed symmetric +assemblies, metal-binding proteins and protein binders. The accuracy of +RFdiffusion is confirmed by the cryogenic electron microscopy structure of a +designed binder in complex with influenza haemagglutinin that is nearly +identical to the design model. In a manner analogous to networks that produce +images from user-specified inputs, RFdiffusion enables the design of diverse +functional proteins from simple molecular specifications.}, issn={1476-4687}, doi={10.1038/s41586-023-06415-8}, url={https://doi.org/10.1038/s41586-023-06415-8} diff --git a/src/REvoDesign/shortcuts/tools/rosetta_tasks.py b/src/REvoDesign/shortcuts/tools/rosetta_tasks.py index acfc539ce..0655c729b 100644 --- a/src/REvoDesign/shortcuts/tools/rosetta_tasks.py +++ b/src/REvoDesign/shortcuts/tools/rosetta_tasks.py @@ -45,7 +45,27 @@ class RosettaLigand(RosettaLigandOriginal, CitableModuleAbstract): doi = {https://doi.org/10.1002/prot.21086}, url = {https://onlinelibrary.wiley.com/doi/abs/10.1002/prot.21086}, eprint = {https://onlinelibrary.wiley.com/doi/pdf/10.1002/prot.21086}, -abstract = {Abstract Protein–small molecule docking algorithms provide a means to model the structure of protein–small molecule complexes in structural detail and play an important role in drug development. In recent years the necessity of simulating protein side-chain flexibility for an accurate prediction of the protein–small molecule interfaces has become apparent, and an increasing number of docking algorithms probe different approaches to include protein flexibility. Here we describe a new method for docking small molecules into protein binding sites employing a Monte Carlo minimization procedure in which the rigid body position and orientation of the small molecule and the protein side-chain conformations are optimized simultaneously. The energy function comprises van der Waals (VDW) interactions, an implicit solvation model, an explicit orientation hydrogen bonding potential, and an electrostatics model. In an evaluation of the scoring function the computed energy correlated with experimental small molecule binding energy with a correlation coefficient of 0.63 across a diverse set of 229 protein– small molecule complexes. The docking method produced lowest energy models with a root mean square deviation (RMSD) smaller than 2 Å in 71 out of 100 protein–small molecule crystal structure complexes (self-docking). In cross-docking calculations in which both protein side-chain and small molecule internal degrees of freedom were varied the lowest energy predictions had RMSDs less than 2 Å in 14 of 20 test cases. Proteins 2006. © 2006 Wiley-Liss, Inc.}, +abstract = {Abstract Protein–small molecule docking algorithms provide a means to model +the structure of protein–small molecule complexes in structural detail and play +an important role in drug development. In recent years the necessity of +simulating protein side-chain flexibility for an accurate prediction of the +protein–small molecule interfaces has become apparent, and an increasing number +of docking algorithms probe different approaches to include protein flexibility. +Here we describe a new method for docking small molecules into protein binding +sites employing a Monte Carlo minimization procedure in which the rigid body +position and orientation of the small molecule and the protein side-chain +conformations are optimized simultaneously. The energy function comprises van +der Waals (VDW) interactions, an implicit solvation model, an explicit +orientation hydrogen bonding potential, and an electrostatics model. In an +evaluation of the scoring function the computed energy correlated with +experimental small molecule binding energy with a correlation coefficient of +0.63 across a diverse set of 229 protein– small molecule complexes. The docking +method produced lowest energy models with a root mean square deviation (RMSD) +smaller than 2 Å in 71 out of 100 protein–small molecule crystal structure +complexes (self-docking). In cross-docking calculations in which both protein +side-chain and small molecule internal degrees of freedom were varied the lowest +energy predictions had RMSDs less than 2 Å in 14 of 20 test cases. Proteins +2006. © 2006 Wiley-Liss, Inc.}, year = {2006} } """, @@ -60,7 +80,14 @@ class RosettaLigand(RosettaLigandOriginal, CitableModuleAbstract): publisher="Springer New York", address="New York, NY", pages="143--155", -abstract="RosettaLigand is premiere software for predicting how a protein and a small molecule interact. Benchmark studies demonstrate that 70{\%} of the top scoring RosettaLigand predicted interfaces are within 2{\AA} RMSD from the crystal structure [1]. The latest release of Rosetta ligand software includes many new features, such as (1) docking of multiple ligands simultaneously, (2) representing ligands as fragments for greater flexibility, (3) redesign of the interface during docking, and (4) an XML script based interface that gives the user full control of the ligand docking protocol.", +abstract="RosettaLigand is premiere software for predicting how a protein and a +small molecule interact. Benchmark studies demonstrate that 70{\%} of the top +scoring RosettaLigand predicted interfaces are within 2{\AA} RMSD from the +crystal structure [1]. The latest release of Rosetta ligand software includes +many new features, such as (1) docking of multiple ligands simultaneously, (2) +representing ligands as fragments for greater flexibility, (3) redesign of the +interface during docking, and (4) an XML script based interface that gives the +user full control of the ligand docking protocol.", isbn="978-1-61779-465-0", doi="10.1007/978-1-61779-465-0_10", url="https://doi.org/10.1007/978-1-61779-465-0_10" @@ -84,7 +111,15 @@ class PROSS(PROSS_Original, CitableModuleAbstract): pages = {123-125}, year = {2020}, month = {12}, -abstract = {Many natural and designed proteins are only marginally stable limiting their usefulness in research and applications. Recently, we described an automated structure and sequence-based design method, called PROSS, for optimizing protein stability and heterologous expression levels that has since been validated on dozens of proteins. Here, we introduce improvements to the method, workflow and presentation, including more accurate sequence analysis, error handling and automated analysis of the quality of the sequence alignment that is used in design calculations.PROSS2 is freely available for academic use at https://pross.weizmann.ac.il. }, +abstract = {Many natural and designed proteins are only marginally stable limiting +their usefulness in research and applications. Recently, we described an +automated structure and sequence-based design method, called PROSS, for +optimizing protein stability and heterologous expression levels that has since +been validated on dozens of proteins. Here, we introduce improvements to the +method, workflow and presentation, including more accurate sequence analysis, +error handling and automated analysis of the quality of the sequence alignment +that is used in design calculations.PROSS2 is freely available for academic use +at https://pross.weizmann.ac.il. }, issn = {1367-4803}, doi = {10.1093/bioinformatics/btaa1071}, url = {https://doi.org/10.1093/bioinformatics/btaa1071}, @@ -93,7 +128,20 @@ class PROSS(PROSS_Original, CitableModuleAbstract): """, "PROSS": """ -@article{10.1016/j.molcel.2016.06.012, author = {Goldenzweig, A. and Goldsmith, M. and Hill, S. E. and Gertman, O. and Laurino, P. and Ashani, Y. and Dym, O. and Unger, T. and Albeck, S. and Prilusky, J. and Lieberman, R. L. and Aharoni, A. and Silman, I. and Sussman, J. L. and Tawfik, D. S. and Fleishman, S. J.}, title = {Automated structure- and sequence-based design of proteins for high bacterial expression and stability}, journal = {Molecular Cell}, year = {2016}, volume = {63}, issue = {2}, pages = {337-346}, doi = {10.1016/j.molcel.2016.06.012} }""", +@article{10.1016/j.molcel.2016.06.012, +author = {Goldenzweig, A. and Goldsmith, M. and Hill, S. E. and Gertman, O. +and Laurino, P. and Ashani, Y. and Dym, O. and Unger, T. and Albeck, S. +and Prilusky, J. and Lieberman, R. L. and Aharoni, A. and Silman, I. +and Sussman, J. L. and Tawfik, D. S. and Fleishman, S. J.}, +title = {Automated structure- and sequence-based design of proteins for high +bacterial expression and stability}, +journal = {Molecular Cell}, +year = {2016}, +volume = {63}, +issue = {2}, +pages = {337-346}, +doi = {10.1016/j.molcel.2016.06.012} +}""", } ) diff --git a/src/REvoDesign/shortcuts/tools/vina_tools.py b/src/REvoDesign/shortcuts/tools/vina_tools.py index b1466b4a8..a56363dfe 100644 --- a/src/REvoDesign/shortcuts/tools/vina_tools.py +++ b/src/REvoDesign/shortcuts/tools/vina_tools.py @@ -9,8 +9,8 @@ """ +import secrets from dataclasses import dataclass -from random import randint from typing import Literal, overload import numpy as np @@ -27,6 +27,15 @@ from ...tools.cgo_utils import LineVertex, Point, PolyLines, Sphere logging = ROOT_LOGGER.getChild(__name__) + + +def _unique_pymol_name(prefix: str) -> str: + name = f"{prefix}_{secrets.randbelow(10001)}" + while name in cmd.get_names(): + name = f"{prefix}_{secrets.randbelow(10001)}" + return name + + ############################################################################## # GetBox Plugin.py -- Draws a box surrounding a selection and gets box information # This script is used to get box information for LeDock, Autodock Vina and AutoDock Vina. @@ -163,7 +172,7 @@ def rebuild(self): self.always_left_corner = bool(self.always_left_corner) self._data = [] - for (idxa, axis), (idxc, colorname) in zip(enumerate("xyz"), enumerate("rgb")): + for axis, colorname in zip("xyz", "rgb"): p2_kwargs = {i: 0.0 for i in "xyz" if i != axis} self._data.extend( Cylinder( @@ -248,7 +257,8 @@ class CgoBox(GraphicObject): def rebuild(self): """ - Post-initialization processing. Deletes existing objects with the same name and ensures all coordinates are floats. + Post-initialization processing. Deletes existing objects with the same + name and ensures all coordinates are floats. """ self.size_xyz = tuple(abs(x) for x in self.p1.delta_xyz(self.p2)) @@ -270,8 +280,10 @@ def to_vina(self): """ Generates a string of parameters for configuring a binding site in AutoDock Vina. """ - center = f"--center_x {self.cen_xyz[0]:.1f} --center_y {self.cen_xyz[1]:.1f} --center_z {self.cen_xyz[2]:.1f}" - size = f"--size_x {self.size_xyz[0]:.1f} --size_y {self.size_xyz[1]:.1f} --size_z {self.size_xyz[2]:.1f}" + center = ( + f"--center_x {self.cen_xyz[0]:.1f} --center_y {self.cen_xyz[1]:.1f} " f"--center_z {self.cen_xyz[2]:.1f}" + ) + size = f"--size_x {self.size_xyz[0]:.1f} --size_y {self.size_xyz[1]:.1f} " f"--size_z {self.size_xyz[2]:.1f}" return f"{center} {size}" @property @@ -307,10 +319,17 @@ def __repr__(self) -> str: """ Provides a string representation of the box, including its coordinates, size, and center. """ - return f"""CgoBox `{self.name}`: -Coordinates: {self.p1.x:.3f} - {self.p2.x:.3f}; {self.p1.y:.3f} - {self.p2.y:.3f}, {self.p1.z:.3f} - {self.p2.z:.3f} -Size: {self.size_xyz[0]:.3f} * {self.size_xyz[1]:.3f} * {self.size_xyz[2]:.3f} = {self.size_xyz[0] * self.size_xyz[1] * self.size_xyz[2]:.3f} -Center: {self.cen_xyz[0]:.3f}, {self.cen_xyz[1]:.3f}, {self.cen_xyz[2]:.3f}""" + coordinates = ( + f"Coordinates: {self.p1.x:.3f} - {self.p2.x:.3f}; " + f"{self.p1.y:.3f} - {self.p2.y:.3f}, {self.p1.z:.3f} - {self.p2.z:.3f}" + ) + size = ( + f"Size: {self.size_xyz[0]:.3f} * {self.size_xyz[1]:.3f} * " + f"{self.size_xyz[2]:.3f} = " + f"{self.size_xyz[0] * self.size_xyz[1] * self.size_xyz[2]:.3f}" + ) + center = f"Center: {self.cen_xyz[0]:.3f}, {self.cen_xyz[1]:.3f}, {self.cen_xyz[2]:.3f}" + return f"CgoBox `{self.name}`:\n{coordinates}\n{size}\n{center}" @classmethod def from_selection( @@ -356,9 +375,7 @@ def from_selection( """ if not box_name: - boxName = "box_" + str(randint(0, 10000)) - while boxName in cmd.get_names(): - boxName = "box_" + str(randint(0, 10000)) + boxName = _unique_pymol_name("box") else: boxName = box_name @@ -375,8 +392,10 @@ def from_selection( return box +# fmt: off @overload -def showbox(box: str, minX: float, maxX: float, minY: float, maxY: float, minZ: float, maxZ: float) -> CgoBox: ... +def showbox(box: str, minX: float, maxX: float, minY: float, maxY: float, minZ: float, maxZ: float) -> CgoBox: + ... @overload @@ -388,7 +407,9 @@ def showbox( maxY: float | str | None = None, minZ: float | str | None = None, maxZ: float | str | None = None, -) -> CgoBox: ... +) -> CgoBox: + ... +# fmt: on def showbox( @@ -425,11 +446,14 @@ def showbox( # Validate that all required coordinates are provided when box is a string if any(x is None for x in [minX, maxX, minY, maxY, minZ, maxZ]): raise ValueError( - "To make a box, you must specify minX, maxX, minY, maxY, minZ, maxZ as valid floats or float-like strings." + "To make a box, you must specify minX, maxX, minY, maxY, minZ, " + "maxZ as valid floats or float-like strings." ) # Create a new CgoBox object from provided parameters box = CgoBox( - name=box, p1=Point(float(minX), float(minY), float(minZ)), p2=Point(float(maxX), float(maxY), float(maxZ)) + name=box, + p1=Point(float(minX), float(minY), float(minZ)), + p2=Point(float(maxX), float(maxY), float(maxZ)), ) # Print box details in different formats @@ -534,9 +558,13 @@ def removeions(): """ Remove specified ions from the molecular model. - This function selects and removes ions such as PO4, SO4, ZN, CA, MG, and CL from the molecular model in the PyMOL environment. + This function selects and removes ions such as PO4, SO4, ZN, CA, MG, and CL + from the molecular model in the PyMOL environment. """ - cmd.select("Ions", "((resn PO4) | (resn SO4) | (resn ZN) | (resn CA) | (resn MG) | (resn CL)) & hetatm") + cmd.select( + "Ions", + "((resn PO4) | (resn SO4) | (resn ZN) | (resn CA) | (resn MG) | (resn CL)) & hetatm", + ) cmd.remove("Ions") cmd.delete("Ions") return @@ -654,7 +682,7 @@ def plot_pca_box(orig_vertices, new_box_name: str = "pca_box"): goc = GOC([]) # Create a PolyLines object for each face - for i, face_indices in enumerate(faces): + for face_indices in faces: face_vertices = [orig_vertices[idx] for idx in face_indices] goc.objects.append( PolyLines(width=2.0, color="cyan", points=LineVertex.from_points(face_vertices), line_type="LINE_LOOP") @@ -669,7 +697,7 @@ def plot_pca_box(orig_vertices, new_box_name: str = "pca_box"): ).load_as(f"pca_axes_{_idx}") # Plot each vertex as a sphere - for i, v in enumerate(orig_vertices): + for v in orig_vertices: goc.objects.append(Sphere(Point(*v), radius=1)) # Rebuild the GOC object and load the new box @@ -683,7 +711,8 @@ def get_pca_box(selection="(sele)", new_box_name: str | None = None, extending=5 Parameters: - selection: A string defining the selection range for which the PCA box is generated. Defaults to "(sele)". - - new_box_name: An optional string specifying the name of the new box. If not provided, a unique name will be generated. + - new_box_name: An optional string specifying the name of the new box. If + not provided, a unique name will be generated. - extending: A float value representing the padding to be added to the oriented bounding box. Defaults to 5.0. Returns: @@ -691,10 +720,7 @@ def get_pca_box(selection="(sele)", new_box_name: str | None = None, extending=5 """ if not new_box_name: # Generate a unique box name if none is provided - boxName = "pca_box_" + str(randint(0, 10000)) - while boxName in cmd.get_names(): - # Ensure the generated box name is unique - boxName = "pca_box_" + str(randint(0, 10000)) + boxName = _unique_pymol_name("pca_box") else: # Use the provided box name boxName = new_box_name diff --git a/src/REvoDesign/shortcuts/utils.py b/src/REvoDesign/shortcuts/utils.py index d13340b73..7924a7f79 100644 --- a/src/REvoDesign/shortcuts/utils.py +++ b/src/REvoDesign/shortcuts/utils.py @@ -160,7 +160,8 @@ def _build_asked_value(entry: dict) -> AskedValue: AskedValue: A fully constructed `AskedValue` object ready for use in a dialog. Raises: - ValueError: If there's an error resolving dynamic choices via [resolve_choice_from](file:///Users/yyy/Documents/protein_design/REvoDesign/src/REvoDesign/shortcuts/utils.py#L106-L159). + ValueError: If there's an error resolving dynamic choices via + `resolve_choice_from`. Any exceptions raised during callable execution will propagate up. """ # Get type diff --git a/src/REvoDesign/sidechain/mutate_runner/DLPacker.py b/src/REvoDesign/sidechain/mutate_runner/DLPacker.py index 37f15cf1b..6ed486ce4 100644 --- a/src/REvoDesign/sidechain/mutate_runner/DLPacker.py +++ b/src/REvoDesign/sidechain/mutate_runner/DLPacker.py @@ -229,7 +229,19 @@ def run_mutate_parallel( doi = {https://doi.org/10.1002/prot.26311}, url = {https://onlinelibrary.wiley.com/doi/abs/10.1002/prot.26311}, eprint = {https://onlinelibrary.wiley.com/doi/pdf/10.1002/prot.26311}, -abstract = {Abstract Prediction of side chain conformations of amino acids in proteins (also termed “packing”) is an important and challenging part of protein structure prediction with many interesting applications in protein design. A variety of methods for packing have been developed but more accurate ones are still needed. Machine learning (ML) methods have recently become a powerful tool for solving various problems in diverse areas of science, including structural biology. In this study, we evaluate the potential of deep neural networks (DNNs) for prediction of amino acid side chain conformations. We formulate the problem as image-to-image transformation and train a U-net style DNN to solve the problem. We show that our method outperforms other physics-based methods by a significant margin: reconstruction RMSDs for most amino acids are about 20\% smaller compared to SCWRL4 and Rosetta Packer with RMSDs for bulky hydrophobic amino acids Phe, Tyr, and Trp being up to 50\% smaller.}, +abstract = {Abstract Prediction of side chain conformations of amino acids in proteins +(also termed “packing”) is an important and challenging part of protein +structure prediction with many interesting applications in protein design. A +variety of methods for packing have been developed but more accurate ones are +still needed. Machine learning (ML) methods have recently become a powerful tool +for solving various problems in diverse areas of science, including structural +biology. In this study, we evaluate the potential of deep neural networks (DNNs) +for prediction of amino acid side chain conformations. We formulate the problem +as image-to-image transformation and train a U-net style DNN to solve the +problem. We show that our method outperforms other physics-based methods by a +significant margin: reconstruction RMSDs for most amino acids are about 20\% +smaller compared to SCWRL4 and Rosetta Packer with RMSDs for bulky hydrophobic +amino acids Phe, Tyr, and Trp being up to 50\% smaller.}, year = {2022} } diff --git a/src/REvoDesign/sidechain/mutate_runner/DunbrackRotamerLib.py b/src/REvoDesign/sidechain/mutate_runner/DunbrackRotamerLib.py index a8771ca66..14f534af3 100644 --- a/src/REvoDesign/sidechain/mutate_runner/DunbrackRotamerLib.py +++ b/src/REvoDesign/sidechain/mutate_runner/DunbrackRotamerLib.py @@ -119,7 +119,8 @@ def run_mutate_parallel(self, mutants: list[Mutant], nproc: int = 2) -> list[str "Dunbrack Rotamer Library": """@Article{Shapovalov2011, author={Shapovalov, Maxim V. and Dunbrack Jr., Roland L.,}, -title={A Smoothed Backbone-Dependent Rotamer Library for Proteins Derived from Adaptive Kernel Density Estimates and Regressions}, +title={A Smoothed Backbone-Dependent Rotamer Library for Proteins Derived from +Adaptive Kernel Density Estimates and Regressions}, journal={Structure}, year={2011}, month={Jun}, @@ -128,7 +129,21 @@ def run_mutate_parallel(self, mutants: list[Mutant], nproc: int = 2) -> list[str volume={19}, number={6}, pages={844-858}, -abstract={Rotamer libraries are used in protein structure determination, prediction, and design. The backbone-dependent rotamer library consists of rotamer frequencies, mean dihedral angles, and variances as a function of the backbone dihedral angles. Structure prediction and design methods that employ backbone flexibility would strongly benefit from smoothly varying probabilities and angles. A new version of the backbone-dependent rotamer library has been developed using adaptive kernel density estimates for the rotamer frequencies and adaptive kernel regression for the mean dihedral angles and variances. This formulation allows for evaluation of the rotamer probabilities, mean angles, and variances as a smooth and continuous function of phi and psi. Continuous probability density estimates for the nonrotameric degrees of freedom of amides, carboxylates, and aromatic side chains have been modeled as a function of the backbone dihedrals and rotamers of the remaining degrees of freedom. New backbone-dependent rotamer libraries at varying levels of smoothing are available from http://dunbrack.fccc.edu.}, +abstract={Rotamer libraries are used in protein structure determination, prediction, +and design. The backbone-dependent rotamer library consists of rotamer +frequencies, mean dihedral angles, and variances as a function of the backbone +dihedral angles. Structure prediction and design methods that employ backbone +flexibility would strongly benefit from smoothly varying probabilities and +angles. A new version of the backbone-dependent rotamer library has been +developed using adaptive kernel density estimates for the rotamer frequencies +and adaptive kernel regression for the mean dihedral angles and variances. This +formulation allows for evaluation of the rotamer probabilities, mean angles, and +variances as a smooth and continuous function of phi and psi. Continuous +probability density estimates for the nonrotameric degrees of freedom of +amides, carboxylates, and aromatic side chains have been modeled as a function +of the backbone dihedrals and rotamers of the remaining degrees of freedom. New +backbone-dependent rotamer libraries at varying levels of smoothing are +available from http://dunbrack.fccc.edu.}, issn={0969-2126}, doi={10.1016/j.str.2011.03.019}, url={https://doi.org/10.1016/j.str.2011.03.019} diff --git a/src/REvoDesign/sidechain/mutate_runner/RosettaMutateRelax.py b/src/REvoDesign/sidechain/mutate_runner/RosettaMutateRelax.py index 7562538bb..29209ba86 100644 --- a/src/REvoDesign/sidechain/mutate_runner/RosettaMutateRelax.py +++ b/src/REvoDesign/sidechain/mutate_runner/RosettaMutateRelax.py @@ -209,12 +209,24 @@ def run_mutate_parallel( List[str]: List of output PDB file paths for all mutants. """ # Refresh node configuration before each run - self.mutate_relax_instance.node = self.node_hint, rosetta_utils.read_rosetta_node_config() + self.mutate_relax_instance.node = ( + self.node_hint, + rosetta_utils.read_rosetta_node_config(), + ) self.mutate_relax_instance.run(mutants, opts=list(self.rosetta_general_opts)) return [os.path.join(self.temp_dir, f"{mutant.short_mutant_id}.pdb") for mutant in mutants] __bibtex__ = rosetta_utils.copy_rosetta_citation( { - "Relax": """@article{10.1002/pro.2389, author = {Conway, P. and Tyka, M. D. and DiMaio, F. and Konerding, D. E. and Baker, D.}, title = {Relaxation of backbone bond geometry improves protein energy landscape modeling}, journal = {Protein Science}, year = {2013}, volume = {23}, issue = {1}, pages = {47-55}, doi = {10.1002/pro.2389} }""" + "Relax": """@article{10.1002/pro.2389, +author = {Conway, P. and Tyka, M. D. and DiMaio, F. and Konerding, D. E. and Baker, D.}, +title = {Relaxation of backbone bond geometry improves protein energy landscape modeling}, +journal = {Protein Science}, +year = {2013}, +volume = {23}, +issue = {1}, +pages = {47-55}, +doi = {10.1002/pro.2389} +}""" } ) diff --git a/src/REvoDesign/sidechain/sidechain_solver.py b/src/REvoDesign/sidechain/sidechain_solver.py index 9ae655ab5..5ed7e2219 100644 --- a/src/REvoDesign/sidechain/sidechain_solver.py +++ b/src/REvoDesign/sidechain/sidechain_solver.py @@ -61,7 +61,7 @@ class SidechainSolverConfig: def reconfigured(self, new_config: "SidechainSolverConfig") -> bool: reconfigured = False if new_config != self: - for (k1, v1), (k2, v2) in zip(asdict(self).items(), asdict(new_config).items()): + for (k1, v1), (_k2, v2) in zip(asdict(self).items(), asdict(new_config).items()): if v1 == v2: continue print_diff(k1, {"Before": v1, "After": v2}) diff --git a/src/REvoDesign/structure/PocketSearcher.py b/src/REvoDesign/structure/PocketSearcher.py index 92e573d17..b11f24bde 100644 --- a/src/REvoDesign/structure/PocketSearcher.py +++ b/src/REvoDesign/structure/PocketSearcher.py @@ -79,7 +79,8 @@ def search_pockets(self): ) cmd.select( substrate_pocket_id, - f"({self.molecule} and c. {self.chain_id}) and (byres ({ligand_sele}) around {self.ligand_radius}) and polymer.protein", + f"({self.molecule} and c. {self.chain_id}) and " + f"(byres ({ligand_sele}) around {self.ligand_radius}) and polymer.protein", ) cmd.select( design_shell_id, @@ -94,13 +95,14 @@ def search_pockets(self): if self.cofactor and self.cofactor_radius > 0: cofact_label, cofact_sele = self.process_multiple_resn(self.cofactor) - logging.debug(f"cofactor info {self.cofactor} ({cofact_label}: {cofact_sele}): {self.cofactor_radius}") + logging.debug(f"cofactor info {self.cofactor} ({cofact_label}: {cofact_sele}): " f"{self.cofactor_radius}") cofactor_pocket_id = cmd.get_unused_name(f"pkt_cof_{cofact_label}_{self.cofactor_radius}_") logging.info(f"Setting cofactor {self.cofactor}: {cofactor_pocket_id}") cmd.select( cofactor_pocket_id, - f"({self.molecule} and c. {self.chain_id}) and (byres ({cofact_sele}) around {self.cofactor_radius}) and polymer.protein", + f"({self.molecule} and c. {self.chain_id}) and " + f"(byres ({cofact_sele}) around {self.cofactor_radius}) and polymer.protein", ) cmd.select( design_shell_id, diff --git a/src/REvoDesign/tools/cgo_utils.py b/src/REvoDesign/tools/cgo_utils.py index 723392f87..3496a4798 100644 --- a/src/REvoDesign/tools/cgo_utils.py +++ b/src/REvoDesign/tools/cgo_utils.py @@ -6,12 +6,16 @@ """ Utilities for pymol.cgo, partially modified from `pymol.cgobuilder` -This module is aimed at building programable application protocol for CGO creating via pymol.cgo module +This module is aimed at building programable application protocol for CGO +creating via pymol.cgo module # docs of cgo: -Lowercase names below should be replaced with floating-point numbers. Generally, the TRIANGLE primitive should only be used only as a last restore since it is much less effective to render than using a series of vertices with a BEGIN/END group. +Lowercase names below should be replaced with floating-point numbers. +Generally, the TRIANGLE primitive should only be used only as a last restore +since it is much less effective to render than using a series of vertices with a +BEGIN/END group. BEGIN, { POINTS | LINES | LINE_LOOP | LINE_STRIP | TRIANGLES | TRIANGLE_STRIP | TRIANGLE_FAN }, @@ -204,14 +208,16 @@ def as_vertex(self): def as_normal(self): """ Generate a CGO normal command for the point - This method inserts the point's coordinates into a CGO normal command, used for specifying normals in PyMOL. + This method inserts the point's coordinates into a CGO normal command, + used for specifying normals in PyMOL. """ return np.insert(cgo.NORMAL, 1, self.array) def move(self, x: float | None = None, y: float | None = None, z: float | None = None) -> "Point": """ Move the point - This method allows the point to be moved along the x, y, and z axes. If a coordinate is not provided, the original value is used. + This method allows the point to be moved along the x, y, and z axes. If + a coordinate is not provided, the original value is used. Parameters: - x: Optional[float] = None, the new x-coordinate, if not provided, the original x-coordinate is used @@ -227,7 +233,8 @@ def move(self, x: float | None = None, y: float | None = None, z: float | None = def as_arrays(points: Iterable["Point"]): """ Convert a collection of points to a numpy array - This static method converts a collection of Point objects into a single numpy array, facilitating batch processing. + This static method converts a collection of Point objects into a single + numpy array, facilitating batch processing. Parameters: - points: Iterable['Point'], a collection of Point objects @@ -241,7 +248,8 @@ def as_arrays(points: Iterable["Point"]): def as_vertexes(points: Iterable["Point"]): """ Convert a collection of points to CGO vertex commands - This static method converts a collection of Point objects into CGO vertex commands, used for batch rendering in PyMOL. + This static method converts a collection of Point objects into CGO vertex + commands, used for batch rendering in PyMOL. Parameters: - points: Iterable['Point'], a collection of Point objects @@ -779,7 +787,8 @@ class LineVertex(GraphicObject): """ Represents a line vertex, inheriting from GraphicObject. - This class is used to define a line drawing element, which can be a point or a Bezier curve, and can include line width and color attributes. + This class is used to define a line drawing element, which can be a point or + a Bezier curve, and can include line width and color attributes. Attributes: - point: A Point or Bezier instance, representing the starting point or control point of the line. @@ -795,7 +804,8 @@ def rebuild(self): """ Rebuilds the line vertex data. - This method initializes the internal data list, and rebuilds the line vertex data based on the width, color, and point type. + This method initializes the internal data list, and rebuilds the line + vertex data based on the width, color, and point type. """ # Initialize the data list self._data = [] @@ -843,7 +853,8 @@ def rebuild(self): """ Rebuilds the sphere's data representation using CGO (Chimera Graphics Object) format. - This method constructs the sphere's data by combining the color information and the sphere's geometric properties. + This method constructs the sphere's data by combining the color + information and the sphere's geometric properties. """ self._data = [ *Color(self.color).as_cgo, # Convert color to CGO format and unpack it into the data list @@ -1012,7 +1023,9 @@ class Cone(GraphicObject): radius_base: The radius of the cone at the base. color_tip: The color of the cone's tip. Default is 'w' for white. color_base: The color of the cone's base. Default is 'g' for green. - caps: A tuple indicating whether to add caps to the tip and/or base. 1 for True, 0 for False. Default is (1, 0), meaning the tip has a cap and the base does not. + caps: A tuple indicating whether to add caps to the tip and/or base. + 1 for True, 0 for False. Default is (1, 0), meaning the tip has a + cap and the base does not. """ tip: Point @@ -1251,7 +1264,8 @@ def rebuild(self): """ Rebuilds the square object's drawing data. - This method initializes a graphic object drawing data (_data) by defining the vertices and colors of the triangles that make up the square. + This method initializes a graphic object drawing data (_data) by defining + the vertices and colors of the triangles that make up the square. It uses the Color utility class to handle color conversion and the Point class for vertex coordinates. """ # Start drawing, defining the drawing mode as triangles @@ -1282,13 +1296,16 @@ class PolyLines(GraphicObject): """ Represents a collection of polylines, inheriting from GraphicObject. - This class is used to define a series of lines connected in a specific way, with common attributes such as width and color. + This class is used to define a series of lines connected in a specific way, + with common attributes such as width and color. Attributes: width (float): The line width. color (str): The line color, represented as a string. points (Iterable[LineVertex]): A collection of line vertices. - line_type (Literal['LINE_STRIP', 'LINE_LOOP', 'TRIANGLE_STRIP', 'TRIANGLE_FAN']): The drawing mode of the line, defaulting to 'LINE_STRIP'. + line_type (Literal['LINE_STRIP', 'LINE_LOOP', 'TRIANGLE_STRIP', + 'TRIANGLE_FAN']): The drawing mode of the line, defaulting to + 'LINE_STRIP'. """ # global @@ -1302,7 +1319,8 @@ def rebuild(self): """ Rebuilds the line data. - This method initializes the line drawing data, including setting the line width, color, and type, and updates the data for each vertex. + This method initializes the line drawing data, including setting the line + width, color, and type, and updates the data for each vertex. """ # Initialize the line drawing data, including line width and color self._data = [ @@ -1322,7 +1340,8 @@ def rebuild(self): @dataclass class Arrow(GraphicObject): """ - Represents an arrow object for visualization in PyMOL, with properties for start and end points, line width, and color. + Represents an arrow object for visualization in PyMOL, with properties for + start and end points, line width, and color. """ start: Point # the start point of the arrow @@ -2135,7 +2154,8 @@ def rebuild(self): """ Rebuilds the data for all graphic objects in the collection. - This method empties the existing graphic object data, then iterates through each graphic object in the collection. + This method empties the existing graphic object data, then iterates + through each graphic object in the collection. If the force_to_rebuild flag is set to True, it calls the rebuild method on each graphic object. Finally, it adds the data of each graphic object to the collection's _data list. """ @@ -2483,7 +2503,7 @@ def rebuild(self): # with timing('writing'): # test_text() -# font_path = "/Users/yyy/Downloads/simhei.ttf" +# font_path = "/path/to/simhei.ttf" # text='Silver Bullet\n\nCool Kid' diff --git a/src/REvoDesign/tools/customized_widgets.py b/src/REvoDesign/tools/customized_widgets.py index 24c64213d..a125a0108 100644 --- a/src/REvoDesign/tools/customized_widgets.py +++ b/src/REvoDesign/tools/customized_widgets.py @@ -99,7 +99,9 @@ def __init__(self, object_name: str | None = None, allow_repeat: bool = False, p super().__init__(parent=parent) if not parent: - self.setFont(CURRENT_FONT or DEFAULT_FONT) + font = CURRENT_FONT or DEFAULT_FONT + if font is not None: + self.setFont(font) self.setObjectName(object_name or "AnonymousWidget") self.allow_repeat = allow_repeat @@ -115,7 +117,7 @@ def __init__(self, object_name: str | None = None, allow_repeat: bool = False, p try: self.check_repeat() except RuntimeError as e: - warnings.warn(issues.REvoDesignWidgetWarning(e)) + warnings.warn(issues.REvoDesignWidgetWarning(e), stacklevel=2) self.destroy() raise RuntimeError(f"a window named {self.objectName()} is already open.") from e @@ -142,7 +144,8 @@ def show(self): def check_repeat(self): """ - Checks if a window with the same name is already open. If found, raises it to the front and raises a RuntimeError. + Checks if a window with the same name is already open. If found, raises + it to the front and raises a RuntimeError. """ from REvoDesign.driver.ui_driver import ConfigBus @@ -154,9 +157,7 @@ def check_repeat(self): # Find windows with the same name the_windows = [ - w - for w in bus.ui.open_windows - if hasattr(w, "objectName") and getattr(w, "objectName")() == self.objectName() + w for w in bus.ui.open_windows if hasattr(w, "objectName") and w.objectName() == self.objectName() ] if any(the_windows): # Raise the existing window to the front @@ -180,7 +181,7 @@ def attach(self): # Ensure the open_windows list exists and add this widget to it if not hasattr(bus.ui, "open_windows"): logging.debug(f"Creating open_windows list under ui for {self.objectName()}") - setattr(bus.ui, "open_windows", []) + bus.ui.open_windows = [] bus.ui.open_windows.append(self) logging.debug(f"Window {self.objectName()} attached.") @@ -1362,38 +1363,48 @@ def getOpenFileNameWithExt(*args, **kwargs): return fname +# fmt: off @overload -def set_widget_value(widget: QtWidgets.QFontComboBox, value: QtGui.QFont | str): ... +def set_widget_value(widget: QtWidgets.QFontComboBox, value: QtGui.QFont | str): + ... @overload -def set_widget_value(widget: QtWidgets.QStackedWidget, value: list): ... +def set_widget_value(widget: QtWidgets.QStackedWidget, value: list): + ... @overload -def set_widget_value(widget: QtWidgets.QProgressBar, value: int | list[int] | tuple[int, int]): ... +def set_widget_value(widget: QtWidgets.QProgressBar, value: int | list[int] | tuple[int, int]): + ... @overload def set_widget_value( widget: QtWidgets.QDoubleSpinBox | QtWidgets.QSpinBox, value: int | float | list[str] | tuple[str, str] -): ... +): + ... @overload -def set_widget_value(widget: MultiCheckableComboBox, value: list | tuple | str | int | float): ... +def set_widget_value(widget: MultiCheckableComboBox, value: list | tuple | str | int | float): + ... @overload -def set_widget_value(widget: QtWidgets.QComboBox, value: list | tuple | dict | str | int | float | bool): ... +def set_widget_value(widget: QtWidgets.QComboBox, value: list | tuple | dict | str | int | float | bool): + ... @overload -def set_widget_value(widget: QtWidgets.QGridLayout, value: str): ... +def set_widget_value(widget: QtWidgets.QGridLayout, value: str): + ... @overload -def set_widget_value(widget: QtWidgets.QLineEdit | QtWidgets.QLCDNumber | QtWidgets.QCheckBox, value: Any): ... +def set_widget_value(widget: QtWidgets.QLineEdit | QtWidgets.QLCDNumber | QtWidgets.QCheckBox, value: Any): + ... +# fmt: on def set_widget_value(widget, value): @@ -1494,25 +1505,30 @@ def set_value_error(widget: QtWidgets.QWidget, value: Any): @overload -def get_widget_value(widget: QtWidgets.QCheckBox) -> bool: ... # type: ignore +def get_widget_value(widget: QtWidgets.QCheckBox) -> bool: # type: ignore + ... @overload def get_widget_value( widget: QtWidgets.QComboBox | QtWidgets.QLineEdit | QtWidgets.QFontComboBox, -) -> str: ... # type: ignore +) -> str: # type: ignore + ... @overload -def get_widget_value(widget: QtWidgets.QDoubleSpinBox | QtWidgets.QLCDNumber) -> float: ... # type: ignore +def get_widget_value(widget: QtWidgets.QDoubleSpinBox | QtWidgets.QLCDNumber) -> float: # type: ignore + ... @overload -def get_widget_value(widget: QtWidgets.QSpinBox | QtWidgets.QProgressBar) -> int: ... # type: ignore +def get_widget_value(widget: QtWidgets.QSpinBox | QtWidgets.QProgressBar) -> int: # type: ignore + ... @overload -def get_widget_value(widget: MultiCheckableComboBox) -> list[str]: ... # type: ignore +def get_widget_value(widget: MultiCheckableComboBox) -> list[str]: # type: ignore + ... def get_widget_value(widget: QtWidgets.QWidget) -> Any: @@ -2500,7 +2516,8 @@ def _load_json_file(self, selected_file): from REvoDesign import __version__ # load back all asked values from a json file - contents_to_load: dict[str, dict[str, Any]] = json.load(open(selected_file)) + with open(selected_file) as selected_file_handle: + contents_to_load: dict[str, dict[str, Any]] = json.load(selected_file_handle) if contents_to_load["metadata"]["__window__"] != self.windowTitle(): logging.error( f"The recipe is made for Dialog `{contents_to_load['metadata']['__window__']}`, " @@ -2667,7 +2684,7 @@ def _remove_row(self, row_layout): Args: row_layout (QHBoxLayout): The row layout to be removed. """ - for i, (layout, key_edit, val_edit) in enumerate(self.row_widgets): + for i, (layout, _key_edit, _val_edit) in enumerate(self.row_widgets): if layout == row_layout: # Remove all widgets in the row for j in reversed(range(layout.count())): @@ -2730,11 +2747,12 @@ def ask_for_multiple_values_as_json() -> str: json_fp = os.path.join("json_multi_input", f"{data_id}.json") os.makedirs(os.path.dirname(json_fp), exist_ok=True) - json.dump( - obj=data.asdict, - fp=open(json_fp, "w"), - indent=4, - ) + with open(json_fp, "w") as json_file: + json.dump( + obj=data.asdict, + fp=json_file, + indent=4, + ) return json_fp diff --git a/src/REvoDesign/tools/download_registry.py b/src/REvoDesign/tools/download_registry.py index 0dcb2faad..938375834 100644 --- a/src/REvoDesign/tools/download_registry.py +++ b/src/REvoDesign/tools/download_registry.py @@ -85,14 +85,18 @@ class FileDownloadRegistry(CitableModuleAbstract): """ A file download registry manager for handling remote file resources. - This class implements automatic file downloading, verification, and cache management based on the pooch library. - It supports specifying file hash values through a registry and provides convenient methods to fetch and verify remote files. + This class implements automatic file downloading, verification, and cache + management based on the pooch library. It supports specifying file hash + values through a registry and provides convenient methods to fetch and + verify remote files. :param name: Module name, used to construct the default data directory path. :param base_url: Base URL for remote files. - :param registry: File registry where keys are filenames and values are corresponding hash values (optional). + :param registry: File registry where keys are filenames and values are + corresponding hash values (optional). :param version: Optional version number for creating versioned data directories. - :param customized_directory: Optional custom download directory path. If not provided, uses the default user data directory. + :param customized_directory: Optional custom download directory path. If not + provided, uses the default user data directory. :param alternative_base_urls: Optional list of alternative base URLs for downloading files. :param retry_count: Number of retries for downloading files. @@ -160,14 +164,25 @@ def _complete_varify_string(a_string: str | None = None, hash_type: str = "md5") return f"{hash_type}:{a_string}" @staticmethod - def preprocess_registry(registry: dict[str, str | None]) -> dict[str, str | None]: + def preprocess_registry(registry: dict[str, str | None]) -> dict[str, str]: """ Preprocess registry to ensure all hash values conform to pooch requirements. :param registry: Original registry. :return: Processed registry. """ - return {k: FileDownloadRegistry._complete_varify_string(v) for k, v in registry.items()} + processed: dict[str, str] = {} + missing_hashes: list[str] = [] + for filename, hash_value in registry.items(): + completed_hash = FileDownloadRegistry._complete_varify_string(hash_value) + if completed_hash is None: + missing_hashes.append(filename) + continue + processed[filename] = completed_hash + if missing_hashes: + missing = ", ".join(sorted(missing_hashes)) + raise ValueError(f"Download registry entries require hashes: {missing}") + return processed @get_cited def setup(self, item: str) -> DownloadedFile: @@ -241,7 +256,9 @@ def prepare_registry_from_md5(md5_contents: str) -> dict[str, str | None]: "Pooch": """ @article{uieda2020, title = {{Pooch}: {A} friend to fetch your data files}, - author = {Leonardo Uieda and Santiago Soler and R{\'{e}}mi Rampin and Hugo van Kemenade and Matthew Turk and Daniel Shapero and Anderson Banihirwe and John Leeman}, + author = {Leonardo Uieda and Santiago Soler and R{\'{e}}mi Rampin and Hugo + van Kemenade and Matthew Turk and Daniel Shapero and Anderson Banihirwe and + John Leeman}, year = {2020}, doi = {10.21105/joss.01943}, url = {https://doi.org/10.21105/joss.01943}, diff --git a/src/REvoDesign/tools/measure_utils.py b/src/REvoDesign/tools/measure_utils.py index c1875c52d..4f16dca8f 100644 --- a/src/REvoDesign/tools/measure_utils.py +++ b/src/REvoDesign/tools/measure_utils.py @@ -313,14 +313,14 @@ def _get_model_atoms(cmd_module, obj: str, state: int = -1): def _atom_coord(a) -> tuple[float, float, float] | None: """Extract ``(x, y, z)`` from a PyMOL atom object.""" if hasattr(a, "coord"): - c = getattr(a, "coord") + c = a.coord if isinstance(c, (list, tuple)) and len(c) >= 3: return (float(c[0]), float(c[1]), float(c[2])) if hasattr(a, "x") and hasattr(a, "y") and hasattr(a, "z"): try: - return (float(getattr(a, "x")), float(getattr(a, "y")), float(getattr(a, "z"))) - except Exception: - pass + return (float(a.x), float(a.y), float(a.z)) + except Exception as exc: + logging.debug("Could not parse atom coordinates from x/y/z attributes: %s", exc) return None @@ -328,9 +328,9 @@ def _atom_unique_id(a) -> int | None: """Extract ``unique_id`` from a PyMOL atom (only uses the explicit attribute).""" if hasattr(a, "unique_id"): try: - return int(getattr(a, "unique_id")) - except Exception: - pass + return int(a.unique_id) + except Exception as exc: + logging.debug("Could not parse atom unique_id: %s", exc) return None diff --git a/src/REvoDesign/tools/mutant_tools.py b/src/REvoDesign/tools/mutant_tools.py index 305f568d4..953b4e652 100644 --- a/src/REvoDesign/tools/mutant_tools.py +++ b/src/REvoDesign/tools/mutant_tools.py @@ -82,7 +82,8 @@ def extract_mutants_from_mutant_id( mutant_string (str): Underscore-seperated mutant string that contains the mutations and score (if possible). _...._ sequences (dict): Wild-type chain: sequence of design molecule - wt_before_chain (bool): Some people write WT residue before chain_id. This helps to correct recognizing of this pattern + wt_before_chain (bool): Some people write WT residue before chain_id. This + helps to correct recognizing of this pattern Returns: tuple: @@ -135,7 +136,7 @@ def extract_mutants_from_mutant_id( _mut_res = _mut.group(2) else: - warnings.warn(issues.BadDataWarning(f"Error while processing mutant id {mut}. ")) + warnings.warn(issues.BadDataWarning(f"Error while processing mutant id {mut}. "), stacklevel=2) continue mutations.append( @@ -222,7 +223,7 @@ def extract_mutant_from_sequences( return if "X" in wt_sequence and not fix_missing: - warnings.warn(issues.ResidueMissingWarning('WT has missing residue masked as "X"!')) + warnings.warn(issues.ResidueMissingWarning('WT has missing residue masked as "X"!'), stacklevel=2) if fix_missing: _mutant_sequence = "" @@ -391,7 +392,8 @@ def extract_mutant_from_pymol_object(pymol_object: str, sequences: RosettaPyProt warnings.warn( issues.BadDataWarning( f"{at.resn} at {at.resi} (chain {chain.chain_id}) is out of range of sequence length." - ) + ), + stacklevel=2, ) continue @@ -485,7 +487,8 @@ def process_mutations(data): def read_profile_design_mutations(filename): - data = json.load(open(filename)) + with open(filename) as mutation_file: + data = json.load(mutation_file) return process_mutations(data) @@ -552,7 +555,7 @@ def quick_mutagenesis(mutant_tree: MutantTree) -> None: nproc = bus.get_value("ui.header_panel.nproc") if mutant_tree.empty: - warnings.warn(issues.NoResultsWarning("Mutant tree is empty!")) + warnings.warn(issues.NoResultsWarning("Mutant tree is empty!"), stacklevel=2) return score_list = mutant_tree.all_mutant_scores @@ -614,7 +617,7 @@ def save_mutant_choices(output_mut_txt_fn: str, mutant_tree: MutantTree): raise issues.NoInputError("No Mutant tree is given!") if mutant_tree.empty: - warnings.warn(issues.NoResultsWarning("mutant tree is empty. save nothing.")) + warnings.warn(issues.NoResultsWarning("mutant tree is empty. save nothing."), stacklevel=2) return mutants_to_save = mutant_tree.all_mutant_ids @@ -629,7 +632,9 @@ def save_mutant_choices(output_mut_txt_fn: str, mutant_tree: MutantTree): return if os.path.exists(output_mut_txt_fn): - warnings.warn(issues.OverridesWarning(f"Mutant table exists and will be overriden! {output_mut_txt_fn}")) + warnings.warn( + issues.OverridesWarning(f"Mutant table exists and will be overriden! {output_mut_txt_fn}"), stacklevel=2 + ) write_input_mutant_table( output_mut_txt_fn, [mt.raw_mutant_id for mt in mutant_tree.all_mutant_objects], diff --git a/src/REvoDesign/tools/package_manager.py b/src/REvoDesign/tools/package_manager.py index bacf7d150..f4dc83c81 100644 --- a/src/REvoDesign/tools/package_manager.py +++ b/src/REvoDesign/tools/package_manager.py @@ -9,6 +9,7 @@ from __future__ import annotations import difflib +import hmac import importlib import importlib.util import io @@ -21,8 +22,11 @@ import re import shutil import socket -import subprocess + +# Used for package-management commands with argv lists and registered process lifecycle handling. +import subprocess # nosec B404 import sys +import tempfile import threading import time import urllib.request @@ -33,6 +37,7 @@ from dataclasses import dataclass, field from datetime import datetime, timezone from functools import cached_property, partial +from pathlib import Path from typing import Any, ClassVar, NoReturn, TypedDict, TypeVar, cast, overload from urllib.error import HTTPError, URLError from urllib.parse import urlparse @@ -201,6 +206,52 @@ def critical(msg: str, *args, **kwargs): THIS_FILE_URL = f"{GIST_BASE_URL}/REvoDesign_PyMOL.py" # Define the URL of the JSON file RICH_TABLE_JSON = f"{GIST_BASE_URL}/REvoDesignExtrasTableRich.json" +MANIFEST_URL = f"{GIST_BASE_URL}/manifest.json" +PACKAGE_MANAGER_BOOTSTRAP_ENV = "REVODESIGN_PM_BOOTSTRAP_DIR" +BOOTSTRAP_RETRY_DELAYS_SECONDS = (0.5, 2.0, 5.0) +MANIFEST_ASSET_MANAGER_UI = "REvoDesign-PyMOL-entry.ui" +MANIFEST_ASSET_EXTRAS_JSON = "REvoDesignExtrasTableRich.json" + +# ponytail: HMAC key prevents a forked manifest from accidentally validating +# against the wrong installation. Not a secret — it ships in the public Gist. +# Value: catches truncation/corruption + accidental cross-project mismatch. +_MANAGER_HMAC_KEY = bytes.fromhex("bd8c4274c76d312b22a0c3d58aad68860b645f1dcc22e67d3d4fe16bf59c6343") + + +def package_manager_bootstrap_dir() -> Path: + bootstrap_dir = os.environ.get(PACKAGE_MANAGER_BOOTSTRAP_ENV) + if bootstrap_dir: + return Path(bootstrap_dir).expanduser() + return Path(tempfile.gettempdir()) / "REvoDesign-manager" + + +def bootstrap_manager_ui_file() -> Path: + return package_manager_bootstrap_dir() / "UI" / "REvoDesign_installer.ui" + + +def bootstrap_extras_json() -> Path: + return package_manager_bootstrap_dir() / "REvoDesignExtrasTableRich.json" + + +def _with_bootstrap_retries(operation: Callable[[], R], description: str) -> R: + last_error: Exception | None = None + for attempt, delay in enumerate((0.0, *BOOTSTRAP_RETRY_DELAYS_SECONDS), start=1): + if delay: + time.sleep(delay) + try: + return operation() + except Exception as e: + last_error = e + logging.warning( + "%s failed on attempt %s/%s: %s", + description, + attempt, + len(BOOTSTRAP_RETRY_DELAYS_SECONDS) + 1, + e, + ) + if last_error is None: + raise RuntimeError(f"{description} failed without an exception") + raise last_error # Define the proxy protocols allowed @@ -235,15 +286,22 @@ def _safe_parse_version(value: str) -> Version | None: return None -def _read_https_url(url: str, timeout: float = 10.0) -> bytes: +def _read_https_url(url: str, timeout: float = 10.0, headers: Mapping[str, str] | None = None) -> bytes: parsed = urlparse(url) if parsed.scheme != "https" or not parsed.netloc: raise ValueError(f"Only https URLs are allowed: {url}") - request = urllib.request.Request(url, headers={"User-Agent": "REvoDesign-PackageManager"}) + request_headers = {"User-Agent": "REvoDesign-PackageManager"} + if headers: + request_headers.update(headers) + request = urllib.request.Request(url, headers=request_headers) with urllib.request.urlopen(request, timeout=timeout) as response: # nosec B310 return response.read() +GITHUB_TAG_CACHE_TTL_SECONDS = 15 * 60 +_GITHUB_TAG_CACHE: dict[str, tuple[float, list[str]]] = {} + + def _python_version_matches(spec: str | None, current_version: str | None) -> bool: """ Check if the current Python version matches the specified version specification. @@ -457,7 +515,7 @@ def find_extras(self, name: str) -> list[ExtrasItem]: return [item for item in self.all_extras if item.name == name] -def fetch_gist_file(ui_file_url: str, save_to_file: str) -> None: +def fetch_gist_file(ui_file_url: str, save_to_file: str, *, timeout: float = 10.0) -> None: """ Fetch the UI file from the given URL, save it to a temporary file, and yield its absolute path. @@ -474,7 +532,7 @@ def fetch_gist_file(ui_file_url: str, save_to_file: str) -> None: try: # Fetch the file content and write it to the temporary file - ui_data = _read_https_url(ui_file_url).decode("utf-8") + ui_data = _read_https_url(ui_file_url, timeout=timeout).decode("utf-8") with open(save_to_file, "w") as ui_handle: ui_handle.write(ui_data) @@ -512,6 +570,92 @@ def fetch_gist_json(url: str) -> dict[str, Any]: return {} +def fetch_required_gist_json(url: str) -> dict[str, Any]: + json_data = fetch_gist_json(url) + if not json_data: + raise URLError(f"Failed to fetch required JSON bootstrap asset: {url}") + return json_data + + +def fetch_bootstrap_manifest() -> dict[str, str]: + manifest = _with_bootstrap_retries( + lambda: fetch_required_gist_json(MANIFEST_URL), + "REvoDesign bootstrap manifest", + ) + if not all(isinstance(key, str) and isinstance(value, str) for key, value in manifest.items()): + raise ValueError("Bootstrap manifest must map asset names to HMAC strings") + return cast(dict[str, str], manifest) + + +def fetch_verified_bootstrap_file( + *, + url: str, + asset_name: str, + save_to_file: str, + manifest: dict[str, str], + description: str, +) -> None: + def _fetch_and_verify() -> None: + fetch_gist_file(ui_file_url=url, save_to_file=save_to_file) + if not verify_manifest({asset_name: save_to_file}, manifest): + raise ValueError(f"{description} failed HMAC verification") + + _with_bootstrap_retries(_fetch_and_verify, description) + + +def load_bootstrap_extras_json(path: Path | None = None) -> dict[str, Any]: + path = path or bootstrap_extras_json() + try: + with path.open(encoding="utf-8") as json_handle: + json_data = json.load(json_handle) + except OSError as e: + logging.error(f"Error loading packaged extras table from {path}: {e}") + return {} + except json.JSONDecodeError as e: + logging.error(f"Error decoding packaged extras table from {path}: {e}") + return {} + + if not isinstance(json_data, dict): + logging.error(f"Error loading packaged extras table from {path}: expected a dictionary.") + return {} + return json_data + + +load_packaged_extras_json = load_bootstrap_extras_json + + +def write_bootstrap_extras_json(remote_data: dict[str, Any]) -> None: + extras_json = bootstrap_extras_json() + cache_dir = str(extras_json.parent) + os.makedirs(cache_dir, exist_ok=True) + with open(str(extras_json), "w", encoding="utf-8") as fh: + json.dump(remote_data, fh) + + +def _compute_hmac(filepath: str) -> str: + """Compute HMAC-SHA256 of a file using the embedded manager key.""" + with open(filepath, "rb") as fh: + return hmac.new(_MANAGER_HMAC_KEY, fh.read(), "sha256").hexdigest() + + +def verify_manifest(assets: dict[str, str], manifest: dict[str, str]) -> bool: + """Verify each asset in `assets` matches its HMAC entry in `manifest`. + + ``assets`` maps filename → temp-file path on disk. + ``manifest`` maps filename → expected HMAC hex digest. + """ + for filename, filepath in assets.items(): + expected = manifest.get(filename) + if expected is None: + logging.error("Manifest missing entry for %s", filename) + return False + actual = _compute_hmac(filepath) + if not hmac.compare_digest(actual, expected): + logging.error("HMAC mismatch for %s", filename) + return False + return True + + R = TypeVar("R") @@ -573,7 +717,7 @@ def stream_reader(pipe: io.IOBase, collector: list[str], label: str): collector.append(line) pipe.close() - process = subprocess.Popen( + process = subprocess.Popen( # nosec B603 cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, @@ -1032,14 +1176,17 @@ class MenuItem: """ A data class representing a menu item. - This class is used to define the properties of a menu item, including its name, associated function, and optional arguments. - The use of the @dataclass decorator automatically generates special methods such as __init__(), __repr__(), and __eq__(). + This class is used to define the properties of a menu item, including its + name, associated function, and optional arguments. + The use of the @dataclass decorator automatically generates special methods + such as __init__(), __repr__(), and __eq__(). The frozen parameter ensures that instances of the class are immutable, enhancing thread safety and consistency. Attributes: name (str): The name of the menu item, used for display and identification. func (Callable): The function associated with the menu item, which is executed when the item is selected. - kwargs (Optional[Mapping]): Optional arguments passed to the associated function when it is executed. Defaults to None. + kwargs (Optional[Mapping]): Optional arguments passed to the associated + function when it is executed. Defaults to None. """ name: str @@ -1071,28 +1218,38 @@ class REvoDesignPackageManager: remote_extra_group_data: ExtrasGroups = None # type: ignore def ensure_ui_file(self, upgrade: bool = False): - ui_file = os.path.abspath( - os.path.join(os.path.dirname(__file__), "REvoDesign-manager", "UI", "REvoDesign_installer.ui") - ) - os.makedirs(os.path.dirname(ui_file), exist_ok=True) + """Return the manager UI file path, fetching from Gist if needed. - # if not exists, preform the first fetch - if not os.path.isfile(ui_file): - logging.info(f"Fetching UI: {UI_FILE_URL} -> {ui_file}") - fetch_gist_file(ui_file_url=UI_FILE_URL, save_to_file=ui_file) - logging.info(f"Fetched UI file for manager: {ui_file}") - return ui_file + Download the manager UI into the writable bootstrap directory. + ``upgrade=True`` still uses the same retrying fetch path. + """ + ui_file_path = bootstrap_manager_ui_file() + cache_dir = str(ui_file_path.parent) + os.makedirs(cache_dir, exist_ok=True) + ui_file = str(ui_file_path) - # otherwise, if the user not requires an upgrade, return - if not upgrade: - logging.debug(f"pre-downloaded UI file found: {ui_file}") - return ui_file + new_ui_handle = tempfile.NamedTemporaryFile(delete=False, suffix=".ui") + new_ui_file = new_ui_handle.name + new_ui_handle.close() - # otherwise, preform the upgrade - new_ui_file = f"{ui_file}.swp" + try: + manifest = fetch_bootstrap_manifest() + fetch_verified_bootstrap_file( + url=UI_FILE_URL, + asset_name=MANIFEST_ASSET_MANAGER_UI, + save_to_file=new_ui_file, + manifest=manifest, + description="REvoDesign manager UI bootstrap", + ) + if upgrade and os.path.isfile(ui_file): + self.upgrade_check(original_file=ui_file, new_file=new_ui_file, title="REvoDesign Manager UI file") + else: + shutil.move(new_ui_file, ui_file) + except Exception: + if os.path.isfile(new_ui_file): + os.remove(new_ui_file) + raise - fetch_gist_file(ui_file_url=UI_FILE_URL, save_to_file=new_ui_file) - self.upgrade_check(original_file=ui_file, new_file=new_ui_file, title="REvoDesign Manager UI file") return ui_file @staticmethod @@ -1100,8 +1257,10 @@ def upgrade_check(original_file: str, new_file: str, title: str): """ Check and apply an upgrade if necessary. - This function compares the original file with a new fetched file, generates a diff file if there are differences, - and prompts the user to confirm whether to apply the upgrade. If the user confirms, the new file replaces the original file. + This function compares the original file with a new fetched file, + generates a diff file if there are differences, and prompts the user to + confirm whether to apply the upgrade. If the user confirms, the new file + replaces the original file. Parameters: - original_file (str): The path to the original file. @@ -1137,9 +1296,12 @@ def upgrade_check(original_file: str, new_file: str, title: str): ':::::Upgrade Summary:::::

' "" "" - f'' - f'' - f'' + '' + f'' + '' + f'' + '' + f'' "
Event-Affected Lines
Added :{num_added_lines}
Changed:{num_chged_lines}
Deleted:{num_deled_lines}
Added :{num_added_lines}
Changed:{num_chged_lines}
Deleted:{num_deled_lines}
" "You must check out these changes carefully.

" f"See all changes in this diff file of {title}.", @@ -1169,13 +1331,38 @@ def self_upgrade(self): if not confirmed: return notify_box("Upgrade cancelled.") - new_py_file = f"{__file__}.swp" + # Download manifest and both assets to a temp dir, verify HMAC, then apply + manifest = fetch_gist_json(MANIFEST_URL) + if not manifest: + notify_box("Cannot verify upgrade: manifest unavailable. " "The upgrade source may be compromised.") + return + + temp_dir = tempfile.mkdtemp() + try: + py_temp = os.path.join(temp_dir, "REvoDesign_PyMOL.py") + ui_temp = os.path.join(temp_dir, "REvoDesign-PyMOL-entry.ui") - fetch_gist_file(THIS_FILE_URL, new_py_file) + fetch_gist_file(THIS_FILE_URL, py_temp) + fetch_gist_file(UI_FILE_URL, ui_temp) - self.upgrade_check(original_file=__file__, new_file=new_py_file, title="REvoDesign Manager") + assets = { + "REvoDesign_PyMOL.py": py_temp, + "REvoDesign-PyMOL-entry.ui": ui_temp, + } + if not verify_manifest(assets, manifest): + notify_box("Upgrade verification failed. " "The downloaded files may be corrupted or tampered with.") + return - self.ensure_ui_file(upgrade=True) + self.upgrade_check(original_file=__file__, new_file=py_temp, title="REvoDesign Manager") + + bootstrap_ui = str(bootstrap_manager_ui_file()) + if os.path.isfile(bootstrap_ui): + self.upgrade_check(original_file=bootstrap_ui, new_file=ui_temp, title="REvoDesign Manager UI file") + else: + os.makedirs(os.path.dirname(bootstrap_ui), exist_ok=True) + shutil.move(ui_temp, bootstrap_ui) + finally: + shutil.rmtree(temp_dir, ignore_errors=True) @contextmanager def freeze_manager(self): @@ -1221,7 +1408,7 @@ def initialize_manager(self): logging.debug("Run pre-fetching tasks... ") - self.refresh_remote_json() + self.load_bootstrap_json() self.pip_installer = run_worker_thread_in_pool(PIPInstaller) @@ -1282,14 +1469,9 @@ def proxy_in_env(self, proxy: str | None = None, mirror: str | None = None) -> d } return proxy_env - def refresh_remote_json(self): - """ - Refreshes the list of available extras by fetching data from a JSON source. - - This method uses a worker thread to fetch extras data with a progress bar indication. - If fetching fails, it shows an error notification and sets up an empty extras list. - """ - d_placeholder = { + @staticmethod + def _extras_placeholder() -> dict[str, Any]: + return { "entities": [ { "name": "No Extras is Fetched", @@ -1298,27 +1480,83 @@ def refresh_remote_json(self): } ], } - remote_data = run_worker_thread_in_pool(worker_function=fetch_gist_json, url=RICH_TABLE_JSON) + def _load_extras_table(self, remote_data: dict[str, Any], *, notify_on_error: bool) -> None: self.notification_channel(remote_data) if not remote_data: - notify_box( - "Error fetching or validating the JSON data. \n" - "Please reconfigure your network and press to try again " - "if you wish to continue installation with extra packages" - ) - - if "entities" not in remote_data: - notify_box("Fetched data is not valid. The data is expected to have an `entities` key.") + if notify_on_error: + notify_box( + "Error fetching or validating the JSON data. \n" + "Please reconfigure your network and press to try again " + "if you wish to continue installation with extra packages" + ) + remote_data = self._extras_placeholder() + elif "entities" not in remote_data: + if notify_on_error: + notify_box("Fetched data is not valid. The data is expected to have an `entities` key.") + remote_data = self._extras_placeholder() - self.remote_extra_group_data = ExtrasGroups.from_dict(remote_data or d_placeholder) + self.remote_extra_group_data = ExtrasGroups.from_dict(remote_data) # Create and position the extra components checkbox list self.extra_checkbox = CheckableListView( self.installer_ui.listView_extras, self.remote_extra_group_data, filter=PLATFORM_INFO ) + def load_bootstrap_json(self): + """ + Fetches the extras table into the writable bootstrap directory. + """ + + def fetch_verified_extras() -> dict[str, Any]: + manifest = fetch_bootstrap_manifest() + extras_handle = tempfile.NamedTemporaryFile(delete=False, suffix=".json") + extras_file = extras_handle.name + extras_handle.close() + try: + fetch_verified_bootstrap_file( + url=RICH_TABLE_JSON, + asset_name=MANIFEST_ASSET_EXTRAS_JSON, + save_to_file=extras_file, + manifest=manifest, + description="REvoDesign extras table bootstrap", + ) + json_data = load_bootstrap_extras_json(Path(extras_file)) + if not json_data: + raise ValueError("Fetched extras bootstrap JSON is empty or invalid") + return json_data + finally: + if os.path.isfile(extras_file): + os.remove(extras_file) + + remote_data = run_worker_thread_in_pool(worker_function=fetch_verified_extras) + if remote_data and "entities" in remote_data: + try: + write_bootstrap_extras_json(remote_data) + except OSError: + logging.warning("Failed to write extras JSON to %s", bootstrap_extras_json()) + self._load_extras_table(remote_data, notify_on_error=False) + + def load_packaged_json(self): + return self.load_bootstrap_json() + + def refresh_remote_json(self): + """ + Refreshes the list of available extras by fetching data from a JSON source. + + This method uses a worker thread to fetch extras data with a progress bar indication. + If fetching fails, it shows an error notification and sets up an empty extras list. + """ + remote_data = run_worker_thread_in_pool(worker_function=fetch_gist_json, url=RICH_TABLE_JSON) + self._load_extras_table(remote_data, notify_on_error=True) + # Save explicit refresh output as the latest bootstrap copy for this session. + if remote_data and "entities" in remote_data: + try: + write_bootstrap_extras_json(remote_data) + except OSError: + logging.warning("Failed to write extras JSON to %s", bootstrap_extras_json()) + @staticmethod def notification_channel(d: dict): """ @@ -1357,7 +1595,8 @@ def collect_diagnostic_data(collect_dummy: bool = False, drop_sensitives=True): if not drop_sensitives: confirmed = decide( title="Agree to collect SENSITIVE data?", - description="[!!!CAUSION!!!]Do you REALLY want to collect diagnostic information INCLUDING ALL SENSITIVE data?\n" + description="[!!!CAUSION!!!]Do you REALLY want to collect diagnostic " + "information INCLUDING ALL SENSITIVE data?\n" "Please DO NOT share this information with anyone else or post it to public channels.", ) if not confirmed: @@ -1443,13 +1682,17 @@ def make_window(self) -> QtWidgets.QDialog: except Exception as e: decided = decide( "UI Error", - "Error Occurs while loading UI file, this UI may out-of-dated.\nCleanup and fetch the latest?", + "Error Occurs while loading UI file, this UI may out-of-dated.\n" + "Re-fetch the latest version from the network?", details=str(e), ) if decided: - os.remove(ui_file) - return self.make_window() - raise RuntimeError(f"Error occurs while loading UI file: {e}.") + # Re-fetch into a temp file first; replace the bootstrap file + # only after the new download succeeds. + ui_file = self.ensure_ui_file(upgrade=True) + self.installer_ui = loadUi(ui_file, dialog) + else: + raise RuntimeError(f"Error occurs while loading UI file: {e}.") from e # add right-click menu on `self.installer_ui.label_header`, # add a item `Upgrade UI` and connect `partial(self.ensure_ui_file, upgrade=True)` @@ -1568,13 +1811,15 @@ def fetch_tags(self): specified GitHub repository, and then sets the result as the value of the `comboBox_version` combo box in the UI. """ - # Run a worker thread to fetch tags with a progress bar - tags = run_worker_thread_in_pool(worker_function=get_github_repo_tags, repo_url=REPO_URL) + # ponytail: non-fatal — offline startup must not show an error popup + try: + tags = run_worker_thread_in_pool(worker_function=get_github_repo_tags, repo_url=REPO_URL) + except Exception: + logging.warning("Failed to fetch GitHub release tags (offline or API error)") + return if tags and isinstance(tags, list): return set_widget_value(self.installer_ui.comboBox_version, tags) - return notify_box(f"Failed to fetch version tags from GitHub repo: \n{REPO_URL}") - # a copy from `REvoDesign/tools/customized_widgets.py` @staticmethod @@ -2186,15 +2431,16 @@ def _terminate_process(process: subprocess.Popen[str]) -> None: if process.poll() is None: try: process.terminate() - except Exception: - pass + except Exception as exc: + logging.debug("Could not terminate process %s: %s", process, exc) try: process.wait(timeout=5) - except Exception: + except Exception as exc: + logging.debug("Process %s did not exit after terminate; killing it: %s", process, exc) try: process.kill() - except Exception: - pass + except Exception as kill_exc: + logging.debug("Could not kill process %s: %s", process, kill_exc) class AbortButtonOverlay(QtCore.QObject): @@ -2340,11 +2586,8 @@ def _connect_signals(self) -> None: def _handle_abort_request(self) -> None: self._aborted = True - try: - cmd.interrupt() - except Exception: - pass self.worker_thread.interrupt() + RunningProcessRegistry.terminate(self.worker_thread) if self.worker_thread.isRunning(): self.worker_thread.quit() if not self.worker_thread.wait(100): @@ -2410,6 +2653,7 @@ class WorkerThread(QtCore.QThread): """Custom worker thread for executing a function in a separate thread.""" result_signal = QtCore.pyqtSignal(list) + error_signal = QtCore.pyqtSignal(object) finished_signal = QtCore.pyqtSignal() interrupt_signal = QtCore.pyqtSignal() notify_signal = QtCore.pyqtSignal(str) @@ -2422,6 +2666,7 @@ def __init__(self, func, args=None, kwargs=None): self.args = args or () self.kwargs = kwargs or {} self.results = None + self.error: BaseException | None = None self.interrupt_signal.connect(self._handle_interrupt) def run(self): @@ -2433,6 +2678,14 @@ def run(self): self.results = [self.func(*self.args, **self.kwargs)] if self.results: self.result_signal.emit(self.results) + except BaseException as exc: # noqa: B036 - worker errors are re-raised by the caller after thread join + self.error = exc + logging.exception("WorkerThread failed while running %s", self.func) + try: + self.error_signal.emit(exc) + self.notify_signal.emit(f"{type(exc).__name__}: {exc}") + except RuntimeError: + pass finally: try: self.finished_signal.emit() @@ -2450,6 +2703,9 @@ def run(self): def handle_result(self): return self.results + def handle_error(self): + return self.error + def interrupt(self): self.requestInterruption() RunningProcessRegistry.terminate(self) @@ -2459,7 +2715,7 @@ def _handle_interrupt(self) -> None: logging.debug("WorkerThread interrupt signal handled.") -def get_github_repo_tags(repo_url) -> list[str]: +def get_github_repo_tags(repo_url: str, *, timeout: float = 5.0) -> list[str]: """ Retrieve all released tags of a GitHub repository using urllib. @@ -2474,43 +2730,56 @@ def get_github_repo_tags(repo_url) -> list[str]: list: A list of tag names for the repository. """ - # Extract the owner and repo name from the URL - parts = repo_url.split("/") - owner = parts[-2] - repo = parts[-1] + parsed = urlparse(repo_url) + parts = [part for part in parsed.path.strip("/").split("/") if part] + if parsed.scheme != "https" or parsed.netloc.lower() != "github.com" or len(parts) < 2: + logging.warning("Unsupported GitHub repository URL: %s", repo_url) + return [] + owner, repo = parts[:2] + repo = repo.removesuffix(".git") # GitHub API URL for listing tags api_url = f"https://api.github.com/repos/{owner}/{repo}/tags" + now = time.monotonic() + cached = _GITHUB_TAG_CACHE.get(api_url) + if cached and now - cached[0] <= GITHUB_TAG_CACHE_TTL_SECONDS: + return list(cached[1]) + + headers = {"Accept": "application/vnd.github+json"} + github_token = os.environ.get("GITHUB_TOKEN", "").strip() + if github_token: + headers["Authorization"] = f"Bearer {github_token}" try: # Send a GET request to the GitHub API - response_data = _read_https_url(api_url).decode() + response_data = _read_https_url(api_url, timeout=timeout, headers=headers).decode() # Parse JSON response data tags = json.loads(response_data) if not isinstance(tags, list): logging.error("Failed to parse GitHub tags response: expected a list.") - return [] + return list(cached[1]) if cached else [] # Extract the name of each tag try: tag_names = [tag["name"] for tag in tags] except (KeyError, TypeError) as e: logging.error(f"Failed to extract tag names from GitHub response: {e}") - return [] + return list(cached[1]) if cached else [] + _GITHUB_TAG_CACHE[api_url] = (time.monotonic(), tag_names) return tag_names except HTTPError as e: # Handle HTTP errors (e.g., repository not found, rate limit exceeded) logging.warning(f"GitHub API returned status code {e.code}") - return [] + return list(cached[1]) if cached else [] except URLError as e: # Handle URL errors (e.g., network issues) logging.error(f"Failed to reach the server. Reason: {e.reason}") - return [] + return list(cached[1]) if cached else [] except json.JSONDecodeError as e: logging.error(f"Failed to decode GitHub tags response as JSON: {e}") - return [] + return list(cached[1]) if cached else [] except TypeError as e: logging.error(f"Failed to parse GitHub tags response: {e}") - return [] + return list(cached[1]) if cached else [] # a minimum copy from `REvoDesign/tools/customized_widgets.py` @@ -2648,7 +2917,7 @@ def execute_on_main_thread(func: Callable[..., GuiResult], *args, **kwargs) -> G def callback(): try: result["value"] = func(*args, **kwargs) - except BaseException as exc: # pragma: no cover - defensive against unexpected UI failures + except BaseException as exc: # noqa: B036 # pragma: no cover - re-raised across the GUI-thread handoff error["error"] = exc finally: semaphore.release() @@ -2663,15 +2932,19 @@ def callback(): # Overload #1: None or Warning => returns bool +# fmt: off @overload -def notify_box(message: str = "", error_type: None | type[Warning] = None, details: str | None = None) -> None: ... +def notify_box(message: str = "", error_type: None | type[Warning] = None, details: str | None = None) -> None: + ... # Overload #2: Exception => NoReturn @overload -def notify_box(message: str, error_type: type[Exception], details: str | None = None) -> NoReturn: ... +def notify_box(message: str, error_type: type[Exception], details: str | None = None) -> NoReturn: + ... +# fmt: on def notify_box( @@ -2718,7 +2991,7 @@ def _show_notification_dialog( # error_type is a Warning => also bool if issubclass(error_type, Warning): - warnings.warn(error_type(message)) + warnings.warn(error_type(message), stacklevel=2) return # Otherwise, raise => NoReturn @@ -2736,6 +3009,7 @@ def raise_error(error_type: type[Exception], message: str) -> NoReturn: raise error_type(message) +# fmt: off @overload def run_worker_thread_in_pool( worker_function: Callable[..., R], @@ -2743,7 +3017,8 @@ def run_worker_thread_in_pool( trigger_buttons: QtWidgets.QPushButton | Iterable[QtWidgets.QPushButton] | None = None, notify_slot: Callable[[str], None] | None = None, **kwargs, -) -> R: ... +) -> R: + ... @overload @@ -2753,7 +3028,9 @@ def run_worker_thread_in_pool( trigger_buttons: QtWidgets.QPushButton | Iterable[QtWidgets.QPushButton] | None = None, notify_slot: Callable[[str], None] | None = None, **kwargs, -) -> R | None: ... +) -> R | None: + ... +# fmt: on def run_worker_thread_in_pool( @@ -2782,6 +3059,9 @@ def run_worker_thread_in_pool( time.sleep(0.01) result = work_thread.handle_result() + error = work_thread.handle_error() + if error is not None: + raise error return result[0] if result else None diff --git a/src/REvoDesign/tools/pymol_utils.py b/src/REvoDesign/tools/pymol_utils.py index 8c5bfc8c2..7216707e2 100644 --- a/src/REvoDesign/tools/pymol_utils.py +++ b/src/REvoDesign/tools/pymol_utils.py @@ -74,7 +74,7 @@ def find_small_molecules_in_protein(sele): Returns None if the selection is not provided. """ if not sele: - warnings.warn(issues.NoInputWarning("Selection for small molecules is not provided.")) + warnings.warn(issues.NoInputWarning("Selection for small molecules is not provided."), stacklevel=2) return None # Return None if the selection is not provided # Retrieve the atoms that belong to small molecules within the selection and extract unique small molecule names @@ -88,13 +88,14 @@ def find_small_molecules_in_protein(sele): "Could not find unique small molecules with standalone chain id. \n" 'A possible fix is calling `alter r. RES, chain=""` to fix the problem \n' "then re-load this session." - ) + ), + stacklevel=2, ) # Return a list of unique small molecule names found within the selection return unique_small_molecules - warnings.warn(issues.FallingBackWarning("Falling back to all `hetatm`")) + warnings.warn(issues.FallingBackWarning("Falling back to all `hetatm`"), stacklevel=2) small_molecules = [at.resn for at in cmd.get_model("hetatm and (not polymer.protein)").atom] unique_small_molecules = list(set(small_molecules)) @@ -165,7 +166,8 @@ def is_distal_residue_pair( - chain_id (str): The chain identifier. - resi_1 (int): The residue number of the first amino acid. - resi_2 (int): The residue number of the second amino acid. - - minimal_distance (float, optional): The minimum distance threshold for residues to be considered distal. Default is 20. + - minimal_distance (float, optional): The minimum distance threshold for + residues to be considered distal. Default is 20. - use_sidechain_angle (bool, optional): Whether to consider the orientation of side chains. Default is False. Returns: @@ -347,7 +349,8 @@ def refresh_all_selections(): Function: refresh_all_selections Usage: selections = refresh_all_selections() - This function refreshes and logs information about all PyMOL selections except 'sele' and those starting with '_align'. + This function refreshes and logs information about all PyMOL selections + except 'sele' and those starting with '_align'. Returns: - list: List of all non-'sele' selections (excluding those starting with '_align') @@ -385,7 +388,8 @@ def is_a_REvoDesign_session(): "session's feature, you should always create seperate" "sessions according to your dataset and merge them " "manually in PyMOL window." - ) + ), + stacklevel=2, ) return check @@ -397,12 +401,15 @@ def make_temperal_input_pdb( resn="", selection="", save_as_format="pdb", - wd=os.getcwd(), + wd=None, reload=True, ): """ Function: make_temperal_input_pdb - Usage: input_file = make_temperal_input_pdb(molecule, chain_id=None, segment_id=None, format='pdb', wd=os.getcwd(), reload=True) + Usage: input_file = make_temperal_input_pdb( + molecule, chain_id=None, segment_id=None, format='pdb', + wd=os.getcwd(), reload=True + ) exi This function generates a temporary input PDB file from the specified molecule selection. It supports selection by chain ID, segment ID, or both. @@ -420,6 +427,9 @@ def make_temperal_input_pdb( Returns: - str: Path to the generated temporary input PDB file """ + if wd is None: + wd = os.getcwd() + os.makedirs(wd, exist_ok=True) input_file = os.path.join( wd, diff --git a/src/REvoDesign/tools/rfdiffusion_tools.py b/src/REvoDesign/tools/rfdiffusion_tools.py index cc562b09e..443f2ce6e 100644 --- a/src/REvoDesign/tools/rfdiffusion_tools.py +++ b/src/REvoDesign/tools/rfdiffusion_tools.py @@ -2,18 +2,30 @@ # Distributed under the terms of the GNU General Public License v3.0. # SPDX-License-Identifier: GPL-3.0-only +import os -import matplotlib import numpy as np from REvoDesign.basic.abc_third_party_module import ThirdPartyModuleAbstract from REvoDesign.bootstrap.set_config import is_package_installed +from REvoDesign.logger import ROOT_LOGGER from REvoDesign.tools.utils import require_installed -try: - matplotlib.use("QtAgg") -except Exception: - pass # Fall back to default backend (e.g. in headless / CI environments) +logging = ROOT_LOGGER.getChild(__name__) + +PREFERRED_INTERACTIVE_BACKEND = "QtAgg" + + +def _prefer_interactive_backend_for_plotting() -> None: + if os.environ.get("MPLBACKEND"): + return + + import matplotlib + + try: + matplotlib.use(PREFERRED_INTERACTIVE_BACKEND, force=False) + except Exception as exc: + logging.debug("Falling back to default matplotlib backend: %s", exc) @require_installed @@ -220,6 +232,8 @@ def plot_potential_field(self, grid_size=200, margin=10, save_to: str = "default Path to save the plot. If 'default.png', the plot will be saved in the current working directory. """ + _prefer_interactive_backend_for_plotting() + import matplotlib.pyplot as plt X, Y, P = self.compute_potential_field(grid_size=grid_size, margin=margin) @@ -247,7 +261,7 @@ def plot_potential_field(self, grid_size=200, margin=10, save_to: str = "default default_size = 90 # Overlay ligand atoms - for idx, (coord, element) in enumerate(zip(self.ligand_coords, self.ligand_elements)): + for coord, element in zip(self.ligand_coords, self.ligand_elements): color = color_map.get(element, default_color) size = size_map.get(element, default_size) plt.scatter(coord[0], coord[1], c=color, s=size, edgecolors="black", linewidth=1.5, zorder=3) @@ -275,6 +289,6 @@ def plot_potential_field(self, grid_size=200, margin=10, save_to: str = "default # # Example usage: # visualizer = SubstratePotentialVisualizer( -# pdb_path="/Users/yyy/Documents/protein_design/REvoDesign-test-data/1SUO.pdb", lig_key="HEM" +# pdb_path="/path/to/1SUO.pdb", lig_key="HEM" # ) # visualizer.plot_potential_field(save_to='HEM.png') diff --git a/src/REvoDesign/tools/rosetta_utils.py b/src/REvoDesign/tools/rosetta_utils.py index 1d1cc5016..d50e44315 100644 --- a/src/REvoDesign/tools/rosetta_utils.py +++ b/src/REvoDesign/tools/rosetta_utils.py @@ -131,11 +131,11 @@ def extra_res_to_opts(ligands_params: list[str] | str) -> list[str]: ligands = [] for _, l in enumerate(ligands_params): if not (isinstance(l, str) and l.endswith(".params")): - warnings.warn(issues.BadDataWarning(f"Invalid Parameter input for ligand - {l}")) + warnings.warn(issues.BadDataWarning(f"Invalid Parameter input for ligand - {l}"), stacklevel=2) continue if not os.path.isfile(l): - warnings.warn(issues.BadDataWarning(f"Ignore nofound ligand - {l}")) + warnings.warn(issues.BadDataWarning(f"Ignore nofound ligand - {l}"), stacklevel=2) continue ligands.extend(["-extra_res_fa" if l.endswith("fa.params") else "-extra_res_cen", os.path.abspath(l)]) @@ -191,7 +191,7 @@ def is_wsl_available(): return wsl_bin is not None except RuntimeError: # If an error occurs, it indicates that WSL may not be available - warnings.warn(issues.PlatformNotSupportedWarning("WSL is not available on this machine.")) + warnings.warn(issues.PlatformNotSupportedWarning("WSL is not available on this machine."), stacklevel=2) # Return False, indicating that WSL is not available return False @@ -218,7 +218,8 @@ def is_docker_available() -> bool: warnings.warn( issues.PlatformNotSupportedWarning( f"Docker is not available(uninstalled or unlaunched) on this machine: {e}" - ) + ), + stacklevel=2, ) return False @@ -290,10 +291,48 @@ def read_rosetta_node_config() -> dict[str, Any]: ROSETTA_COMMON_CITATION: dict[str, str | tuple] = { - "Rosetta": """@article{10.1038/s41592-020-0848-2, author = {Leman, J. K. and Weitzner, B. D. and Lewis, S. M. and Adolf‐Bryfogle, J. and Alam, N. and Alford, R. F. and Aprahamian, M. L. and Baker, D. and Barlow, K. A. and Barth, P. and Basanta, B. and Bender, B. J. and Blacklock, K. and Bonet, J. and Boyken, S. E. and Bradley, P. and Bystroff, C. and Conway, P. and Cooper, S. and Correia, B. E. and Coventry, B. and Das, R. and Jong, R. M. d. and DiMaio, F. and Dsilva, L. and Dunbrack, R. L. and Ford, A. S. and Frenz, B. and Fu, D. and Geniesse, C. and Goldschmidt, L. and Gowthaman, R. and Gray, J. J. and Gront, D. and Guffy, S. L. and Horowitz, S. and Huang, P. and Huber, T. and Jacobs, T. M. and Jeliazkov, J. R. and Johnson, D. K. and Kappel, K. and Karanicolas, J. and Khakzad, H. and Khar, K. R. and Khare, S. D. and Khatib, F. and Khramushin, A. and King, C. and Kleffner, R. and Koepnick, B. and Kortemme, T. and Kuenze, G. and Kuhlman, B. and Kuroda, D. and Labonte, J. W. and Lai, J. and Lapidoth, G. and Leaver‐Fay, A. and Lindert, S. and Linsky, T. W. and London, N. and Lubin, J. H. and Lyskov, S. and Maguire, J. B. and Malmström, L. and Marcos, E. and Marcu, O. and Marze, N. and Meiler, J. and Moretti, R. and Mulligan, V. K. and Nerli, S. and Norn, C. and Ó’Conchúir, S. and Ollikainen, N. and Ovchinnikov, S. and Pacella, M. S. and Pan, X. and Park, H. and Pavlovicz, R. E. and Pethe, M. A. and Pierce, B. G. and Pilla, K. B. and Raveh, B. and Renfrew, P. D. and Burman, S. S. R. and Rubenstein, A. B. and Sauer, M. F. and Scheck, A. and Schief, W. R. and Schueler‐Furman, O. and Sedan, Y. and Sevy, A. M. and Sgourakis, N. G. and Shi, L. and Siegel, J. B. and Silva, D. and Smith, S. T. and Song, Y. and Stein, A. and Szegedy, M. and Teets, F. D. and Thyme, S. B. and Wang, R. Y. and Watkins, A. M. and Zimmerman, L. and Bonneau, R.}, title = {Macromolecular modeling and design in rosetta: recent methods and frameworks}, journal = {Nature Methods}, year = {2020}, volume = {17}, issue = {7}, pages = {665-680}, doi = {10.1038/s41592-020-0848-2} }""", + "Rosetta": """@article{10.1038/s41592-020-0848-2, +author = {Leman, J. K. and Weitzner, B. D. and Lewis, S. M. +and Adolf‐Bryfogle, J. and Alam, N. and Alford, R. F. +and Aprahamian, M. L. and Baker, D. and Barlow, K. A. and Barth, P. +and Basanta, B. and Bender, B. J. and Blacklock, K. and Bonet, J. +and Boyken, S. E. and Bradley, P. and Bystroff, C. and Conway, P. +and Cooper, S. and Correia, B. E. and Coventry, B. and Das, R. +and Jong, R. M. d. and DiMaio, F. and Dsilva, L. and Dunbrack, R. L. +and Ford, A. S. and Frenz, B. and Fu, D. and Geniesse, C. +and Goldschmidt, L. and Gowthaman, R. and Gray, J. J. and Gront, D. +and Guffy, S. L. and Horowitz, S. and Huang, P. and Huber, T. +and Jacobs, T. M. and Jeliazkov, J. R. and Johnson, D. K. +and Kappel, K. and Karanicolas, J. and Khakzad, H. and Khar, K. R. +and Khare, S. D. and Khatib, F. and Khramushin, A. and King, C. +and Kleffner, R. and Koepnick, B. and Kortemme, T. and Kuenze, G. +and Kuhlman, B. and Kuroda, D. and Labonte, J. W. and Lai, J. +and Lapidoth, G. and Leaver‐Fay, A. and Lindert, S. and Linsky, T. W. +and London, N. and Lubin, J. H. and Lyskov, S. and Maguire, J. B. +and Malmström, L. and Marcos, E. and Marcu, O. and Marze, N. +and Meiler, J. and Moretti, R. and Mulligan, V. K. and Nerli, S. +and Norn, C. and Ó’Conchúir, S. and Ollikainen, N. and Ovchinnikov, S. +and Pacella, M. S. and Pan, X. and Park, H. and Pavlovicz, R. E. +and Pethe, M. A. and Pierce, B. G. and Pilla, K. B. and Raveh, B. +and Renfrew, P. D. and Burman, S. S. R. and Rubenstein, A. B. +and Sauer, M. F. and Scheck, A. and Schief, W. R. +and Schueler‐Furman, O. and Sedan, Y. and Sevy, A. M. +and Sgourakis, N. G. and Shi, L. and Siegel, J. B. and Silva, D. +and Smith, S. T. and Song, Y. and Stein, A. and Szegedy, M. +and Teets, F. D. and Thyme, S. B. and Wang, R. Y. +and Watkins, A. M. and Zimmerman, L. and Bonneau, R.}, +title = {Macromolecular modeling and design in rosetta: recent methods and frameworks}, +journal = {Nature Methods}, +year = {2020}, +volume = {17}, +issue = {7}, +pages = {665-680}, +doi = {10.1038/s41592-020-0848-2} +}""", "Rosetta3": """ @incollection{LEAVERFAY2011545, -title = {Chapter nineteen - Rosetta3: An Object-Oriented Software Suite for the Simulation and Design of Macromolecules}, +title = {Chapter nineteen - Rosetta3: An Object-Oriented Software Suite for the +Simulation and Design of Macromolecules}, editor = {Michael L. Johnson and Ludwig Brand}, series = {Methods in Enzymology}, publisher = {Academic Press}, @@ -304,27 +343,75 @@ def read_rosetta_node_config() -> dict[str, Any]: issn = {0076-6879}, doi = {https://doi.org/10.1016/B978-0-12-381270-4.00019-6}, url = {https://www.sciencedirect.com/science/article/pii/B9780123812704000196}, -author = {Andrew Leaver-Fay and Michael Tyka and Steven M. Lewis and Oliver F. Lange and James Thompson and Ron Jacak and Kristian W. Kaufman and P. Douglas Renfrew and Colin A. Smith and Will Sheffler and Ian W. Davis and Seth Cooper and Adrien Treuille and Daniel J. Mandell and Florian Richter and Yih-En Andrew Ban and Sarel J. Fleishman and Jacob E. Corn and David E. Kim and Sergey Lyskov and Monica Berrondo and Stuart Mentzer and Zoran Popović and James J. Havranek and John Karanicolas and Rhiju Das and Jens Meiler and Tanja Kortemme and Jeffrey J. Gray and Brian Kuhlman and David Baker and Philip Bradley}, -abstract = {We have recently completed a full rearchitecturing of the Rosetta molecular modeling program, generalizing and expanding its existing functionality. The new architecture enables the rapid prototyping of novel protocols by providing easy-to-use interfaces to powerful tools for molecular modeling. The source code of this rearchitecturing has been released as Rosetta3 and is freely available for academic use. At the time of its release, it contained 470,000 lines of code. Counting currently unpublished protocols at the time of this writing, the source includes 1,285,000 lines. Its rapid growth is a testament to its ease of use. This chapter describes the requirements for our new architecture, justifies the design decisions, sketches out central classes, and highlights a few of the common tasks that the new software can perform.} +author = {Andrew Leaver-Fay and Michael Tyka and Steven M. Lewis +and Oliver F. Lange and James Thompson and Ron Jacak and Kristian W. Kaufman +and P. Douglas Renfrew and Colin A. Smith and Will Sheffler and Ian W. Davis +and Seth Cooper and Adrien Treuille and Daniel J. Mandell and Florian Richter +and Yih-En Andrew Ban and Sarel J. Fleishman and Jacob E. Corn +and David E. Kim and Sergey Lyskov and Monica Berrondo and Stuart Mentzer +and Zoran Popović and James J. Havranek and John Karanicolas and Rhiju Das +and Jens Meiler and Tanja Kortemme and Jeffrey J. Gray and Brian Kuhlman +and David Baker and Philip Bradley}, +abstract = {We have recently completed a full rearchitecturing of the Rosetta +molecular modeling program, generalizing and expanding its existing +functionality. The new architecture enables the rapid prototyping of novel +protocols by providing easy-to-use interfaces to powerful tools for molecular +modeling. The source code of this rearchitecturing has been released as Rosetta3 +and is freely available for academic use. At the time of its release, it +contained 470,000 lines of code. Counting currently unpublished protocols at the +time of this writing, the source includes 1,285,000 lines. Its rapid growth is a +testament to its ease of use. This chapter describes the requirements for our +new architecture, justifies the design decisions, sketches out central classes, +and highlights a few of the common tasks that the new software can perform.} }""", "RosettaScripts": """ @article{10.1371/journal.pone.0020161, doi = {10.1371/journal.pone.0020161}, - author = {Fleishman, Sarel J. AND Leaver-Fay, Andrew AND Corn, Jacob E. AND Strauch, Eva-Maria AND Khare, Sagar D. AND Koga, Nobuyasu AND Ashworth, Justin AND Murphy, Paul AND Richter, Florian AND Lemmon, Gordon AND Meiler, Jens AND Baker, David}, + author = {Fleishman, Sarel J. AND Leaver-Fay, Andrew AND Corn, Jacob E. + AND Strauch, Eva-Maria AND Khare, Sagar D. AND Koga, Nobuyasu + AND Ashworth, Justin AND Murphy, Paul AND Richter, Florian + AND Lemmon, Gordon AND Meiler, Jens AND Baker, David}, journal = {PLOS ONE}, publisher = {Public Library of Science}, - title = {RosettaScripts: A Scripting Language Interface to the Rosetta Macromolecular Modeling Suite}, + title = {RosettaScripts: A Scripting Language Interface to the Rosetta + Macromolecular Modeling Suite}, year = {2011}, month = {06}, volume = {6}, url = {https://doi.org/10.1371/journal.pone.0020161}, pages = {1-10}, - abstract = {Macromolecular modeling and design are increasingly useful in basic research, biotechnology, and teaching. However, the absence of a user-friendly modeling framework that provides access to a wide range of modeling capabilities is hampering the wider adoption of computational methods by non-experts. RosettaScripts is an XML-like language for specifying modeling tasks in the Rosetta framework. RosettaScripts provides access to protocol-level functionalities, such as rigid-body docking and sequence redesign, and allows fast testing and deployment of complex protocols without need for modifying or recompiling the underlying C++ code. We illustrate these capabilities with RosettaScripts protocols for the stabilization of proteins, the generation of computationally constrained libraries for experimental selection of higher-affinity binding proteins, loop remodeling, small-molecule ligand docking, design of ligand-binding proteins, and specificity redesign in DNA-binding proteins.}, + abstract = {Macromolecular modeling and design are increasingly useful in basic + research, biotechnology, and teaching. However, the absence of a + user-friendly modeling framework that provides access to a wide range of + modeling capabilities is hampering the wider adoption of computational + methods by non-experts. RosettaScripts is an XML-like language for + specifying modeling tasks in the Rosetta framework. RosettaScripts provides + access to protocol-level functionalities, such as rigid-body docking and + sequence redesign, and allows fast testing and deployment of complex + protocols without need for modifying or recompiling the underlying C++ code. + We illustrate these capabilities with RosettaScripts protocols for the + stabilization of proteins, the generation of computationally constrained + libraries for experimental selection of higher-affinity binding proteins, + loop remodeling, small-molecule ligand docking, design of ligand-binding + proteins, and specificity redesign in DNA-binding proteins.}, number = {6}, }""", "Rosetta All-Atom Energy Function": """ -@article{10.1021/acs.jctc.7b00125, author = {Alford, R. F. and Leaver‐Fay, A. and Jeliazkov, J. R. and O’Meara, M. J. and DiMaio, F. and Park, H. and Shapovalov, M. V. and Renfrew, P. D. and Mulligan, V. K. and Kappel, K. and Labonte, J. W. and Pacella, M. S. and Bonneau, R. and Bradley, P. and Dunbrack, R. L. and Das, R. and Baker, D. and Kuhlman, B. and Kortemme, T. and Gray, J. J.}, title = {The rosetta all-atom energy function for macromolecular modeling and design}, journal = {Journal of Chemical Theory and Computation}, year = {2017}, volume = {13}, issue = {6}, pages = {3031-3048}, doi = {10.1021/acs.jctc.7b00125} }""", +@article{10.1021/acs.jctc.7b00125, +author = {Alford, R. F. and Leaver‐Fay, A. and Jeliazkov, J. R. +and O’Meara, M. J. and DiMaio, F. and Park, H. and Shapovalov, M. V. +and Renfrew, P. D. and Mulligan, V. K. and Kappel, K. and Labonte, J. W. +and Pacella, M. S. and Bonneau, R. and Bradley, P. and Dunbrack, R. L. +and Das, R. and Baker, D. and Kuhlman, B. and Kortemme, T. and Gray, J. J.}, +title = {The rosetta all-atom energy function for macromolecular modeling and design}, +journal = {Journal of Chemical Theory and Computation}, +year = {2017}, +volume = {13}, +issue = {6}, +pages = {3031-3048}, +doi = {10.1021/acs.jctc.7b00125} +}""", } # TODO: refactor as general citation merging function diff --git a/src/REvoDesign/tools/safe_pickle.py b/src/REvoDesign/tools/safe_pickle.py index ee189f983..2d88f709f 100644 --- a/src/REvoDesign/tools/safe_pickle.py +++ b/src/REvoDesign/tools/safe_pickle.py @@ -4,7 +4,9 @@ import importlib import io -import pickle # nosec B403 - this module provides a restricted unpickler wrapper + +# This module provides a restricted unpickler wrapper. +import pickle # nosec B403 from collections import OrderedDict, defaultdict, deque from collections.abc import Iterable from typing import BinaryIO diff --git a/src/REvoDesign/tools/ssl_certificates.py b/src/REvoDesign/tools/ssl_certificates.py index f8d75f27a..9d3b8bb8f 100644 --- a/src/REvoDesign/tools/ssl_certificates.py +++ b/src/REvoDesign/tools/ssl_certificates.py @@ -4,12 +4,14 @@ """ -This module contains functions and classes related to managing SSL certificates and generating unique identifiers (UUIDs). +This module contains functions and classes related to managing SSL certificates +and generating unique identifiers (UUIDs). """ import datetime import os import platform +import secrets import ssl from dataclasses import dataclass from typing import Literal @@ -21,6 +23,39 @@ logging = ROOT_LOGGER.getChild(__name__) +DEFAULT_CERT_VALIDITY_DAYS = 7 +CERT_VALIDITY_ENV = "REVODESIGN_SSL_CERT_VALIDITY_DAYS" +PRIVATE_DIR_MODE = 0o700 +PRIVATE_KEY_MODE = 0o600 +CERTIFICATE_MODE = 0o644 + + +def _certificate_validity_days() -> int: + raw_days = os.environ.get(CERT_VALIDITY_ENV) + if raw_days is None: + return DEFAULT_CERT_VALIDITY_DAYS + + try: + days = int(raw_days) + except ValueError as exc: + raise ValueError(f"{CERT_VALIDITY_ENV} must be a positive integer.") from exc + + if days <= 0: + raise ValueError(f"{CERT_VALIDITY_ENV} must be a positive integer.") + + return days + + +def _write_pem_file(path: str, payload: bytes, mode: int) -> None: + flags = os.O_WRONLY | os.O_CREAT | os.O_TRUNC + fd = os.open(path, flags, mode) + try: + with os.fdopen(fd, "wb") as handle: + handle.write(payload) + finally: + os.chmod(path, mode) + + @dataclass class SSLCertificateManager: peer_roles = Literal["server", "client"] @@ -32,7 +67,8 @@ def __post_init__(self): self.cache_dir: str = set_cache_dir() self.crt_dir: str = os.path.join(self.cache_dir, "crts") - os.makedirs(self.crt_dir, exist_ok=True) + os.makedirs(self.crt_dir, mode=PRIVATE_DIR_MODE, exist_ok=True) + os.chmod(self.crt_dir, PRIVATE_DIR_MODE) self.crt_path = os.path.join(self.crt_dir, f"{self.role}.crt") self.key_path = os.path.join(self.crt_dir, f"{self.role}.key") @@ -69,7 +105,8 @@ def get_certificate(self): Function: get_certificate Usage: get_certificate() - This function checks for the existence of an SSL certificate and generates a new one if it doesn't exist or has expired. + This function checks for the existence of an SSL certificate and + generates a new one if it doesn't exist or has expired. Returns: @@ -93,6 +130,7 @@ def get_certificate(self): logging.warning("Certificate has expired. Generating a new certificate.") self.create_new_certificate() else: + self._enforce_file_permissions() logging.info("Certificate is still valid.") def create_new_certificate(self): @@ -100,12 +138,14 @@ def create_new_certificate(self): Function: create_new_certificate Usage: create_new_certificate() - This function creates a new SSL certificate and private key if they do not exist or if the certificate has expired. + This function creates a new SSL certificate and private key if they do + not exist or if the certificate has expired. Returns: - None """ role = self.role + validity_days = _certificate_validity_days() # Get node information from OS or set to 'Unknown' if not available @@ -132,15 +172,20 @@ def create_new_certificate(self): cert.get_subject().CN = _cn[:64] # Set serial number, validity period, issuer, public key, and sign the certificate - cert.set_serial_number(1000) + cert.set_serial_number(secrets.randbits(159) + 1) cert.gmtime_adj_notBefore(0) - cert.gmtime_adj_notAfter(7 * 24 * 60 * 60) + cert.gmtime_adj_notAfter(validity_days * 24 * 60 * 60) cert.set_issuer(cert.get_subject()) cert.set_pubkey(k) cert.sign(k, "sha256") # Write the certificate and private key to files in PEM format - with open(self.crt_path, "wb") as f: - f.write(crypto.dump_certificate(crypto.FILETYPE_PEM, cert)) - with open(self.key_path, "wb") as f: - f.write(crypto.dump_privatekey(crypto.FILETYPE_PEM, k)) + _write_pem_file(self.crt_path, crypto.dump_certificate(crypto.FILETYPE_PEM, cert), CERTIFICATE_MODE) + _write_pem_file(self.key_path, crypto.dump_privatekey(crypto.FILETYPE_PEM, k), PRIVATE_KEY_MODE) + + def _enforce_file_permissions(self): + os.chmod(self.crt_dir, PRIVATE_DIR_MODE) + if os.path.exists(self.crt_path): + os.chmod(self.crt_path, CERTIFICATE_MODE) + if os.path.exists(self.key_path): + os.chmod(self.key_path, PRIVATE_KEY_MODE) diff --git a/src/REvoDesign/tools/system_tools.py b/src/REvoDesign/tools/system_tools.py index a924a1503..642c51474 100644 --- a/src/REvoDesign/tools/system_tools.py +++ b/src/REvoDesign/tools/system_tools.py @@ -24,7 +24,8 @@ def check_mac_rosetta2(): This function first checks if the current operating system is macOS. If not, it returns immediately. It then determines if the machine is an ARM-based Mac and whether it is recognized as an x86_64 architecture. - If both conditions are met, a warning is issued indicating that the environment is using Rosetta 2, which may impact performance. + If both conditions are met, a warning is issued indicating that the + environment is using Rosetta 2, which may impact performance. Returns: None @@ -49,7 +50,8 @@ def check_mac_rosetta2(): issues.AppleSiliconRosetta2Warning( "Oops! You are in Rosetta-translated PyMOL bundle from official channel. " "This might limit the performance of joblib, causing MutantVisualizer slower." - ) + ), + stacklevel=2, ) diff --git a/src/REvoDesign/tools/utils.py b/src/REvoDesign/tools/utils.py index bb14c3181..7d52711a2 100644 --- a/src/REvoDesign/tools/utils.py +++ b/src/REvoDesign/tools/utils.py @@ -13,7 +13,7 @@ import importlib import inspect import os -import random +import secrets import shutil import string import sys @@ -23,6 +23,7 @@ from collections.abc import Callable, Iterable from functools import partial, wraps from itertools import pairwise +from pathlib import PurePosixPath, PureWindowsPath from typing import Any, Literal, ParamSpec, TypeVar, cast, overload import matplotlib @@ -160,17 +161,16 @@ def resolve_lambda_expression(expression: str, as_partial: bool | None = None) - def resolve_typed_arg(arg: str) -> Any: if arg.lower() == "true": return True - elif arg.lower() == "false": + if arg.lower() == "false": return False # integer - elif arg.isdigit(): + if arg.isdigit(): return int(arg) # float - elif arg.replace(".", "", 1).isdigit(): + if arg.replace(".", "", 1).isdigit(): return float(arg) # string fallback - else: - return arg + return arg def resolve_default_value(typing: type) -> Any: @@ -182,6 +182,7 @@ def resolve_default_value(typing: type) -> Any: return 0.0 if typing == str: return "" + return None def resolve_dotted_config_item(config_string: str) -> Any: @@ -190,7 +191,8 @@ def resolve_dotted_config_item(config_string: str) -> Any: The input string can be in one of the following formats: 1. `"CFG::"` - Specifies both the configuration section and item. - 2. `"CFG:"` - Specifies only the configuration item, with no section(default for the main config section). + 2. `"CFG:"` - Specifies only the configuration item, with + no section(default for the main config section). Args: config_string (str): The dotted configuration string to resolve. @@ -316,6 +318,59 @@ def __init__(self, *args, **kwargs): MethodKind = Literal["InstanceMethod", "ClassMethod", "StaticMethod", "Function"] +def _resolve_function_owner(method: Callable) -> tuple[type | None, str, str]: + qualname = getattr(method, "__qualname__", "") + if "." not in qualname: + return None, "", qualname + + owner_path, func_name = qualname.rsplit(".", 1) + module = inspect.getmodule(method) + if module is None: + return None, owner_path, func_name + + obj: Any = module + for part in owner_path.split("."): + obj = getattr(obj, part, None) + if obj is None: + return None, owner_path, func_name + + return (obj if isinstance(obj, type) else None), owner_path, func_name + + +def _inspect_bound_method(method: Callable) -> MethodKind | None: + self_obj = getattr(method, "__self__", None) + if self_obj is None: + return None + if isinstance(self_obj, type): + logging.debug(f"{method!r} is a classmethod bound to {self_obj!r}") + return "ClassMethod" + logging.debug(f"{method!r} is an instancemethod bound to {self_obj!r}") + return "InstanceMethod" + + +def _inspect_function(method: Callable) -> MethodKind | None: + if not inspect.isfunction(method): + return None + + owner_cls, owner_path, func_name = _resolve_function_owner(method) + if not owner_path: + logging.debug(f"{method!r} is a plain top-level function") + return "Function" + + if owner_cls is None: + logging.debug(f"{method!r} is a dotted qualname but class is not resolvable") + return "StaticMethod" + + attr = owner_cls.__dict__.get(func_name) + if isinstance(attr, staticmethod): + logging.debug(f"{method!r} is a staticmethod on {owner_cls!r}") + return "StaticMethod" + if inspect.isfunction(attr): + logging.debug(f"{method!r} is a function on {owner_cls!r}") + return "Function" + return None + + def inspect_method_types(method: Callable) -> MethodKind: """ Inspect the type of a callable: instance method, class method, @@ -335,76 +390,15 @@ def inspect_method_types(method: Callable) -> MethodKind: if not callable(method): raise issues.UnexpectedWorkflowError(f"Cannot inspect method type for non-callable {method!r}") - # --- 1. Bound methods: have a non-None __self__ ------------------------- - # obj.method -> bound instance method ( __self__ is instance ) - # cls.method -> bound class method ( __self__ is class ) - if hasattr(method, "__self__") and getattr(method, "__self__", None) is not None: - self_obj = method.__self__ # type: ignore[attr-defined] - # Bound to a class -> classmethod - if isinstance(self_obj, type): - logging.debug(f"{method!r} is a classmethod bound to {self_obj!r}") - return "ClassMethod" - logging.debug(f"{method!r} is an instancemethod bound to {self_obj!r}") - # Bound to an instance -> instance method - return "InstanceMethod" - - # --- 2. Unbound functions (no binding) ----------------------------------- - # At this point, typical Python methods accessed from the class are plain - # function objects (unbound). We must check how they are stored on the class. - if inspect.isfunction(method): - qualname = getattr(method, "__qualname__", "") - logging.debug(f"Checking unbound function {method}: {getattr(method, '__module__', '')!r}") - - # Try to resolve the owning class from qualname: "MyClass.method" - if "." in qualname: - owner_path, func_name = qualname.rsplit(".", 1) - - owner_cls = None - module = inspect.getmodule(method) - if module is not None: - obj = module - # Walk the dotted path "Outer.Inner" etc. - for part in owner_path.split("."): - if not hasattr(obj, part): - obj = None - break - obj = getattr(obj, part) - if isinstance(obj, type): - owner_cls = obj - - if owner_cls is not None: - # Look at the raw attribute on the class dict - attr = owner_cls.__dict__.get(func_name) - # Static method is stored as a staticmethod descriptor - if isinstance(attr, staticmethod): - logging.debug(f"{method!r} is a staticmethod on {owner_cls!r}") - return "StaticMethod" - # A normal "def" on the class body: unbound instance method - if inspect.isfunction(attr): - logging.debug(f"{method!r} is a function on {owner_cls!r}") - return "Function" - - # Fallback: dotted qualname but class not resolvable - # (e.g. local classes with '' in qualname). - # In your semantics we still treat this as a static-style method. - logging.debug(f"{method!r} is a dotted qualname but class is not resolvable") - return "StaticMethod" - - # No dot in qualname: a plain top-level function - logging.debug(f"{method!r} is a plain top-level function") - return "Function" + detected = _inspect_bound_method(method) or _inspect_function(method) + if detected is not None: + return detected - # --- 3. Rare cases: method-like objects not covered above --------------- if inspect.ismethod(method): - # This branch is mostly for odd cases / builtins. - self_obj = method.__self__ - if isinstance(self_obj, type): - logging.debug(f"{method!r} is a method-like classmethod object") - return "ClassMethod" - logging.debug(f"{method!r} is a method-like instancemethod object") - return "InstanceMethod" - - # --- 4. Give up ---------------------------------------------------------- + detected = _inspect_bound_method(method) + if detected is not None: + return detected + logging.error(f"Giving up for inpection on method type! {method!r}") raise TypeError(f"Cannot inspect method type for {method!r}") @@ -483,10 +477,11 @@ def _cite_for_cls(cls_or_obj: type[CitableModuleAbstract] | CitableModuleAbstrac cls_or_obj: A class or instance object of type CitableModuleAbstract """ try: - if not (hasattr(cls_or_obj, "cite") and callable(getattr(cls_or_obj, "cite"))): + cite_method = getattr(cls_or_obj, "cite", None) + if not callable(cite_method): name = getattr(cls_or_obj, "__name__", type(cls_or_obj).__name__) raise TypeError(f"{name} is not citable or missing cite()") - cls_or_obj.cite() + cite_method() except Exception as e: logging.warning(f"Ignore cite() error: {e}") @@ -596,7 +591,6 @@ def data_generator(): """ current_batch = [] for data_point in inputs_data_generator: - # print(f"Send data {data_point}") current_batch.append(data_point) if len(current_batch) == batch_size: yield current_batch @@ -607,79 +601,92 @@ def data_generator(): yield current_batch +def _safe_member_target(base_dir: str, member_name: str) -> str: + normalized = member_name.replace("\\", "/") + posix_path = PurePosixPath(normalized) + windows_path = PureWindowsPath(normalized) + if ( + not normalized + or normalized == "." + or posix_path.is_absolute() + or windows_path.is_absolute() + or windows_path.drive + or any(part == ".." for part in posix_path.parts) + ): + raise ValueError(f"Unsafe archive member path: {member_name}") + target_path = os.path.realpath(os.path.join(base_dir, *posix_path.parts)) + base_path = os.path.realpath(base_dir) + try: + if os.path.commonpath([base_path, target_path]) != base_path: + raise ValueError(f"Unsafe archive member path: {member_name}") + except ValueError as exc: + raise ValueError(f"Unsafe archive member path: {member_name}") from exc + return target_path + + +def _extract_zip_safely(zip_ref: zipfile.ZipFile, base_dir: str) -> None: + for member in zip_ref.infolist(): + target_path = _safe_member_target(base_dir, member.filename) + if member.is_dir(): + os.makedirs(target_path, exist_ok=True) + continue + os.makedirs(os.path.dirname(target_path), exist_ok=True) + with zip_ref.open(member, "r") as source, open(target_path, "wb") as dest: + shutil.copyfileobj(source, dest) + + +def _extract_tar_safely(tar_ref: tarfile.TarFile, base_dir: str) -> None: + for member in tar_ref.getmembers(): + target_path = _safe_member_target(base_dir, member.name) + if member.isdir(): + os.makedirs(target_path, exist_ok=True) + continue + if not member.isfile(): + logging.warning("Skipping non-regular archive member: %s", member.name) + continue + stream = tar_ref.extractfile(member) + if stream is None: + raise ValueError(f"Failed to read archive member: {member.name}") + os.makedirs(os.path.dirname(target_path), exist_ok=True) + with contextlib.closing(stream) as source, open(target_path, "wb") as dest: + shutil.copyfileobj(source, dest) + + +def _extract_archive_contents(archive_file: str, extract_to: str) -> None: + if archive_file.endswith(".zip"): + with zipfile.ZipFile(archive_file, "r") as zip_ref: + _extract_zip_safely(zip_ref, extract_to) + return + + tar_modes = { + ".tar": "r:", + ".tar.gz": "r:*", + ".tgz": "r:*", + ".tar.bz2": "r:bz2", + ".tbz": "r:bz2", + ".tar.xz": "r:xz", + } + for suffix, mode in tar_modes.items(): + if archive_file.endswith(suffix): + with tarfile.open(archive_file, mode) as tar_ref: + _extract_tar_safely(tar_ref, extract_to) + return + + raise ValueError(f"Unsupported archive format: {archive_file}") + + def extract_archive(archive_file: str, extract_to: str): """ - Extracts the contents of an archive file (zip, tar.gz, tar.bz2, tar.xz, or rar) to a specified directory. + Extracts an archive file (zip, tar, tar.gz, tar.bz2, or tar.xz) to a directory. Args: archive_file (str): Path to the archive file. extract_to (str): Directory where the contents will be extracted. """ - - def _safe_member_target(base_dir: str, member_name: str) -> str: - normalized = member_name.replace("\\", "/") - if normalized.startswith("/") or normalized.startswith("../") or "/../" in f"/{normalized}": - raise ValueError(f"Unsafe archive member path: {member_name}") - target_path = os.path.realpath(os.path.join(base_dir, member_name)) - base_path = os.path.realpath(base_dir) - try: - if os.path.commonpath([base_path, target_path]) != base_path: - raise ValueError(f"Unsafe archive member path: {member_name}") - except ValueError as exc: - raise ValueError(f"Unsafe archive member path: {member_name}") from exc - return target_path - - def _extract_zip_safely(zip_ref: zipfile.ZipFile, base_dir: str) -> None: - for member in zip_ref.infolist(): - target_path = _safe_member_target(base_dir, member.filename) - if member.is_dir(): - os.makedirs(target_path, exist_ok=True) - continue - os.makedirs(os.path.dirname(target_path), exist_ok=True) - with zip_ref.open(member, "r") as source, open(target_path, "wb") as dest: - shutil.copyfileobj(source, dest) - - def _extract_tar_safely(tar_ref: tarfile.TarFile, base_dir: str) -> None: - for member in tar_ref.getmembers(): - target_path = _safe_member_target(base_dir, member.name) - if member.isdir(): - os.makedirs(target_path, exist_ok=True) - continue - if not member.isfile(): - logging.warning("Skipping non-regular archive member: %s", member.name) - continue - stream = tar_ref.extractfile(member) - if stream is None: - raise ValueError(f"Failed to read archive member: {member.name}") - os.makedirs(os.path.dirname(target_path), exist_ok=True) - with contextlib.closing(stream) as source, open(target_path, "wb") as dest: - shutil.copyfileobj(source, dest) - try: with timing(f"extracting {archive_file} to {extract_to}"): os.makedirs(extract_to, exist_ok=True) - if archive_file.endswith(".zip"): - with zipfile.ZipFile(archive_file, "r") as zip_ref: - _extract_zip_safely(zip_ref, extract_to) - - elif archive_file.endswith(".tar"): - with tarfile.open(archive_file, "r:") as tar_ref: - _extract_tar_safely(tar_ref, extract_to) - - elif archive_file.endswith((".tar.gz", ".tgz")): - with tarfile.open(archive_file, "r:*") as tar_ref: - _extract_tar_safely(tar_ref, extract_to) - - elif archive_file.endswith((".tar.bz2", ".tbz")): - with tarfile.open(archive_file, "r:bz2") as tar_ref: - _extract_tar_safely(tar_ref, extract_to) - - elif archive_file.endswith(".tar.xz"): - with tarfile.open(archive_file, "r:xz") as tar_ref: - _extract_tar_safely(tar_ref, extract_to) - - else: - raise ValueError(f"Unsupported archive format: {archive_file}") + _extract_archive_contents(archive_file, extract_to) except Exception as e: logging.error(f"Error extracting {archive_file}: {str(e)}") raise ValueError(f"Failed to extract {archive_file}: {e}") from e @@ -821,8 +828,7 @@ def generate_strong_password(length=16): # Define the characters to use for generating the password password_characters = string.ascii_letters + string.digits + "!#$%&*+-./:?@^_~" - # Generate the password using random characters from the defined set - generated_password = "".join(random.choice(password_characters) for _ in range(length)) + generated_password = "".join(secrets.choice(password_characters) for _ in range(length)) return generated_password diff --git a/tests/bootstrap/test_set_config.py b/tests/bootstrap/test_set_config.py index f882d3dcb..5be6c9f07 100644 --- a/tests/bootstrap/test_set_config.py +++ b/tests/bootstrap/test_set_config.py @@ -14,7 +14,6 @@ _iter_yaml_rel_paths, experiment_config, is_package_installed, - list_all_config_files, reload_config_file, save_configuration, set_cache_dir, @@ -117,7 +116,7 @@ def test_main_config_upgrade(user_decide, patch_config_user_data): patch("REvoDesign.bootstrap.set_config.shutil.copytree") as mock_copytree, ): - set_REvoDesign_config_file() + set_REvoDesign_config_file(prompt_user=True) mock_decide.assert_called_once() mock_copytree.assert_called_once() if user_decide: @@ -126,6 +125,40 @@ def test_main_config_upgrade(user_decide, patch_config_user_data): mock_rmtree.assert_not_called() +def test_main_config_upgrade_headless_keeps_existing_tree(patch_config_user_data): + with ( + patch("REvoDesign.bootstrap.set_config.decide") as mock_decide, + patch("REvoDesign.bootstrap.set_config.os.path.isfile", return_value=False), + patch("REvoDesign.bootstrap.set_config.os.path.isdir", return_value=True), + patch("REvoDesign.bootstrap.set_config.os.listdir", return_value=["master.yaml", "a_dir"]), + patch("REvoDesign.bootstrap.set_config.shutil.rmtree") as mock_rmtree, + patch("REvoDesign.bootstrap.set_config.shutil.copytree") as mock_copytree, + ): + set_REvoDesign_config_file(prompt_user=False) + + mock_decide.assert_not_called() + mock_rmtree.assert_not_called() + mock_copytree.assert_called_once() + + +def test_main_config_upgrade_env_reset_is_explicit(monkeypatch, patch_config_user_data): + monkeypatch.setenv("REVODESIGN_RESET_CONFIG", "1") + with ( + patch("REvoDesign.bootstrap.set_config.decide") as mock_decide, + patch("REvoDesign.bootstrap.set_config.os.path.isfile", return_value=False), + patch("REvoDesign.bootstrap.set_config.os.path.isdir", return_value=True), + patch("REvoDesign.bootstrap.set_config.os.listdir", return_value=["master.yaml", "a_dir"]), + patch("REvoDesign.bootstrap.set_config.os.path.exists", return_value=True), + patch("REvoDesign.bootstrap.set_config.shutil.rmtree") as mock_rmtree, + patch("REvoDesign.bootstrap.set_config.shutil.copytree") as mock_copytree, + ): + set_REvoDesign_config_file(prompt_user=False) + + mock_decide.assert_not_called() + mock_rmtree.assert_called_once() + mock_copytree.assert_called_once() + + def test_iter_yaml_rel_paths(tmp_path): base_dir = tmp_path / "configs" base_dir.mkdir() diff --git a/tests/conftest.py b/tests/conftest.py index 8987b62cb..b03e472bd 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -10,6 +10,7 @@ from __future__ import annotations import gc +import hmac import importlib.util import json import os @@ -23,7 +24,7 @@ from typing import Literal from unittest.mock import MagicMock, patch -# import hydra as unmocked_hydra +import platformdirs import psutil import pytest from _pytest.nodes import Item @@ -35,9 +36,36 @@ from RosettaPy.node import NodeHintT from RosettaPy.utils import tmpdir_manager +# --------------------------------------------------------------------------- +# Keep import-time REvoDesign bootstrap away from the user's real config/cache tree. +# MUST run BEFORE any REvoDesign module import, which triggers bootstrap config resolution +# via ``REVODESIGN_CONFIG_FILE = set_REvoDesign_config_file()``. +# --------------------------------------------------------------------------- +TEST_ROOT = os.path.abspath(".") +DATA_DIRNAME = os.path.join(TEST_ROOT, "mock", "user_data", "REvoDesign") +CACHE_DIRNAME = os.path.join(TEST_ROOT, "..", "tests", "downloaded", "cache") +EXPECTED_MAIN_CONFIG_FILE = os.path.join(DATA_DIRNAME, "config", "main.yaml") +os.makedirs(DATA_DIRNAME, exist_ok=True) +os.makedirs(CACHE_DIRNAME, exist_ok=True) + + +def _static_platform_dir(dirname: str): + def _impl(*args, **kwargs): + Path(dirname).mkdir(parents=True, exist_ok=True) + return dirname + + return _impl + + +platformdirs.user_data_dir = _static_platform_dir(DATA_DIRNAME) +platformdirs.user_cache_dir = _static_platform_dir(CACHE_DIRNAME) + +# import hydra as unmocked_hydra + +# isort: split -- REvoDesign imports below are safe now that platformdirs is patched +import REvoDesign.tools.package_manager as package_manager from REvoDesign import REvoDesignPlugin from REvoDesign.basic.abc_singleton import reset_singletons -from REvoDesign.bootstrap import EXPERIMENTS_CONFIG_DIR from REvoDesign.bootstrap.set_config import ConfigConverter, reload_config_file, set_REvoDesign_config_file from REvoDesign.common import MutantTree from REvoDesign.driver.ui_driver import ConfigBus @@ -48,6 +76,13 @@ from .data import TestData from .data.test_data import KeyData +REPO_DIR = os.path.join(os.path.dirname(__file__), "..") +TESTS_DIR = Path(__file__).resolve().parent +REPO_ROOT_PATH = TESTS_DIR.parent +MSA_ROOT = TESTS_DIR / "data" / "msa" +MINIUC_ROOT = MSA_ROOT / "miniuc" +TESTMINIUC_ROOT = MSA_ROOT / "testminiuc" + def _resolve_pytest_qt_api() -> str: backend = getattr(pymol_qt, "PYQT_NAME", "") @@ -64,23 +99,6 @@ def _resolve_pytest_qt_api() -> str: PYTEST_QT_API = os.environ.setdefault("PYTEST_QT_API", _resolve_pytest_qt_api()) -REPO_DIR = os.path.join(os.path.dirname(__file__), "..") -TESTS_DIR = Path(__file__).resolve().parent -REPO_ROOT_PATH = TESTS_DIR.parent -MSA_ROOT = TESTS_DIR / "data" / "msa" -MINIUC_ROOT = MSA_ROOT / "miniuc" -TESTMINIUC_ROOT = MSA_ROOT / "testminiuc" - -# mostly mock on data and cache to isolated from user's production system -TEST_ROOT = os.path.abspath(".") -DATA_DIRNAME = os.path.join(TEST_ROOT, "mock", "user_data", "REvoDesign") -CACHE_DIRNAME = os.path.join(TEST_ROOT, "..", "tests", "downloaded", "cache") - -EXPECTED_MAIN_CONFIG_FILE = os.path.join(DATA_DIRNAME, "config", "main.yaml") -os.makedirs(DATA_DIRNAME, exist_ok=True) -os.makedirs(CACHE_DIRNAME, exist_ok=True) - - def copy_config_tree(): print("copy config tree to mock user data") shutil.copytree( @@ -116,28 +134,13 @@ def check_real_config_dir(i: int): """ A checkpoint to check whether the test suite has created the real config dir. """ - the_dir = "/Users/yyy/Library/Application Support/REvoDesign/config" - if os.path.exists(the_dir): - raise RuntimeError(f"[{i}]The test suite has created this dir! {the_dir}") + real_config_dir = Path.home() / "Library" / "Application Support" / "REvoDesign" / "config" + if real_config_dir.exists(): + raise RuntimeError(f"[{i}]The test suite has created this dir! {real_config_dir}") # check_real_config_dir(cck.i) -# TODO: isolate from user's system by mocking the platformdirs module. - -# def _static_platform_dir(dirname: str): -# def _impl(*args, **kwargs): -# if not os.path.exists(dirname): -# os.makedirs(dirname) -# return dirname - -# return _impl - - -# platformdirs.user_data_dir = _static_platform_dir(DATA_DIRNAME) -# platformdirs.user_cache_dir = _static_platform_dir(CACHE_DIRNAME) - - # check_real_config_dir(cck.i) # with ( @@ -225,6 +228,29 @@ def mock_fetch_json(url): return d +def _bootstrap_manifest() -> dict[str, str]: + """Generate a valid HMAC manifest from canonical repo files for testing.""" + key = package_manager._MANAGER_HMAC_KEY + files: dict[str, str] = { + "REvoDesign-PyMOL-entry.ui": os.path.join(REPO_DIR, "src/REvoDesign/UI/REvoDesign-PyMOL-entry.ui"), + "REvoDesignExtrasTableRich.json": os.path.join(REPO_DIR, "jsons/REvoDesignExtrasTableRich.json"), + } + return {name: hmac.new(key, open(path, "rb").read(), "sha256").hexdigest() for name, path in files.items()} + + +def _mock_fetch_verified_bootstrap_file( + *, url: str, asset_name: str, save_to_file: str, manifest: dict[str, str], description: str +) -> None: + """Copy canonical repo file instead of fetching from network.""" + if asset_name == package_manager.MANIFEST_ASSET_MANAGER_UI: + src = os.path.join(REPO_DIR, "src/REvoDesign/UI/REvoDesign-PyMOL-entry.ui") + elif asset_name == package_manager.MANIFEST_ASSET_EXTRAS_JSON: + src = os.path.join(REPO_DIR, "jsons/REvoDesignExtrasTableRich.json") + else: + raise ValueError(f"Unknown bootstrap asset: {asset_name}") + shutil.copy(src, save_to_file) + + @pytest.fixture(scope="function") def pm_plugin(qtbot: qtbot.QtBot) -> REvoDesignPackageManager: @@ -236,7 +262,15 @@ def pm_plugin(qtbot: qtbot.QtBot) -> REvoDesignPackageManager: patch( "REvoDesign.tools.package_manager.get_github_repo_tags", side_effect=lambda repo_url: ["1.0.0", "1.0.1", "1.0.2"], - ) as mock_notify_box, + ) as mock_fetch_tags, + patch( + "REvoDesign.tools.package_manager.fetch_bootstrap_manifest", + side_effect=_bootstrap_manifest, + ) as mock_bootstrap_manifest, + patch( + "REvoDesign.tools.package_manager.fetch_verified_bootstrap_file", + side_effect=_mock_fetch_verified_bootstrap_file, + ) as mock_verified_bootstrap, ): plugin = REvoDesignPackageManager() @@ -447,7 +481,8 @@ def load_session_and_check( print(f"nproc: {nproc}") if self.in_which_runner.get("CIRCLECI") and nproc > self.test_data.nproc_circleci: print( - f"Fix nproc to reduce performance for CircleCI: {nproc} {os.cpu_count()}-> {self.test_data.nproc_circleci}" + "Fix nproc to reduce performance for CircleCI: " + f"{nproc} {os.cpu_count()}-> {self.test_data.nproc_circleci}" ) set_widget_value( self.plugin.ui.spinBox_nproc, diff --git a/tests/data/ddg_csv.json b/tests/data/ddg_csv.json index 22069a106..dd2764fae 100644 --- a/tests/data/ddg_csv.json +++ b/tests/data/ddg_csv.json @@ -5,12 +5,12 @@ "__date__": "2025-03-12 17:20:50" }, "__asked_values__": { - "profile": "/Users/yyy/Documents/protein_design/REvoDesign/tests/data/pytia_ddg/1SUO_pred_mask.csv", + "profile": "tests/data/pytia_ddg/1SUO_pred_mask.csv", "profile_type": "CSV", "prefer_lower_score": true, "keep_missing": true, - "residue_range": "/Users/yyy/Documents/protein_design/REvoDesign/tests/data/surface_residue_records/1SUO_residues_cutoff_30.0.txt", + "residue_range": "tests/data/surface_residue_records/1SUO_residues_cutoff_30.0.txt", "view_highlight": "orient", "view_highlight_nbr": 6 } -} \ No newline at end of file +} diff --git a/tests/data/kinetics/openkinetics_1SUO/manifest.json b/tests/data/kinetics/openkinetics_1SUO/manifest.json index 47ac3d48a..363867f84 100644 --- a/tests/data/kinetics/openkinetics_1SUO/manifest.json +++ b/tests/data/kinetics/openkinetics_1SUO/manifest.json @@ -19,8 +19,8 @@ "revodesign_commit_hash": "83e1f78f8349efad62af8c01aab88625ecd6d237", "schema_assumption_note": "Official API docs could not be fetched automatically during implementation; the collector uses the documented fallback endpoint pattern from plan/openkinetics.md.", "secrets_policy": "Authorization headers and API keys are redacted and never stored", - "source_mutation_file": "/Users/yyy/Documents/protein_design/REvoDesign/tests/data/mutations/1SUO.surf.entro.mutagenesis.besthits.mut.txt", - "source_pdb_file": "/Users/yyy/Documents/protein_design/REvoDesign/tests/data/pdb/1SUO.pdb", + "source_mutation_file": "tests/data/mutations/1SUO.surf.entro.mutagenesis.besthits.mut.txt", + "source_pdb_file": "tests/data/pdb/1SUO.pdb", "source_structure_id": "1SUO", "substrate_metadata_source": "rdkit", "substrate_smiles": "ClC1CCC(C2CNCN2)CC1" diff --git a/tests/data/test_data.py b/tests/data/test_data.py index 2060045d6..ca52a40d3 100644 --- a/tests/data/test_data.py +++ b/tests/data/test_data.py @@ -6,8 +6,6 @@ import os from dataclasses import dataclass -import pytest - @dataclass class TestData: @@ -384,10 +382,9 @@ def expand_zip(self, compressed_file: str): expanded_dirs = os.listdir(dist_dir) if not expanded_dirs: - import zipfile + from REvoDesign.tools.utils import extract_archive - with zipfile.ZipFile(compressed_file, mode="r") as z: - z.extractall(path=dist_dir) + extract_archive(archive_file=compressed_file, extract_to=dist_dir) extracted_files = os.listdir(dist_dir) return dist_dir, extracted_files @@ -395,7 +392,10 @@ def expand_zip(self, compressed_file: str): @dataclass() class TestDataOnLocalMac(TestData): - test_data_repo: str = "/Users/yyy/Documents/protein_design/REvoDesignTutorial/" + test_data_repo: str = os.environ.get( + "REVODESIGN_TUTORIAL_DATA", + os.path.expanduser("~/Documents/protein_design/REvoDesignTutorial"), + ) if __name__ == "__main__": diff --git a/tests/driver/test_environ_register.py b/tests/driver/test_environ_register.py new file mode 100644 index 000000000..4476055fe --- /dev/null +++ b/tests/driver/test_environ_register.py @@ -0,0 +1,71 @@ +# Copyright (c) 2026 The REvoDesign Developers. +# Distributed under the terms of the GNU General Public License v3.0. +# SPDX-License-Identifier: GPL-3.0-only + +from __future__ import annotations + +from REvoDesign.driver import environ_register + + +class FakeConfig: + def __init__(self): + self.reload_count = 0 + + def reload(self): + self.reload_count += 1 + + +class FakeBus: + _instance = object() + + def __init__(self, variables, override_existing=False): + self.cfg_group = {"environ": FakeConfig()} + self.variables = variables + self.override_existing = override_existing + + def get_value(self, cfg_item, converter=None, **_kwargs): + if cfg_item == "variables": + return dict(self.variables) + if cfg_item == "override_existing": + return self.override_existing + raise AssertionError(f"unexpected config item: {cfg_item}") + + +def _patch_bus(monkeypatch, bus): + class FakeConfigBus: + _instance = object() + + def __new__(cls): + return bus + + monkeypatch.setattr(environ_register, "ConfigBus", FakeConfigBus) + + +def test_register_environment_variables_preserves_existing_process_values(monkeypatch): + monkeypatch.setenv("REVODESIGN_TEST_EXISTING", "from-shell") + monkeypatch.delenv("REVODESIGN_TEST_NEW", raising=False) + bus = FakeBus( + { + "REVODESIGN_TEST_EXISTING": "from-yaml", + "REVODESIGN_TEST_NEW": "new-value", + "REVODESIGN_TEST_SKIPPED": None, + } + ) + _patch_bus(monkeypatch, bus) + + environ_register.register_environment_variables() + + assert bus.cfg_group["environ"].reload_count == 1 + assert environ_register.os.environ["REVODESIGN_TEST_EXISTING"] == "from-shell" + assert environ_register.os.environ["REVODESIGN_TEST_NEW"] == "new-value" + assert "REVODESIGN_TEST_SKIPPED" not in environ_register.os.environ + + +def test_register_environment_variables_can_explicitly_override(monkeypatch): + monkeypatch.setenv("REVODESIGN_TEST_EXISTING", "from-shell") + bus = FakeBus({"REVODESIGN_TEST_EXISTING": "from-yaml"}, override_existing=True) + _patch_bus(monkeypatch, bus) + + environ_register.register_environment_variables() + + assert environ_register.os.environ["REVODESIGN_TEST_EXISTING"] == "from-yaml" diff --git a/tests/driver/test_file_dialog.py b/tests/driver/test_file_dialog.py new file mode 100644 index 000000000..48fa45973 --- /dev/null +++ b/tests/driver/test_file_dialog.py @@ -0,0 +1,91 @@ +# Copyright (c) 2026 The REvoDesign Developers. +# Distributed under the terms of the GNU General Public License v3.0. +# SPDX-License-Identifier: GPL-3.0-only + +from __future__ import annotations + +import zipfile +from pathlib import Path + +import pytest + +from REvoDesign.driver.file_dialog import FileDialog, flatten_compressed_files + + +def _write_zip(path: Path, member_name: str = "payload.txt", content: str = "payload") -> None: + with zipfile.ZipFile(path, "w") as zip_file: + zip_file.writestr(member_name, content) + + +def test_flatten_compressed_files_extracts_to_temporary_directory(tmp_path): + archive = tmp_path / "input.zip" + _write_zip(archive) + + extracted_dir = Path(flatten_compressed_files(str(archive), target_dir=str(tmp_path))) + + assert extracted_dir.is_dir() + assert extracted_dir.parent == tmp_path + assert extracted_dir.name.startswith("revodesign-input.zip-") + assert not (tmp_path / "expanded_compressed_files").exists() + assert (extracted_dir / "payload.txt").read_text() == "payload" + + +def test_flatten_compressed_files_rejects_oversized_archive(tmp_path): + archive = tmp_path / "input.zip" + _write_zip(archive) + + with pytest.raises(ValueError, match="Archive is too large"): + flatten_compressed_files(str(archive), target_dir=str(tmp_path), max_archive_bytes=1) + + +def test_flatten_compressed_files_rejects_large_extracted_payload(tmp_path): + archive = tmp_path / "compressed.zip" + payload = "A" * 4096 + with zipfile.ZipFile(archive, "w", compression=zipfile.ZIP_DEFLATED) as zip_file: + zip_file.writestr("payload.txt", payload) + + assert archive.stat().st_size < len(payload) + with pytest.raises(ValueError, match="Archive expands too large"): + flatten_compressed_files( + str(archive), + target_dir=str(tmp_path), + max_archive_bytes=archive.stat().st_size + 1, + ) + + +def test_file_dialog_cleanup_archive_dirs_removes_temporary_extractions(tmp_path): + temp_dir = tmp_path / "revodesign-input.zip-test" + temp_dir.mkdir() + (temp_dir / "payload.txt").write_text("payload") + + dialog = object.__new__(FileDialog) + dialog._archive_temp_dirs = {str(temp_dir)} + + dialog.cleanup_archive_dirs() + + assert not temp_dir.exists() + assert dialog._archive_temp_dirs == set() + + +def test_browse_filename_reopens_archive_temp_directory(monkeypatch, tmp_path): + archive = tmp_path / "input.zip" + _write_zip(archive) + extracted_dir = tmp_path / "revodesign-input.zip-test" + picked_file = extracted_dir / "payload.txt" + dialogs = [] + + def mock_get_open_file_name(_window, _title, directory="", **_kwargs): + dialogs.append(directory) + return str(archive) if len(dialogs) == 1 else str(picked_file) + + monkeypatch.setattr("REvoDesign.driver.file_dialog.decide", lambda **_kwargs: True) + monkeypatch.setattr("REvoDesign.driver.file_dialog.flatten_compressed_files", lambda _path: str(extracted_dir)) + monkeypatch.setattr("REvoDesign.driver.file_dialog.getOpenFileNameWithExt", mock_get_open_file_name) + + dialog = object.__new__(FileDialog) + dialog.window = None + dialog._archive_temp_dirs = set() + + assert dialog.browse_filename() == str(picked_file) + assert dialogs == ["", str(extracted_dir)] + assert dialog._archive_temp_dirs == {str(extracted_dir)} diff --git a/tests/editor/test_editor_server.py b/tests/editor/test_editor_server.py index f367439a2..273107949 100644 --- a/tests/editor/test_editor_server.py +++ b/tests/editor/test_editor_server.py @@ -86,6 +86,20 @@ def initialize_server(): # Test cases +def test_initialize_token_does_not_log_secret(monkeypatch, caplog): + from REvoDesign.editor.monaco import server + + monkeypatch.setattr(server.secrets, "token_urlsafe", lambda _size: "sensitive-editor-token") + + with caplog.at_level("INFO", logger=server.logging.name): + token = server.initialize_token() + + assert token == "sensitive-editor-token" + assert ConfigStore().get("editor.token") == "sensitive-editor-token" + assert "sensitive-editor-token" not in caplog.text + assert "Generated editor authentication token." in caplog.text + + def test_favicon(test_client): """Test that the favicon endpoint returns the correct file.""" response = test_client.get("/favicon.svg") diff --git a/tests/editor/test_monaco.py b/tests/editor/test_monaco.py index 2487dccc5..bc82ab2f7 100644 --- a/tests/editor/test_monaco.py +++ b/tests/editor/test_monaco.py @@ -11,7 +11,7 @@ import pytest -from REvoDesign.editor.monaco.monaco import MonacoEditorManager, edit_file_with_monaco +from REvoDesign.editor.monaco.monaco import MonacoEditorManager, _npm_integrity_to_pooch_hash, edit_file_with_monaco @pytest.fixture @@ -73,6 +73,10 @@ def test_copy_html_template(test_tmp_dir): mock_copy.assert_called_once_with(manager.html_template_path, os.path.join(version_dir, "index.html")) +def test_npm_integrity_to_pooch_hash(): + assert _npm_integrity_to_pooch_hash("sha512-YWJjMTIz") == "sha512:616263313233" + + def test_edit_file_with_monaco(mock_server_control, mock_config_store, test_tmp_dir, test_worker): file_path = os.path.join(test_tmp_dir, "file.txt") diff --git a/tests/menu/shortcuts/test_mutation_effect_predictor.py b/tests/menu/shortcuts/test_mutation_effect_predictor.py index ba98f8b7b..f9cd8e20e 100644 --- a/tests/menu/shortcuts/test_mutation_effect_predictor.py +++ b/tests/menu/shortcuts/test_mutation_effect_predictor.py @@ -3,7 +3,9 @@ # SPDX-License-Identifier: GPL-3.0-only +import pandas as pd import pytest +from RosettaPy.common.mutation import RosettaPyProteinSequence from REvoDesign.shortcuts.tools.mutation_effect_predictors import ThermoMpnnPredictor, shortcut_thermompnn from tests.conftest import TestWorker @@ -23,13 +25,23 @@ ], ) @pytest.mark.skipif(not ThermoMpnnPredictor.installed, reason="ThermoMpnnPredictor not installed") -def test_shortcut_thermompnn(job_id, mode, threshold, long_dist, ss_penalty, test_worker: TestWorker): +def test_shortcut_thermompnn(job_id, mode, threshold, long_dist, ss_penalty, test_worker: TestWorker, monkeypatch): pdb = "../tests/data/6zcy_lig.pdb" test_worker.test_id = test_worker.method_name() test_worker.load_session_and_check(customized_session=pdb) save_dir = "predictors/thermompnn" + def _mock_init(self, pdb, save_dir, prefix, chains, mode, batch_size, threshold, distance, ss_penalty, device): + del save_dir, chains, batch_size, threshold, distance, ss_penalty, device + self.prefix = prefix + self.mode = mode + self.sequence = RosettaPyProteinSequence.from_pdb(pdb) + + monkeypatch.setattr(ThermoMpnnPredictor, "__init__", _mock_init) + monkeypatch.setattr(ThermoMpnnPredictor, "run", lambda self: pd.DataFrame({"ddG": [-1.0], "Mutation": ["MA1A"]})) + monkeypatch.setattr(ThermoMpnnPredictor, "cleanup", lambda self: None) + shortcut_thermompnn( pdb=pdb, save_dir=save_dir, diff --git a/tests/shortcuts/test_function_utils.py b/tests/shortcuts/test_function_utils.py new file mode 100644 index 000000000..592a58ae5 --- /dev/null +++ b/tests/shortcuts/test_function_utils.py @@ -0,0 +1,99 @@ +# Copyright (c) 2026 The REvoDesign Developers. +# Distributed under the terms of the GNU General Public License v3.0. +# SPDX-License-Identifier: GPL-3.0-only + +from __future__ import annotations + +import subprocess + +from REvoDesign.shortcuts import function_utils + + +class DummyProcess: + def __init__(self, returncode: int | None = None): + self.returncode = returncode + self.terminated = False + self.killed = False + self.wait_calls: list[float | None] = [] + + def poll(self): + return self.returncode + + def terminate(self): + self.terminated = True + self.returncode = 0 + + def kill(self): + self.killed = True + self.returncode = 0 + + def wait(self, timeout=None): + self.wait_calls.append(timeout) + return self.returncode + + +class HangingProcess(DummyProcess): + def wait(self, timeout=None): + self.wait_calls.append(timeout) + if not self.killed: + raise subprocess.TimeoutExpired(cmd="pymol", timeout=timeout) + return 0 + + +def test_visualize_conformer_sdf_tracks_preview_process(monkeypatch, tmp_path): + sdf_file = tmp_path / "ligand.sdf" + sdf_file.write_text("mock", encoding="utf-8") + process = DummyProcess() + calls = [] + + monkeypatch.setattr(function_utils.shutil, "which", lambda name: f"/usr/bin/{name}") + + def fake_popen(*args, **kwargs): + calls.append((args, kwargs)) + return process + + monkeypatch.setattr(function_utils.subprocess, "Popen", fake_popen) + function_utils._PREVIEW_PROCESSES.clear() + + function_utils.visualize_conformer_sdf(str(sdf_file), "New Window") + + assert process in function_utils._PREVIEW_PROCESSES + assert calls[0][0][0] == [ + "/usr/bin/pymol", + "-xi", + str(tmp_path / "ligand_load_to_preview.pml"), + ] + assert calls[0][1]["stdin"] is subprocess.DEVNULL + + function_utils.cleanup_conformer_preview_processes() + + +def test_visualize_conformer_sdf_prunes_finished_preview_processes(monkeypatch, tmp_path): + sdf_file = tmp_path / "ligand.sdf" + sdf_file.write_text("mock", encoding="utf-8") + finished = DummyProcess(returncode=0) + running = DummyProcess() + + monkeypatch.setattr(function_utils.shutil, "which", lambda name: f"/usr/bin/{name}") + monkeypatch.setattr(function_utils.subprocess, "Popen", lambda *args, **kwargs: running) + function_utils._PREVIEW_PROCESSES.clear() + function_utils._PREVIEW_PROCESSES.add(finished) + + function_utils.visualize_conformer_sdf(str(sdf_file), "New Window") + + assert finished not in function_utils._PREVIEW_PROCESSES + assert running in function_utils._PREVIEW_PROCESSES + + function_utils.cleanup_conformer_preview_processes() + + +def test_cleanup_conformer_preview_processes_kills_stubborn_process(): + process = HangingProcess() + function_utils._PREVIEW_PROCESSES.clear() + function_utils._PREVIEW_PROCESSES.add(process) + + function_utils.cleanup_conformer_preview_processes(timeout=0.01) + + assert process.terminated is True + assert process.killed is True + assert process not in function_utils._PREVIEW_PROCESSES diff --git a/tests/shortcuts/test_rfdiffusion_tasks.py b/tests/shortcuts/test_rfdiffusion_tasks.py new file mode 100644 index 000000000..214a668a9 --- /dev/null +++ b/tests/shortcuts/test_rfdiffusion_tasks.py @@ -0,0 +1,83 @@ +# Copyright (c) 2026 The REvoDesign Developers. +# Distributed under the terms of the GNU General Public License v3.0. +# SPDX-License-Identifier: GPL-3.0-only + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest + +from REvoDesign import issues +from REvoDesign.shortcuts.tools import rfdiffusion_tasks + + +def test_dgl_solver_does_not_run_pip_without_approval(monkeypatch): + solver = rfdiffusion_tasks.DglSolver(installed=False) + monkeypatch.setattr(solver, "fetch_cuda_version_before_install", lambda: None) + monkeypatch.setattr(rfdiffusion_tasks, "_has_qapplication", lambda: True) + monkeypatch.delenv(rfdiffusion_tasks.DGL_INSTALL_APPROVAL_ENV, raising=False) + monkeypatch.setattr(rfdiffusion_tasks, "decide", lambda **_kwargs: False) + + def fail_run_command(_cmd): + raise AssertionError("pip install should not run without approval") + + monkeypatch.setattr(rfdiffusion_tasks, "run_command", fail_run_command) + + assert solver.install() is False + assert solver.installed is False + + +def test_dgl_solver_declines_without_qapplication(monkeypatch): + solver = rfdiffusion_tasks.DglSolver(installed=False) + monkeypatch.setattr(solver, "fetch_cuda_version_before_install", lambda: None) + monkeypatch.setattr(rfdiffusion_tasks, "_has_qapplication", lambda: False) + monkeypatch.delenv(rfdiffusion_tasks.DGL_INSTALL_APPROVAL_ENV, raising=False) + + def fail_decide(**_kwargs): + raise AssertionError("headless DGL install approval should not open a dialog") + + def fail_run_command(_cmd): + raise AssertionError("pip install should not run without approval") + + monkeypatch.setattr(rfdiffusion_tasks, "decide", fail_decide) + monkeypatch.setattr(rfdiffusion_tasks, "run_command", fail_run_command) + + assert solver.install() is False + assert solver.installed is False + + +def test_dgl_solver_env_approval_runs_pip(monkeypatch): + solver = rfdiffusion_tasks.DglSolver(installed=False) + commands = [] + monkeypatch.setattr(solver, "fetch_cuda_version_before_install", lambda: None) + monkeypatch.setenv(rfdiffusion_tasks.DGL_INSTALL_APPROVAL_ENV, "1") + monkeypatch.setattr(rfdiffusion_tasks, "decide", lambda **_kwargs: False) + monkeypatch.setattr( + rfdiffusion_tasks, + "run_command", + lambda cmd: commands.append(cmd) or SimpleNamespace(returncode=0, stderr=""), + ) + + assert solver.install() is True + assert solver.installed is True + assert commands == [ + [ + rfdiffusion_tasks.sys.executable, + "-m", + "pip", + "install", + "dgl==2.2.1", + "-f", + "https://data.dgl.ai/wheels/repo.html", + ] + ] + + +def test_ensure_dgl_raises_when_installation_is_declined(monkeypatch): + solver = rfdiffusion_tasks.DglSolver(installed=False) + monkeypatch.setattr(rfdiffusion_tasks, "DglSolver", lambda: solver) + monkeypatch.setattr(solver, "install", lambda: False) + + with pytest.raises(issues.MissingExternalToolError, match="DGL installation was cancelled"): + rfdiffusion_tasks.RfDiffusion.ensure_dgl() diff --git a/tests/test_config_isolation.py b/tests/test_config_isolation.py new file mode 100644 index 000000000..533d25217 --- /dev/null +++ b/tests/test_config_isolation.py @@ -0,0 +1,15 @@ +# Copyright (c) 2026 The REvoDesign Developers. +# Distributed under the terms of the GNU General Public License v3.0. +# SPDX-License-Identifier: GPL-3.0-only + +from __future__ import annotations + +from pathlib import Path + +from REvoDesign.bootstrap import REVODESIGN_CONFIG_FILE + +from .conftest import EXPECTED_MAIN_CONFIG_FILE + + +def test_import_time_config_bootstrap_uses_test_data_dir(): + assert Path(REVODESIGN_CONFIG_FILE).resolve() == Path(EXPECTED_MAIN_CONFIG_FILE).resolve() diff --git a/tests/test_revodesign_workdir.py b/tests/test_revodesign_workdir.py new file mode 100644 index 000000000..a395c7dc3 --- /dev/null +++ b/tests/test_revodesign_workdir.py @@ -0,0 +1,34 @@ +# Copyright (c) 2026 The REvoDesign Developers. +# Distributed under the terms of the GNU General Public License v3.0. +# SPDX-License-Identifier: GPL-3.0-only + +from __future__ import annotations + +from pathlib import Path + + +class RecordingBus: + def __init__(self): + self.values = {} + + def set_value(self, key, value): + self.values[key] = value + + +def test_set_working_directory_does_not_mutate_process_cwd(app, monkeypatch, tmp_path): + from REvoDesign.REvoDesign import REvoDesignPlugin + + original_cwd = tmp_path / "original" + target_dir = tmp_path / "target" + original_cwd.mkdir() + target_dir.mkdir() + monkeypatch.chdir(original_cwd) + + plugin = REvoDesignPlugin() + plugin.bus = RecordingBus() + + plugin.set_working_directory(str(target_dir)) + + assert plugin.PWD == str(target_dir) + assert plugin.bus.values["work_dir"] == str(target_dir) + assert Path.cwd() == original_cwd.resolve() diff --git a/tests/tools/customized_widgets/test_shortcuts_util.py b/tests/tools/customized_widgets/test_shortcuts_util.py index 34ffa96d7..c84804653 100644 --- a/tests/tools/customized_widgets/test_shortcuts_util.py +++ b/tests/tools/customized_widgets/test_shortcuts_util.py @@ -9,7 +9,7 @@ from REvoDesign import issues from REvoDesign.basic.extensions import resolve_extension from REvoDesign.common.file_extensions import ExtColl -from REvoDesign.shortcuts.tools.esm2 import list_all_esm_variant_predict_model_names +from REvoDesign.shortcuts.tools.esm2 import ESM1V_WEIGHTS, list_all_esm_variant_predict_model_names from REvoDesign.shortcuts.utils import REGISTRY_DIR, _build_asked_value, resolve_choice_from, resolve_default_value from REvoDesign.tools.utils import resolve_dotted_function @@ -78,6 +78,13 @@ def test_resolve_choice_from(input_str, expected_type, expected_val): assert result == expected_val +def test_esm_model_registry_uses_hashes(): + for model_file in list_all_esm_variant_predict_model_names(): + registry_hash = ESM1V_WEIGHTS.registry[model_file] + assert registry_hash.startswith("md5:") + assert len(registry_hash) == len("md5:") + 32 + + @pytest.mark.parametrize( "input_str, error_type", [ diff --git a/tests/tools/test_file_downloader.py b/tests/tools/test_file_downloader.py index 5afb441a0..08f9c3146 100644 --- a/tests/tools/test_file_downloader.py +++ b/tests/tools/test_file_downloader.py @@ -35,7 +35,6 @@ def test_complete_varify_string(self, a_string, hash_type, expected): {"file1.txt": "md5:abc123", "file2.txt": "def456"}, {"file1.txt": "md5:abc123", "file2.txt": "md5:def456"}, ), - ({"file1.txt": None, "file2.txt": ""}, {"file1.txt": None, "file2.txt": None}), ({"file1.txt": "sha256:abc123"}, {"file1.txt": "sha256:abc123"}), ], ) @@ -44,6 +43,24 @@ def test_preprocess_registry(self, registry, expected): result = FileDownloadRegistry.preprocess_registry(registry) assert result == expected + def test_preprocess_registry_rejects_missing_hashes(self): + """Test preprocess_registry rejects entries that would disable verification.""" + with pytest.raises(ValueError, match="file1.txt, file2.txt"): + FileDownloadRegistry.preprocess_registry({"file1.txt": None, "file2.txt": ""}) + + @patch("REvoDesign.tools.download_registry.pooch.create") + def test_init_rejects_registry_entries_without_hashes(self, mock_pooch_create): + """Test unsigned registry entries fail before Pooch can download them.""" + with pytest.raises(ValueError, match="unsafe.pkl"): + FileDownloadRegistry( + name="test_module", + base_url="http://example.com", + registry={"safe.txt": "md5:abc123", "unsafe.pkl": None}, + customized_directory="/custom/path", + ) + + mock_pooch_create.assert_not_called() + @pytest.mark.parametrize( "md5_contents, expected", [ diff --git a/tests/tools/test_package_manager.py b/tests/tools/test_package_manager.py index 6cae257f0..33167c924 100644 --- a/tests/tools/test_package_manager.py +++ b/tests/tools/test_package_manager.py @@ -3,6 +3,7 @@ # SPDX-License-Identifier: GPL-3.0-only +import hmac import io import json import os @@ -22,13 +23,22 @@ from REvoDesign.tools.package_manager import ( GitSolver, PIPInstaller, + REvoDesignPackageManager, + _compute_hmac, _python_version_matches, + bootstrap_extras_json, + bootstrap_manager_ui_file, + fetch_bootstrap_manifest, fetch_gist_file, fetch_gist_json, + fetch_verified_bootstrap_file, filter_sensitive_data, get_github_repo_tags, + load_bootstrap_extras_json, + package_manager_bootstrap_dir, run_command, solve_installation_config, + verify_manifest, ) @@ -65,6 +75,22 @@ def test_pm_fetch_gist_file_url_error(): fetch_gist_file(mock_url, "temp_file.ui") +def test_pm_fetch_gist_file_uses_timeout(monkeypatch, tmp_path): + calls = [] + + def mock_read(url, **kwargs): + calls.append((url, kwargs)) + return b"mock UI content" + + output = tmp_path / "manager.ui" + monkeypatch.setattr(package_manager, "_read_https_url", mock_read) + + fetch_gist_file("https://example.com/file.ui", str(output), timeout=3) + + assert output.read_text() == "mock UI content" + assert calls == [("https://example.com/file.ui", {"timeout": 3})] + + # Test for fetch_gist_json @@ -101,6 +127,121 @@ def test_pm_fetch_gist_json_error(): assert result == {} +def test_pm_load_bootstrap_extras_json(tmp_path): + registry = tmp_path / "REvoDesignExtrasTableRich.json" + registry.write_text(json.dumps({"entities": []})) + + assert load_bootstrap_extras_json(registry) == {"entities": []} + + +def test_pm_bootstrap_paths_use_env(monkeypatch, tmp_path): + monkeypatch.setenv(package_manager.PACKAGE_MANAGER_BOOTSTRAP_ENV, str(tmp_path)) + + assert package_manager_bootstrap_dir() == tmp_path + assert bootstrap_manager_ui_file() == tmp_path / "UI" / "REvoDesign_installer.ui" + assert bootstrap_extras_json() == tmp_path / "REvoDesignExtrasTableRich.json" + + +def test_pm_ensure_ui_file_refetches_even_when_previous_bootstrap_exists(monkeypatch, tmp_path): + monkeypatch.setenv(package_manager.PACKAGE_MANAGER_BOOTSTRAP_ENV, str(tmp_path)) + ui_file = tmp_path / "UI" / "REvoDesign_installer.ui" + ui_file.parent.mkdir() + ui_file.write_text("") + fresh_ui = b"" + manifest = { + package_manager.MANIFEST_ASSET_MANAGER_UI: hmac.new( + package_manager._MANAGER_HMAC_KEY, fresh_ui, "sha256" + ).hexdigest() + } + + def mock_fetch(ui_file_url, save_to_file, **kwargs): + with open(save_to_file, "wb") as fh: + fh.write(fresh_ui) + + monkeypatch.setattr(package_manager, "fetch_bootstrap_manifest", lambda: manifest) + monkeypatch.setattr(package_manager, "fetch_gist_file", mock_fetch) + + plugin = REvoDesignPackageManager() + + assert plugin.ensure_ui_file() == str(ui_file) + assert ui_file.read_text() == "" + + +def test_pm_ensure_ui_file_retries_then_writes_bootstrap_file(monkeypatch, tmp_path): + monkeypatch.setenv(package_manager.PACKAGE_MANAGER_BOOTSTRAP_ENV, str(tmp_path)) + ui_file = tmp_path / "UI" / "REvoDesign_installer.ui" + fetched_ui = b"" + manifest = { + package_manager.MANIFEST_ASSET_MANAGER_UI: hmac.new( + package_manager._MANAGER_HMAC_KEY, fetched_ui, "sha256" + ).hexdigest() + } + + fetch_calls = [] + + def mock_fetch(ui_file_url, save_to_file, **kwargs): + fetch_calls.append(save_to_file) + if len(fetch_calls) == 1: + raise urllib.error.URLError("temporary outage") + with open(save_to_file, "wb") as fh: + fh.write(fetched_ui) + + monkeypatch.setattr(package_manager, "fetch_bootstrap_manifest", lambda: manifest) + monkeypatch.setattr(package_manager, "fetch_gist_file", mock_fetch) + monkeypatch.setattr(package_manager.time, "sleep", lambda _seconds: None) + + plugin = REvoDesignPackageManager() + + result = plugin.ensure_ui_file() + assert result == str(ui_file) + assert len(fetch_calls) == 2 + assert open(ui_file).read() == "" + + +def test_pm_load_bootstrap_json_fetches_and_writes_bootstrap_file(monkeypatch, tmp_path, qtbot): + monkeypatch.setenv(package_manager.PACKAGE_MANAGER_BOOTSTRAP_ENV, str(tmp_path)) + plugin = REvoDesignPackageManager() + list_view = QtWidgets.QListView() + qtbot.addWidget(list_view) + plugin.installer_ui = SimpleNamespace(listView_extras=list_view) + remote_data = { + "entities": [ + { + "name": "Fetched Extras", + "description": "Remote extras registry", + "extras": [], + } + ], + } + remote_bytes = json.dumps(remote_data).encode() + manifest = { + package_manager.MANIFEST_ASSET_EXTRAS_JSON: hmac.new( + package_manager._MANAGER_HMAC_KEY, remote_bytes, "sha256" + ).hexdigest() + } + fetch_calls = [] + + def mock_fetch(ui_file_url, save_to_file, **kwargs): + fetch_calls.append(save_to_file) + with open(save_to_file, "wb") as fh: + fh.write(b'{"entities": []}' if len(fetch_calls) == 1 else remote_bytes) + + monkeypatch.setattr( + package_manager, + "run_worker_thread_in_pool", + lambda worker_function, **_kwargs: worker_function(), + ) + monkeypatch.setattr(package_manager, "fetch_bootstrap_manifest", lambda: manifest) + monkeypatch.setattr(package_manager, "fetch_gist_file", mock_fetch) + monkeypatch.setattr(package_manager.time, "sleep", lambda _seconds: None) + + plugin.load_bootstrap_json() + + assert len(fetch_calls) == 2 + assert plugin.remote_extra_group_data.entities[0].name == "Fetched Extras" + assert json.loads(bootstrap_extras_json().read_text()) == remote_data + + def test_pm_run_command_success(): """ Test run_command for a successful command execution. @@ -449,30 +590,33 @@ def test_pm_ensure_package(pip_installer, mocker): class TestGetGithubRepoTags: - def test_pm_valid_repo_url(self): - # Test a valid repository URL, expecting a list of tags + def test_pm_valid_repo_url(self, monkeypatch): + package_manager._GITHUB_TAG_CACHE.clear() repo_url = "https://github.com/BradyAJohnston/MolecularNodes" + monkeypatch.setattr( + package_manager, + "_read_https_url", + lambda url, **_kwargs: json.dumps([{"name": "v1.0.0"}, {"name": "v1.1.0"}]).encode(), + ) + tags = get_github_repo_tags(repo_url) - if not tags: - pytest.skip("Live GitHub tags are unavailable in the current test environment.") - assert isinstance(tags, list) - assert len(tags) > 0 - for tag in tags: - assert isinstance(tag, str) + + assert tags == ["v1.0.0", "v1.1.0"] def test_pm_invalid_repo_url(self): - # Test an invalid repository URL, expecting an empty list - repo_url = "https://github.com/nonexistent/repo" + package_manager._GITHUB_TAG_CACHE.clear() + repo_url = "https://example.com/nonexistent/repo" tags = get_github_repo_tags(repo_url) assert isinstance(tags, list) assert len(tags) == 0 def test_pm_http_error(self, monkeypatch): - # Test handling of HTTPError - def mock_urlopen(*args, **kwargs): - raise HTTPError(args[0], 404, "Not Found", None, None) + package_manager._GITHUB_TAG_CACHE.clear() + + def mock_read(url, **_kwargs): + raise HTTPError(url, 404, "Not Found", None, None) - monkeypatch.setattr(urllib.request, "urlopen", mock_urlopen) + monkeypatch.setattr(package_manager, "_read_https_url", mock_read) repo_url = "https://github.com/BradyAJohnston/MolecularNodes" tags = get_github_repo_tags(repo_url) @@ -480,11 +624,12 @@ def mock_urlopen(*args, **kwargs): assert len(tags) == 0 def test_pm_url_error(self, monkeypatch): - # Test handling of URLError - def mock_urlopen(*args, **kwargs): + package_manager._GITHUB_TAG_CACHE.clear() + + def mock_read(*_args, **_kwargs): raise URLError("Failed to reach the server") - monkeypatch.setattr(urllib.request, "urlopen", mock_urlopen) + monkeypatch.setattr(package_manager, "_read_https_url", mock_read) repo_url = "https://github.com/BradyAJohnston/MolecularNodes" tags = get_github_repo_tags(repo_url) @@ -492,32 +637,52 @@ def mock_urlopen(*args, **kwargs): assert len(tags) == 0 -def test_pm_get_github_repo_tags_success(): - """Test successful retrieval of tags.""" +def test_pm_get_github_repo_tags_uses_timeout_auth_and_cache(monkeypatch): + calls = [] - result = get_github_repo_tags("https://github.com/BradyAJohnston/MolecularNodes") - if not result: - pytest.skip("Live GitHub tags are unavailable in the current test environment.") - assert result, "Github repository tags should not be empty" + def mock_read(url, **kwargs): + calls.append((url, kwargs)) + return json.dumps([{"name": "v2.0.0"}]).encode() + monkeypatch.setenv("GITHUB_TOKEN", "token-123") + monkeypatch.setattr(package_manager, "_read_https_url", mock_read) + package_manager._GITHUB_TAG_CACHE.clear() -def test_pm_get_github_repo_tags_http_error(): - """Test handling of an HTTPError.""" - with patch( - "urllib.request.urlopen", - side_effect=urllib.error.HTTPError( - url="https://api.github.com/repos/test_owner/test_repo/tags", code=404, msg="Not Found", hdrs=None, fp=None - ), - ): - result = get_github_repo_tags("https://github.com/test_owner/test_repo") - assert result == [] + result = get_github_repo_tags("https://github.com/test_owner/test_repo", timeout=3) + assert result == ["v2.0.0"] + assert calls == [ + ( + "https://api.github.com/repos/test_owner/test_repo/tags", + { + "timeout": 3, + "headers": { + "Accept": "application/vnd.github+json", + "Authorization": "Bearer token-123", + }, + }, + ) + ] + cached_at, cached_tags = package_manager._GITHUB_TAG_CACHE["https://api.github.com/repos/test_owner/test_repo/tags"] + assert cached_at > 0 + assert cached_tags == ["v2.0.0"] + + assert get_github_repo_tags("https://github.com/test_owner/test_repo", timeout=3) == ["v2.0.0"] + assert len(calls) == 1 + + +def test_pm_get_github_repo_tags_falls_back_to_cached_tags(monkeypatch): + api_url = "https://api.github.com/repos/test_owner/test_repo/tags" + package_manager._GITHUB_TAG_CACHE[api_url] = (0.0, ["cached-tag"]) + + def mock_read(url, **_kwargs): + raise urllib.error.URLError(reason="Network unreachable") -def test_pm_get_github_repo_tags_url_error(): - """Test handling of a URLError.""" - with patch("urllib.request.urlopen", side_effect=urllib.error.URLError(reason="Network unreachable")): - result = get_github_repo_tags("https://github.com/test_owner/test_repo") - assert result == [] + monkeypatch.setattr(package_manager, "_read_https_url", mock_read) + + result = get_github_repo_tags("https://github.com/test_owner/test_repo") + + assert result == ["cached-tag"] @pytest.mark.parametrize( @@ -732,6 +897,22 @@ def test_pm_thread_manage_kill_entry(qtbot, monkeypatch): button = QtWidgets.QPushButton("Killable") qtbot.addWidget(button) + class FakeProcess: + terminated = False + killed = False + + def poll(self): + return None + + def terminate(self): + self.terminated = True + + def wait(self, timeout=None): + return 0 + + def kill(self): + self.killed = True + def cancellable_task(): thread = package_manager.QtCore.QThread.currentThread() start = time.monotonic() @@ -743,6 +924,9 @@ def cancellable_task(): worker_thread = package_manager.WorkerThread(cancellable_task) manager = package_manager.ThreadExecutionManager(worker_thread, description="Kill", trigger_buttons=button) + fake_process = FakeProcess() + package_manager.RunningProcessRegistry._worker_processes[id(worker_thread)] = {fake_process} + package_manager.RunningProcessRegistry._process_worker[id(fake_process)] = id(worker_thread) assert manager.abort_overlays mock_cmd = SimpleNamespace(interrupt=MagicMock()) @@ -760,5 +944,142 @@ def cancellable_task(): result = worker_thread.handle_result() assert manager.abort_overlays == [] - assert mock_cmd.interrupt.called + assert not mock_cmd.interrupt.called + assert fake_process.terminated + assert not fake_process.killed assert result is None or result[0] == "aborted" + + +def test_worker_thread_emits_and_records_errors(qtbot): + def failing_task(): + raise RuntimeError("worker exploded") + + worker_thread = package_manager.WorkerThread(failing_task) + errors = [] + notifications = [] + finished = [] + worker_thread.error_signal.connect(errors.append) + worker_thread.notify_signal.connect(notifications.append) + worker_thread.finished_signal.connect(lambda: finished.append(True)) + + worker_thread.start() + + qtbot.waitUntil(lambda: bool(finished), timeout=2000) + worker_thread.wait(2000) + + assert worker_thread.handle_result() is None + assert isinstance(worker_thread.handle_error(), RuntimeError) + assert len(errors) == 1 + assert isinstance(errors[0], RuntimeError) + assert notifications == ["RuntimeError: worker exploded"] + + +def test_run_worker_thread_in_pool_reraises_worker_errors(qtbot, monkeypatch): + monkeypatch.setattr(package_manager, "refresh_window", lambda: None) + + def failing_task(): + raise RuntimeError("worker exploded") + + with pytest.raises(RuntimeError, match="worker exploded"): + package_manager.run_worker_thread_in_pool(failing_task, notify_slot=lambda _message: None) + + +# ── HMAC manifest verification ──────────────────────────────────────────── + + +def test_pm_verify_manifest_all_match(monkeypatch, tmp_path): + key = os.urandom(32) + monkeypatch.setattr(package_manager, "_MANAGER_HMAC_KEY", key) + files = {} + manifest = {} + for name in ("a.py", "b.ui"): + path = tmp_path / name + path.write_bytes(os.urandom(64)) + files[name] = str(path) + manifest[name] = hmac.new(key, path.read_bytes(), "sha256").hexdigest() + + assert verify_manifest(files, manifest) + + +def test_pm_verify_manifest_mismatch(monkeypatch, tmp_path): + key = os.urandom(32) + monkeypatch.setattr(package_manager, "_MANAGER_HMAC_KEY", key) + path = tmp_path / "a.py" + path.write_bytes(b"original") + manifest = {"a.py": hmac.new(key, b"tampered", "sha256").hexdigest()} + + assert not verify_manifest({"a.py": str(path)}, manifest) + + +def test_pm_verify_manifest_missing_entry(monkeypatch, tmp_path): + key = os.urandom(32) + monkeypatch.setattr(package_manager, "_MANAGER_HMAC_KEY", key) + path = tmp_path / "a.py" + path.write_bytes(b"data") + + assert not verify_manifest({"a.py": str(path)}, {}) + + +def test_pm_fetch_bootstrap_manifest_rejects_invalid_schema(monkeypatch): + monkeypatch.setattr(package_manager, "fetch_required_gist_json", lambda _url: {"asset.py": 123}) + + with pytest.raises(ValueError, match="Bootstrap manifest must map asset names to HMAC strings"): + fetch_bootstrap_manifest() + + +def test_pm_fetch_verified_bootstrap_file_retries_on_hmac_mismatch(monkeypatch, tmp_path): + key = os.urandom(32) + monkeypatch.setattr(package_manager, "_MANAGER_HMAC_KEY", key) + monkeypatch.setattr(package_manager.time, "sleep", lambda _seconds: None) + output = tmp_path / "asset.ui" + expected_payload = b"expected asset" + manifest = {"asset.ui": hmac.new(key, expected_payload, "sha256").hexdigest()} + fetch_payloads = iter((b"tampered asset", expected_payload)) + fetch_calls = [] + + def mock_fetch(ui_file_url, save_to_file, **kwargs): + fetch_calls.append(ui_file_url) + with open(save_to_file, "wb") as file_handle: + file_handle.write(next(fetch_payloads)) + + monkeypatch.setattr(package_manager, "fetch_gist_file", mock_fetch) + + fetch_verified_bootstrap_file( + url="https://example.com/asset.ui", + asset_name="asset.ui", + save_to_file=str(output), + manifest=manifest, + description="test asset", + ) + + assert fetch_calls == ["https://example.com/asset.ui", "https://example.com/asset.ui"] + assert output.read_bytes() == expected_payload + + +def test_pm_compute_hmac_deterministic(monkeypatch, tmp_path): + key = os.urandom(32) + monkeypatch.setattr(package_manager, "_MANAGER_HMAC_KEY", key) + path = tmp_path / "a.py" + path.write_bytes(b"hello") + + assert _compute_hmac(str(path)) == _compute_hmac(str(path)) + + +# ── fetch_tags graceful degradation ─────────────────────────────────────── + + +def test_pm_fetch_tags_silent_on_network_error(monkeypatch, qtbot): + plugin = REvoDesignPackageManager() + list_view = QtWidgets.QListView() + qtbot.addWidget(list_view) + plugin.installer_ui = SimpleNamespace(listView_extras=list_view, comboBox_version=QtWidgets.QComboBox()) + monkeypatch.setattr( + package_manager, + "get_github_repo_tags", + MagicMock(side_effect=URLError("offline")), + ) + + # Must not raise, must not call notify_box + with patch("REvoDesign.tools.package_manager.notify_box") as mock_notify: + plugin.fetch_tags() + mock_notify.assert_not_called() diff --git a/tests/tools/test_rfdiffusion_tools.py b/tests/tools/test_rfdiffusion_tools.py new file mode 100644 index 000000000..4918f9be4 --- /dev/null +++ b/tests/tools/test_rfdiffusion_tools.py @@ -0,0 +1,46 @@ +# Copyright (c) 2026 The REvoDesign Developers. +# Distributed under the terms of the GNU General Public License v3.0. +# SPDX-License-Identifier: GPL-3.0-only + +from __future__ import annotations + +import importlib +import sys + +import matplotlib + + +def _reload_rfdiffusion_tools(): + sys.modules.pop("REvoDesign.tools.rfdiffusion_tools", None) + return importlib.import_module("REvoDesign.tools.rfdiffusion_tools") + + +def test_import_does_not_force_matplotlib_backend(monkeypatch): + calls = [] + monkeypatch.setattr(matplotlib, "use", lambda *args, **kwargs: calls.append((args, kwargs))) + + _reload_rfdiffusion_tools() + + assert calls == [] + + +def test_plotting_can_prefer_interactive_backend_lazily(monkeypatch): + module = _reload_rfdiffusion_tools() + calls = [] + monkeypatch.delenv("MPLBACKEND", raising=False) + monkeypatch.setattr(matplotlib, "use", lambda *args, **kwargs: calls.append((args, kwargs))) + + module._prefer_interactive_backend_for_plotting() + + assert calls == [((module.PREFERRED_INTERACTIVE_BACKEND,), {"force": False})] + + +def test_plotting_honors_existing_mplbackend(monkeypatch): + module = _reload_rfdiffusion_tools() + calls = [] + monkeypatch.setenv("MPLBACKEND", "Agg") + monkeypatch.setattr(matplotlib, "use", lambda *args, **kwargs: calls.append((args, kwargs))) + + module._prefer_interactive_backend_for_plotting() + + assert calls == [] diff --git a/tests/tools/test_ssl_certificates.py b/tests/tools/test_ssl_certificates.py new file mode 100644 index 000000000..cf0a69abd --- /dev/null +++ b/tests/tools/test_ssl_certificates.py @@ -0,0 +1,74 @@ +# Copyright (c) 2026 The REvoDesign Developers. +# Distributed under the terms of the GNU General Public License v3.0. +# SPDX-License-Identifier: GPL-3.0-only + +from __future__ import annotations + +import datetime +import stat + +import pytest +from OpenSSL import crypto + +import REvoDesign +from REvoDesign.tools.ssl_certificates import ( + CERT_VALIDITY_ENV, + CERTIFICATE_MODE, + PRIVATE_DIR_MODE, + PRIVATE_KEY_MODE, + SSLCertificateManager, +) + + +def _mode(path): + return stat.S_IMODE(path.stat().st_mode) + + +def _load_cert(path): + return crypto.load_certificate(crypto.FILETYPE_PEM, path.read_bytes()) + + +def test_new_certificate_uses_random_serial_and_private_permissions(monkeypatch, tmp_path): + monkeypatch.setattr(REvoDesign, "set_cache_dir", lambda: str(tmp_path)) + + manager = SSLCertificateManager(role="server") + manager.create_new_certificate() + first_serial = _load_cert(tmp_path / "crts" / "server.crt").get_serial_number() + + manager.create_new_certificate() + second_serial = _load_cert(tmp_path / "crts" / "server.crt").get_serial_number() + + assert first_serial != 1000 + assert second_serial != 1000 + assert first_serial != second_serial + assert _mode(tmp_path / "crts") == PRIVATE_DIR_MODE + assert _mode(tmp_path / "crts" / "server.key") == PRIVATE_KEY_MODE + assert _mode(tmp_path / "crts" / "server.crt") == CERTIFICATE_MODE + + +def test_certificate_validity_can_be_configured(monkeypatch, tmp_path): + monkeypatch.setattr(REvoDesign, "set_cache_dir", lambda: str(tmp_path)) + monkeypatch.setenv(CERT_VALIDITY_ENV, "3") + + manager = SSLCertificateManager(role="server") + manager.create_new_certificate() + + cert = _load_cert(tmp_path / "crts" / "server.crt") + expires_at = datetime.datetime.strptime( + cert.get_notAfter().decode("utf-8"), + "%Y%m%d%H%M%SZ", + ).replace(tzinfo=datetime.timezone.utc) + now = datetime.datetime.now(datetime.timezone.utc) + + assert datetime.timedelta(days=2, hours=23) < expires_at - now <= datetime.timedelta(days=3, minutes=1) + + +@pytest.mark.parametrize("value", ["0", "-1", "invalid"]) +def test_certificate_validity_rejects_invalid_values(monkeypatch, tmp_path, value): + monkeypatch.setattr(REvoDesign, "set_cache_dir", lambda: str(tmp_path)) + monkeypatch.setenv(CERT_VALIDITY_ENV, value) + + manager = SSLCertificateManager(role="server") + + with pytest.raises(ValueError, match=CERT_VALIDITY_ENV): + manager.create_new_certificate() diff --git a/tests/tools/test_utils.py b/tests/tools/test_utils.py index b61e5f70d..a55242c1c 100644 --- a/tests/tools/test_utils.py +++ b/tests/tools/test_utils.py @@ -231,6 +231,16 @@ def test_extract_zip_rejects_path_traversal(temp_extract_dir): assert not os.path.exists(outside_path) +@pytest.mark.parametrize("member_name", ["/tmp/evil.txt", r"C:\tmp\evil.txt"]) +def test_extract_zip_rejects_absolute_member_paths(temp_extract_dir, member_name): + archive_path = os.path.join(temp_extract_dir, "unsafe_absolute.zip") + with zipfile.ZipFile(archive_path, "w") as zipf: + zipf.writestr(member_name, "malicious") + + with pytest.raises(ValueError, match="Unsafe archive member path"): + extract_archive(archive_path, temp_extract_dir) + + def test_extract_tar_rejects_path_traversal(temp_extract_dir): archive_path = os.path.join(temp_extract_dir, "unsafe.tar.gz") payload = b"malicious" @@ -246,6 +256,19 @@ def test_extract_tar_rejects_path_traversal(temp_extract_dir): assert not os.path.exists(outside_path) +@pytest.mark.parametrize("member_name", ["/tmp/evil.txt", r"C:\tmp\evil.txt"]) +def test_extract_tar_rejects_absolute_member_paths(temp_extract_dir, member_name): + archive_path = os.path.join(temp_extract_dir, "unsafe_absolute.tar.gz") + payload = b"malicious" + info = tarfile.TarInfo(name=member_name) + info.size = len(payload) + with tarfile.open(archive_path, "w:gz") as tar: + tar.addfile(info, io.BytesIO(payload)) + + with pytest.raises(ValueError, match="Unsafe archive member path"): + extract_archive(archive_path, temp_extract_dir) + + def test_get_color_uniform_range(): # Test when min_value equals max_value color = get_color("viridis", 0.5, 0.5, 0.5) @@ -439,6 +462,21 @@ def test_generate_strong_password_character_set(): assert set(password).issubset(allowed_characters) +def test_generate_strong_password_uses_cryptographic_rng(monkeypatch): + """Password generation must not fall back to the predictable random module.""" + choices = [] + + def fake_choice(characters): + choices.append(characters) + return "A" + + monkeypatch.setattr("REvoDesign.tools.utils.secrets.choice", fake_choice) + + assert generate_strong_password(16) == "A" * 16 + assert len(choices) == 16 + assert all("A" in characters for characters in choices) + + def test_generate_strong_password_value_error(): """Test that a ValueError is raised for invalid lengths.""" with pytest.raises(ValueError):