diff --git a/pyrit/registry/components/initializer_registry.py b/pyrit/registry/components/initializer_registry.py index 786c198670..02d3e0b1a0 100644 --- a/pyrit/registry/components/initializer_registry.py +++ b/pyrit/registry/components/initializer_registry.py @@ -285,7 +285,10 @@ def _load_module_from_path(*, file_path: Path, module_name: str) -> ModuleType: @staticmethod def _module_defined_initializers( - *, module: ModuleType, base_class: type[PyRITInitializer] + *, + module: ModuleType, + base_class: type[PyRITInitializer], + strict: bool = False, ) -> list[type[PyRITInitializer]]: """ Find concrete ``PyRITInitializer`` subclasses defined in *module*. @@ -296,21 +299,39 @@ def _module_defined_initializers( Args: module: The imported module to scan. base_class: The ``PyRITInitializer`` base class. + strict: Raise if an explicitly declared export cannot be loaded. Returns: list[type]: Concrete initializer classes defined in the module. """ - return [ - attr - for attr_name in dir(module) + static_members = dict(inspect.getmembers_static(module)) + member_names = set(static_members) + declared_exports = vars(module).get("__all__", ()) + if isinstance(declared_exports, (list, tuple)): + member_names.update(name for name in declared_exports if isinstance(name, str)) + + initializers: list[type[PyRITInitializer]] = [] + for attr_name in sorted(member_names): + if attr_name in static_members: + attr = static_members[attr_name] + else: + try: + attr = getattr(module, attr_name) + except Exception as error: + if strict: + raise + logger.warning(f"Failed to load declared export '{attr_name}' from {module.__name__}: {error}") + continue if ( - inspect.isclass(attr := getattr(module, attr_name)) + inspect.isclass(attr) and issubclass(attr, base_class) and attr is not base_class and not inspect.isabstract(attr) and attr.__module__ == module.__name__ - ) - ] + ): + initializers.append(attr) + + return initializers def create_from_script_paths( self, *, script_paths: Sequence[str | Path], strict: bool = False @@ -353,7 +374,11 @@ def create_from_script_paths( file_path=script_path, module_name=f"init_script_{script_path.stem}" ) file_instances: list[PyRITInitializer] = [] - for init_cls in self._module_defined_initializers(module=module, base_class=PyRITInitializer): + for init_cls in self._module_defined_initializers( + module=module, + base_class=PyRITInitializer, + strict=strict, + ): try: file_instances.append(init_cls()) logger.debug(f"Found and instantiated {init_cls.__name__} in {script_path.name}") @@ -500,7 +525,11 @@ def _load_custom_initializer_class( raise ValueError(f"Failed to write initializer script: {error}") from error try: module = self._load_module_from_path(file_path=script_path, module_name=f"custom_initializer.{name}") - discovered_classes = self._module_defined_initializers(module=module, base_class=PyRITInitializer) + discovered_classes = self._module_defined_initializers( + module=module, + base_class=PyRITInitializer, + strict=True, + ) except Exception as error: raise ValueError(f"Failed to load initializer script '{name}': {error}") from error diff --git a/tests/unit/registry/test_initializer_registry.py b/tests/unit/registry/test_initializer_registry.py index 71556acfdd..f8a5a0ecd9 100644 --- a/tests/unit/registry/test_initializer_registry.py +++ b/tests/unit/registry/test_initializer_registry.py @@ -340,6 +340,25 @@ def _write_initializer_script(directory: Path, filename: str, *class_names: str) return script_path +_BROKEN_DECLARED_EXPORT_SCRIPT = """ +from pyrit.setup.pyrit_initializer import PyRITInitializer + +class GoodInitializer(PyRITInitializer): + async def initialize_async(self) -> None: + pass + +__all__ = ["GoodInitializer", "BrokenInitializer"] + +def __dir__(): + return __all__ + +def __getattr__(name): + if name == "BrokenInitializer": + raise RuntimeError("missing required dependency") + raise AttributeError(name) +""" + + def test_create_from_script_paths_loads_multiple_classes(lazy_registry): """Test that all initializer subclasses defined in a file are instantiated.""" with tempfile.TemporaryDirectory() as temp_dir: @@ -352,6 +371,33 @@ def test_create_from_script_paths_loads_multiple_classes(lazy_registry): assert lazy_registry.get_class_names() == [] +def test_create_from_script_paths_skips_broken_declared_export_when_not_strict(lazy_registry): + """Best-effort script loading keeps valid initializers when a declared export is unavailable.""" + with tempfile.TemporaryDirectory() as temp_dir: + path = Path(temp_dir) / "broken_export.py" + path.write_text(_BROKEN_DECLARED_EXPORT_SCRIPT) + + instances = lazy_registry.create_from_script_paths(script_paths=[path], strict=False) + + assert [type(instance).__name__ for instance in instances] == ["GoodInitializer"] + + +def test_create_from_script_paths_preserves_broken_declared_export_error_when_strict(lazy_registry): + """Strict script loading must not silently accept an incomplete script.""" + with tempfile.TemporaryDirectory() as temp_dir: + path = Path(temp_dir) / "broken_export.py" + path.write_text(_BROKEN_DECLARED_EXPORT_SCRIPT) + + with pytest.raises(RuntimeError, match="missing required dependency"): + lazy_registry.create_from_script_paths(script_paths=[path], strict=True) + + +def test_register_from_content_preserves_broken_declared_export_error(lazy_registry): + """Custom-script validation must reject a declared export that cannot be loaded.""" + with pytest.raises(ValueError, match="missing required dependency"): + lazy_registry.register_from_content(name="broken_export", script_content=_BROKEN_DECLARED_EXPORT_SCRIPT) + + def test_create_from_script_paths_rejects_non_python_file(lazy_registry): """Test that a non-.py path raises ValueError.""" with tempfile.TemporaryDirectory() as temp_dir: @@ -526,6 +572,53 @@ def test_discover_directory_registers_and_lists_metadata(): assert any(m.registry_name == "solo" for m in metadata) +def test_single_file_discovery_skips_unrelated_dynamic_attributes( + tmp_path: Path, caplog: pytest.LogCaptureFixture +) -> None: + """Test that a single-file discovery path ignores unrelated dynamic module attributes.""" + script = _SOLO_SCRIPT + ( + "\ndef __dir__():\n" + " return ['SoloInitializer', 'dynamic_attribute']\n\n" + "def __getattr__(name):\n" + " raise RuntimeError(f'dynamic attribute evaluated: {name}')\n" + ) + script_path = tmp_path / "dynamic.py" + script_path.write_text(script, encoding="utf-8") + + registry = InitializerRegistry(discovery_path=script_path, lazy_discovery=False) + + assert registry.get_class_names() == ["solo"] + assert "dynamic attribute evaluated" not in caplog.text + + +@pytest.mark.parametrize("include_lazy_export_in_dir", [False, True]) +def test_single_file_discovery_registers_declared_lazy_exports( + tmp_path: Path, include_lazy_export_in_dir: bool, caplog: pytest.LogCaptureFixture +) -> None: + script = _SOLO_SCRIPT + ( + "\n__all__ = ['SoloInitializer', 'LazyInitializer', 'non_class']\n" + "def __dir__():\n" + f" return ['SoloInitializer', 'unrelated'] + " + f"{['LazyInitializer'] if include_lazy_export_in_dir else []!r}\n" + "def __getattr__(name):\n" + " if name == 'LazyInitializer':\n" + " class LazyInitializer(PyRITInitializer):\n" + " async def initialize_async(self) -> None:\n" + " pass\n" + " return LazyInitializer\n" + " if name == 'non_class':\n" + " return 42\n" + " raise RuntimeError(f'unexpected dynamic lookup: {name}')\n" + ) + script_path = tmp_path / "lazy.py" + script_path.write_text(script, encoding="utf-8") + + registry = InitializerRegistry(discovery_path=script_path, lazy_discovery=False) + + assert set(registry.get_class_names()) == {"lazy", "solo"} + assert not caplog.records + + def test_discover_directory_does_not_evaluate_dynamic_module_attributes(caplog): """Test that discovery inspects ordinary module members without invoking dynamic attributes.""" script = (