Fix pyfuse.install_package_as

Author: codeSamuraii

pyfuse/graph/analyzer.py

@@ -121,15 +121,27 @@ def get_module_assignments(func: Callable[..., object]) -> dict[str, str]:
     return assignments
 
 
+def _is_install_package_as_call(expr: ast.expr) -> bool:
+    """Match ``install_package_as`` or ``pyfuse.install_package_as``."""
+    if isinstance(expr, ast.Name):
+        return expr.id == "install_package_as"
+    if isinstance(expr, ast.Attribute):
+        return expr.attr == "install_package_as"
+    return False
+
+
 def _parse_install_package_as(node: ast.With) -> str | None:
-    """Return the package name if *node* is ``with install_package_as(...)``."""
+    """Return the package name if *node* is ``with install_package_as(...)``.
+
+    Accepts both the bare form (``install_package_as("foo")``) and the
+    attribute form (``pyfuse.install_package_as("foo")``).
+    """
     if len(node.items) != 1:
         return None
     ctx = node.items[0].context_expr
     if not (
         isinstance(ctx, ast.Call)
-        and isinstance(ctx.func, ast.Name)
-        and ctx.func.id == "install_package_as"
+        and _is_install_package_as_call(ctx.func)
         and len(ctx.args) == 1
         and isinstance(ctx.args[0], ast.Constant)
         and isinstance(ctx.args[0].value, str)

tests/test_deps.py

@@ -1,3 +1,5 @@
+import ast
+
 import pytest
 
 from pyfuse.worker.deps import (
@@ -9,6 +11,7 @@
     install_package_as,
     is_installed,
 )
+from pyfuse.graph.analyzer import _parse_install_package_as
 from pyfuse.core.models import FunctionNode, ImportInfo
 from pyfuse.graph.store import Store
 
@@ -156,6 +159,32 @@ def test_package_none_from_dict_without_key(self) -> None:
         assert restored.package is None
 
 
+# -- _parse_install_package_as ----------------------------------------------
+
+class TestParseInstallPackageAs:
+    def _parse(self, source: str) -> str | None:
+        tree = ast.parse(source)
+        with_node = tree.body[0]
+        assert isinstance(with_node, ast.With)
+        return _parse_install_package_as(with_node)
+
+    def test_bare_form(self) -> None:
+        assert self._parse('with install_package_as("PyYAML"):\n    pass\n') == "PyYAML"
+
+    def test_attribute_form(self) -> None:
+        """``with pyfuse.install_package_as(...)`` is also recognized."""
+        assert (
+            self._parse('with pyfuse.install_package_as("python-multipart"):\n    pass\n')
+            == "python-multipart"
+        )
+
+    def test_unrelated_with_block(self) -> None:
+        assert self._parse('with open("x") as f:\n    pass\n') is None
+
+    def test_missing_argument(self) -> None:
+        assert self._parse('with install_package_as():\n    pass\n') is None
+
+
 # -- _collect_package_hints --------------------------------------------------
 
 class TestCollectPackageHints:

tests/test_venv.py

@@ -202,6 +202,20 @@ def test_install_package_as(self, tmp_path: pytest.TempPathFactory) -> None:
         assert "yaml" not in packages
         assert "dateutil" not in packages
 
+    def test_install_package_as_attribute_form(
+        self, tmp_path: pytest.TempPathFactory,
+    ) -> None:
+        """``with pyfuse.install_package_as(...)`` is also recognized."""
+        script = tmp_path / "test_ipa_attr.py"  # type: ignore[operator]
+        script.write_text(
+            "import pyfuse\n"
+            'with pyfuse.install_package_as("python-multipart"):\n'
+            "    import multipart\n"
+        )
+        packages = _detect_script_packages(str(script))
+        assert "python-multipart" in packages
+        assert "multipart" not in packages
+
 
 class TestDetectPyfuseExtras:
     def test_connect_redis(self, tmp_path: pytest.TempPathFactory) -> None: