From 49894189f0c1c2ef80b722fca07961465262cb8c Mon Sep 17 00:00:00 2001 From: xraymemory Date: Thu, 2 Jul 2026 14:39:48 -0400 Subject: [PATCH] fix(imports,tests): support bare require_* decorators; clean up protpardelle temp dirs #285: the require_* decorator factories treated a bare @require_x as require_x(func), binding the function to the message arg so it was never wrapped. Detect a callable first arg and re-dispatch, so both @require_x and @require_x('msg') work. Applied to require_boltz/protenix/rf3/protpardelle/any_model. #286: the protpardelle test env-var setup used os.environ.setdefault(..., tempfile.mkdtemp(...)), which eagerly created (and leaked) a temp dir on every run even when the var was already set. Guard the mkdtemp on the var being unset and register atexit shutil.rmtree cleanup, in both conftest.py and test_protpardelle_wrapper.py. Fixes #285. Fixes #286. --- src/sampleworks/utils/imports.py | 25 +++++++++++++++++++ tests/models/protpardelle/conftest.py | 16 ++++++++---- .../protpardelle/test_protpardelle_wrapper.py | 16 ++++++++---- 3 files changed, 47 insertions(+), 10 deletions(-) diff --git a/src/sampleworks/utils/imports.py b/src/sampleworks/utils/imports.py index e646cacf..5d80cd48 100644 --- a/src/sampleworks/utils/imports.py +++ b/src/sampleworks/utils/imports.py @@ -76,6 +76,11 @@ def require_boltz(message: str | None = None) -> Callable[[F], F]: ... def custom_function(): ... pass """ + if callable(message): + # Bare ``@require_boltz`` usage: the decorated function arrives as ``message``. + # Re-dispatch so both ``@require_boltz`` and ``@require_boltz("msg")`` work. + return require_boltz()(message) + default_message = "Boltz model wrapper is not available. Install with: pixi install -e boltz" def decorator(func: F) -> F: @@ -119,6 +124,11 @@ def require_protenix(message: str | None = None) -> Callable[[F], F]: ... def custom_function(): ... pass """ + if callable(message): + # Bare ``@require_protenix`` usage: the decorated function arrives as ``message``. + # Re-dispatch so both ``@require_protenix`` and ``@require_protenix("msg")`` work. + return require_protenix()(message) + default_message = ( "Protenix model wrapper is not available. Install with: pixi install -e protenix" ) @@ -164,6 +174,11 @@ def require_rf3(message: str | None = None) -> Callable[[F], F]: ... def custom_function(): ... pass """ + if callable(message): + # Bare ``@require_rf3`` usage: the decorated function arrives as ``message``. + # Re-dispatch so both ``@require_rf3`` and ``@require_rf3("msg")`` work. + return require_rf3()(message) + default_message = "RF3 model wrapper is not available. Install with: pixi install -e rf3" def decorator(func: F) -> F: @@ -207,6 +222,11 @@ def require_protpardelle(message: str | None = None) -> Callable[[F], F]: ... def custom_function(): ... pass """ + if callable(message): + # Bare ``@require_protpardelle`` usage: the decorated function arrives as ``message``. + # Re-dispatch so both ``@require_protpardelle`` and ``@require_protpardelle("msg")`` work. + return require_protpardelle()(message) + default_message = ( "Protpardelle model wrapper is not available. Install with: pixi install -e protpardelle" ) @@ -252,6 +272,11 @@ def require_any_model(message: str | None = None) -> Callable[[F], F]: ... def custom_function(): ... pass """ + if callable(message): + # Bare ``@require_any_model`` usage: the decorated function arrives as ``message``. + # Re-dispatch so both ``@require_any_model`` and ``@require_any_model("msg")`` work. + return require_any_model()(message) + default_message = ( "No model wrappers are available. " "Please install at least one model wrapper with the appropriate feature group: " diff --git a/tests/models/protpardelle/conftest.py b/tests/models/protpardelle/conftest.py index 5b0ba8e8..6b6335f7 100644 --- a/tests/models/protpardelle/conftest.py +++ b/tests/models/protpardelle/conftest.py @@ -6,7 +6,9 @@ imported, then build a small randomly-initialized ``ai-allatom`` model. """ +import atexit import os +import shutil import tempfile import textwrap from pathlib import Path @@ -14,11 +16,15 @@ import pytest -# Must be set before any `import protpardelle...` happens. Respect an -# externally configured directory (e.g. when real weights are available). -os.environ.setdefault( - "PROTPARDELLE_MODEL_PARAMS", tempfile.mkdtemp(prefix="protpardelle_model_params_") -) +# Must be set before any `import protpardelle...` happens. Respect an externally +# configured directory (e.g. when real weights are available); otherwise create a +# throwaway dir and remove it at process exit so tests don't leak temp dirs. +# `setdefault` can't own the created dir (it eagerly evaluates mkdtemp even when the +# var is already set), so guard explicitly and register cleanup only when we created it. +if "PROTPARDELLE_MODEL_PARAMS" not in os.environ: + _model_params_dir = tempfile.mkdtemp(prefix="protpardelle_model_params_") + os.environ["PROTPARDELLE_MODEL_PARAMS"] = _model_params_dir + atexit.register(shutil.rmtree, _model_params_dir, ignore_errors=True) protpardelle_models = pytest.importorskip( "protpardelle.core.models", reason="Protpardelle not installed in this environment" diff --git a/tests/models/protpardelle/test_protpardelle_wrapper.py b/tests/models/protpardelle/test_protpardelle_wrapper.py index 8567d4fb..acef03cb 100644 --- a/tests/models/protpardelle/test_protpardelle_wrapper.py +++ b/tests/models/protpardelle/test_protpardelle_wrapper.py @@ -5,17 +5,23 @@ needing downloaded weights. """ +import atexit import os +import shutil import tempfile import pytest -# Ensure the model-params directory exists before protpardelle is imported, -# mirroring conftest (import order across conftests is not guaranteed). -os.environ.setdefault( - "PROTPARDELLE_MODEL_PARAMS", tempfile.mkdtemp(prefix="protpardelle_model_params_") -) +# Ensure the model-params directory exists before protpardelle is imported, mirroring +# conftest (import order across conftests is not guaranteed). Whichever module runs first +# creates the throwaway dir and registers its cleanup; the other sees the var already set +# and does nothing. Guard explicitly rather than `setdefault`, which would eagerly leak a +# mkdtemp even when the var is present. +if "PROTPARDELLE_MODEL_PARAMS" not in os.environ: + _model_params_dir = tempfile.mkdtemp(prefix="protpardelle_model_params_") + os.environ["PROTPARDELLE_MODEL_PARAMS"] = _model_params_dir + atexit.register(shutil.rmtree, _model_params_dir, ignore_errors=True) pytest.importorskip( "protpardelle.core.models", reason="Protpardelle not installed in this environment"