FIX avoid dynamic attribute evaluation in single-file initializer discovery

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
 

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 = (