fix: isolate path-backed python source modules
This commit is contained in:
@@ -3,8 +3,12 @@ from __future__ import annotations
|
|||||||
import sys
|
import sys
|
||||||
from collections.abc import Mapping, Sequence
|
from collections.abc import Mapping, Sequence
|
||||||
from dataclasses import dataclass
|
from dataclasses import dataclass
|
||||||
|
from hashlib import sha256
|
||||||
from importlib import import_module
|
from importlib import import_module
|
||||||
|
from importlib import util as importlib_util
|
||||||
|
from importlib.machinery import ModuleSpec
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
|
from types import ModuleType
|
||||||
from typing import Any, Protocol
|
from typing import Any, Protocol
|
||||||
|
|
||||||
from wf_authoring import NodeSpec, node
|
from wf_authoring import NodeSpec, node
|
||||||
@@ -43,9 +47,11 @@ def load_python_source(
|
|||||||
path: str | Path | None = None,
|
path: str | Path | None = None,
|
||||||
) -> CapabilitySource:
|
) -> CapabilitySource:
|
||||||
"""Load a trusted in-process Python source from a module registry object."""
|
"""Load a trusted in-process Python source from a module registry object."""
|
||||||
if path is not None:
|
module_obj = _import_source_module(
|
||||||
_ensure_import_path(Path(path))
|
source_id=source_id,
|
||||||
module_obj = import_module(module)
|
module=module,
|
||||||
|
path=Path(path) if path is not None else None,
|
||||||
|
)
|
||||||
if not hasattr(module_obj, registry):
|
if not hasattr(module_obj, registry):
|
||||||
raise ValueError(f"missing registry object {registry!r} in module {module!r}")
|
raise ValueError(f"missing registry object {registry!r} in module {module!r}")
|
||||||
raw_registry = getattr(module_obj, registry)
|
raw_registry = getattr(module_obj, registry)
|
||||||
@@ -101,6 +107,114 @@ def _coerce_specs(raw_registry: object) -> list[NodeSpec[Any, Any]]:
|
|||||||
return specs
|
return specs
|
||||||
|
|
||||||
|
|
||||||
|
def _import_source_module(
|
||||||
|
*,
|
||||||
|
source_id: str,
|
||||||
|
module: str,
|
||||||
|
path: Path | None,
|
||||||
|
) -> ModuleType:
|
||||||
|
"""Import a source module without letting same-name local modules collide.
|
||||||
|
|
||||||
|
Configured Python sources commonly use local names like `ops` under their
|
||||||
|
own source root. Importing those through `sys.path` and `import_module`
|
||||||
|
would share `sys.modules["ops"]` across unrelated examples/tests. When the
|
||||||
|
module file exists under the configured root, load it under a stable
|
||||||
|
synthetic name derived from the source id and root path.
|
||||||
|
"""
|
||||||
|
if path is None:
|
||||||
|
return import_module(module)
|
||||||
|
|
||||||
|
resolved_root = path.resolve()
|
||||||
|
module_file = _module_file_under_root(resolved_root, module)
|
||||||
|
if module_file is None:
|
||||||
|
_ensure_import_path(resolved_root)
|
||||||
|
return import_module(module)
|
||||||
|
|
||||||
|
synthetic_root = _synthetic_module_root(
|
||||||
|
source_id=source_id,
|
||||||
|
root=resolved_root,
|
||||||
|
)
|
||||||
|
_ensure_synthetic_packages(
|
||||||
|
synthetic_root=synthetic_root,
|
||||||
|
root=resolved_root,
|
||||||
|
module=module,
|
||||||
|
)
|
||||||
|
# Keep legacy absolute sibling imports (`import helper`) working where
|
||||||
|
# possible, but source isolation only applies to imports made through the
|
||||||
|
# synthetic package (`from . import helper`, `from ..shared import value`).
|
||||||
|
# Python's normal absolute import cache can still collide for local helper
|
||||||
|
# names; robust sources should use package-relative imports.
|
||||||
|
_ensure_import_path(resolved_root)
|
||||||
|
|
||||||
|
synthetic_name = f"{synthetic_root}.{module}"
|
||||||
|
cached = sys.modules.get(synthetic_name)
|
||||||
|
if cached is not None:
|
||||||
|
return cached
|
||||||
|
|
||||||
|
submodule_search_locations = (
|
||||||
|
[str(module_file.parent)] if module_file.name == "__init__.py" else None
|
||||||
|
)
|
||||||
|
spec = importlib_util.spec_from_file_location(
|
||||||
|
synthetic_name,
|
||||||
|
module_file,
|
||||||
|
submodule_search_locations=submodule_search_locations,
|
||||||
|
)
|
||||||
|
if spec is None or spec.loader is None:
|
||||||
|
raise ImportError(f"cannot load Python source module {module!r}")
|
||||||
|
|
||||||
|
module_obj = importlib_util.module_from_spec(spec)
|
||||||
|
sys.modules[synthetic_name] = module_obj
|
||||||
|
spec.loader.exec_module(module_obj)
|
||||||
|
return module_obj
|
||||||
|
|
||||||
|
|
||||||
|
def _module_file_under_root(root: Path, module: str) -> Path | None:
|
||||||
|
module_path = root.joinpath(*module.split("."))
|
||||||
|
file_path = module_path.with_suffix(".py")
|
||||||
|
if file_path.is_file():
|
||||||
|
return file_path
|
||||||
|
package_path = module_path / "__init__.py"
|
||||||
|
if package_path.is_file():
|
||||||
|
return package_path
|
||||||
|
return None
|
||||||
|
|
||||||
|
|
||||||
|
def _synthetic_module_root(*, source_id: str, root: Path) -> str:
|
||||||
|
digest = sha256(f"{source_id}\0{root}".encode("utf-8")).hexdigest()[:16]
|
||||||
|
safe_source = source_id.replace(".", "_").replace("-", "_")
|
||||||
|
return f"_wf_source_{safe_source}_{digest}"
|
||||||
|
|
||||||
|
|
||||||
|
def _ensure_synthetic_packages(
|
||||||
|
*,
|
||||||
|
synthetic_root: str,
|
||||||
|
root: Path,
|
||||||
|
module: str,
|
||||||
|
) -> None:
|
||||||
|
"""Create package shells so source modules can use relative imports."""
|
||||||
|
_ensure_package_module(synthetic_root, root)
|
||||||
|
parts = module.split(".")[:-1]
|
||||||
|
current_name = synthetic_root
|
||||||
|
current_path = root
|
||||||
|
for part in parts:
|
||||||
|
current_name = f"{current_name}.{part}"
|
||||||
|
current_path = current_path / part
|
||||||
|
_ensure_package_module(current_name, current_path)
|
||||||
|
|
||||||
|
|
||||||
|
def _ensure_package_module(name: str, path: Path) -> None:
|
||||||
|
existing = sys.modules.get(name)
|
||||||
|
if existing is not None:
|
||||||
|
existing.__path__ = [str(path)] # type: ignore[attr-defined]
|
||||||
|
return
|
||||||
|
package = ModuleType(name)
|
||||||
|
package.__file__ = str(path)
|
||||||
|
package.__package__ = name
|
||||||
|
package.__path__ = [str(path)] # type: ignore[attr-defined]
|
||||||
|
package.__spec__ = ModuleSpec(name, loader=None, is_package=True)
|
||||||
|
sys.modules[name] = package
|
||||||
|
|
||||||
|
|
||||||
def _ensure_import_path(path: Path) -> None:
|
def _ensure_import_path(path: Path) -> None:
|
||||||
"""Keep configured trusted source roots importable for deferred imports."""
|
"""Keep configured trusted source roots importable for deferred imports."""
|
||||||
resolved = str(path.resolve())
|
resolved = str(path.resolve())
|
||||||
|
|||||||
@@ -86,6 +86,152 @@ registry = [echo]
|
|||||||
assert set(source.capabilities.node_specs) == {"local.ops.echo"}
|
assert set(source.capabilities.node_specs) == {"local.ops.echo"}
|
||||||
|
|
||||||
|
|
||||||
|
def test_load_python_source_isolates_same_module_name_across_roots(
|
||||||
|
tmp_path: Path,
|
||||||
|
) -> None:
|
||||||
|
left_root = tmp_path / "left"
|
||||||
|
right_root = tmp_path / "right"
|
||||||
|
left_root.mkdir()
|
||||||
|
right_root.mkdir()
|
||||||
|
source_template = """
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
from pydantic import BaseModel
|
||||||
|
from wf_authoring import node
|
||||||
|
|
||||||
|
class Input(BaseModel):
|
||||||
|
text: str
|
||||||
|
|
||||||
|
class Output(BaseModel):
|
||||||
|
value: str
|
||||||
|
|
||||||
|
@node(name="{name}")
|
||||||
|
def operation(payload: Input) -> Output:
|
||||||
|
return Output(value=payload.text)
|
||||||
|
|
||||||
|
registry = [operation]
|
||||||
|
"""
|
||||||
|
(left_root / "ops.py").write_text(
|
||||||
|
source_template.format(name="left"), encoding="utf-8"
|
||||||
|
)
|
||||||
|
(right_root / "ops.py").write_text(
|
||||||
|
source_template.format(name="right"), encoding="utf-8"
|
||||||
|
)
|
||||||
|
|
||||||
|
left = load_python_source(
|
||||||
|
source_id="local.left",
|
||||||
|
path=left_root,
|
||||||
|
module="ops",
|
||||||
|
registry="registry",
|
||||||
|
)
|
||||||
|
right = load_python_source(
|
||||||
|
source_id="local.right",
|
||||||
|
path=right_root,
|
||||||
|
module="ops",
|
||||||
|
registry="registry",
|
||||||
|
)
|
||||||
|
|
||||||
|
assert set(left.capabilities.node_specs) == {"local.left.left"}
|
||||||
|
assert set(right.capabilities.node_specs) == {"local.right.right"}
|
||||||
|
|
||||||
|
|
||||||
|
def test_load_python_source_supports_relative_imports_under_source_root(
|
||||||
|
tmp_path: Path,
|
||||||
|
) -> None:
|
||||||
|
source_root = tmp_path / "relative_source"
|
||||||
|
package_root = source_root / "pkg"
|
||||||
|
package_root.mkdir(parents=True)
|
||||||
|
(package_root / "__init__.py").write_text("", encoding="utf-8")
|
||||||
|
(package_root / "labels.py").write_text(
|
||||||
|
'LABEL = "relative"\n',
|
||||||
|
encoding="utf-8",
|
||||||
|
)
|
||||||
|
(package_root / "ops.py").write_text(
|
||||||
|
"""
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
from pydantic import BaseModel
|
||||||
|
from wf_authoring import node
|
||||||
|
|
||||||
|
from .labels import LABEL
|
||||||
|
|
||||||
|
class Input(BaseModel):
|
||||||
|
text: str
|
||||||
|
|
||||||
|
class Output(BaseModel):
|
||||||
|
value: str
|
||||||
|
|
||||||
|
@node(name=LABEL)
|
||||||
|
def operation(payload: Input) -> Output:
|
||||||
|
return Output(value=payload.text)
|
||||||
|
|
||||||
|
registry = [operation]
|
||||||
|
""",
|
||||||
|
encoding="utf-8",
|
||||||
|
)
|
||||||
|
|
||||||
|
source = load_python_source(
|
||||||
|
source_id="local.relative",
|
||||||
|
path=source_root,
|
||||||
|
module="pkg.ops",
|
||||||
|
registry="registry",
|
||||||
|
)
|
||||||
|
|
||||||
|
assert set(source.capabilities.node_specs) == {"local.relative.relative"}
|
||||||
|
|
||||||
|
|
||||||
|
def test_load_python_source_supports_delayed_relative_imports(
|
||||||
|
tmp_path: Path,
|
||||||
|
) -> None:
|
||||||
|
source_root = tmp_path / "delayed_relative_source"
|
||||||
|
package_root = source_root / "pkg"
|
||||||
|
package_root.mkdir(parents=True)
|
||||||
|
(package_root / "__init__.py").write_text("", encoding="utf-8")
|
||||||
|
(package_root / "labels.py").write_text(
|
||||||
|
'LABEL = "delayed"\n',
|
||||||
|
encoding="utf-8",
|
||||||
|
)
|
||||||
|
(package_root / "registry.py").write_text(
|
||||||
|
"""
|
||||||
|
def build_registry():
|
||||||
|
from .ops import operation
|
||||||
|
|
||||||
|
return [operation]
|
||||||
|
""",
|
||||||
|
encoding="utf-8",
|
||||||
|
)
|
||||||
|
(package_root / "ops.py").write_text(
|
||||||
|
"""
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
from pydantic import BaseModel
|
||||||
|
from wf_authoring import node
|
||||||
|
|
||||||
|
from .labels import LABEL
|
||||||
|
|
||||||
|
class Input(BaseModel):
|
||||||
|
text: str
|
||||||
|
|
||||||
|
class Output(BaseModel):
|
||||||
|
value: str
|
||||||
|
|
||||||
|
@node(name=LABEL)
|
||||||
|
def operation(payload: Input) -> Output:
|
||||||
|
return Output(value=payload.text)
|
||||||
|
""",
|
||||||
|
encoding="utf-8",
|
||||||
|
)
|
||||||
|
|
||||||
|
source = load_python_source(
|
||||||
|
source_id="local.delayed",
|
||||||
|
path=source_root,
|
||||||
|
module="pkg.registry",
|
||||||
|
registry="build_registry",
|
||||||
|
)
|
||||||
|
|
||||||
|
assert set(source.capabilities.node_specs) == {"local.delayed.delayed"}
|
||||||
|
|
||||||
|
|
||||||
def test_load_python_source_propagates_enabled_flag() -> None:
|
def test_load_python_source_propagates_enabled_flag() -> None:
|
||||||
source = load_python_source(
|
source = load_python_source(
|
||||||
source_id="local.ops",
|
source_id="local.ops",
|
||||||
|
|||||||
Reference in New Issue
Block a user