Rewrite type references in stubs to include prefix

Stubs are placed in the scyjava.types namespace by default,
but have imports with a toplevel namespace; for example,
when generating stubs for org.scijava:scijava-common, the file:

    src/scijava/types/org/scijava/__init__.pyi

would declare imports:

    import java.lang
    import java.util

    import org.scijava.annotations
    import org.scijava.app
    import org.scijava.cache
    import org.scijava.command

and so on. But these imports yield errors when browsing in an IDE,
and prevent the type checker from resolving all the types properly.

This commit makes the following changes to address the problem:

- Add a python_package_prefix parameter to generate_stubs() function
- Add a _rewrite_stub_imports() function that:
  - Rewrites import foo.bar.X → import {python_package_prefix}.foo.bar.X
  - Rewrites type references similarly to have the python_package_prefix
- Update the CLI to automatically pass
  python_package_prefix="scyjava.types" when using the default location
- If --output-python-path gives a different prefix, use that instead
- Add a new test_stubgen_type_references to validate rewriting behavior

So then the following example above gets rewritten to be:

    import java.lang
    import java.util

    import scyjava.stubs.org.scijava.annotations
    import scyjava.stubs.org.scijava.app
    import scyjava.stubs.org.scijava.cache
    import scyjava.stubs.org.scijava.command

Note that the java.lang and java.util imports are not rewritten,
because they are not among the classes whose stubs are being generated.

With this change, IDEs now autocomplete expressions involving these Java
classes correctly, even when chained; for example:

    import scyjava.stubs.org.scijava.Context

    Context().getServiceIndex().getA

shows getAll as a valid completion, because getAll() is a member method
of a supertype of ServiceIndex, the class returned by getServiceIndex().
Before this patch, such chained type completions did not work.

Co-authored-by: Claude <noreply@anthropic.com>

Author: ctrueden

src/scyjava/_stubs/_cli.py

@@ -102,6 +102,11 @@ def main() -> None:
     if not output_dir.exists():
         output_dir.mkdir(parents=True, exist_ok=True)
 
+    # Determine the Python package prefix for import rewriting
+    python_package_prefix = args.output_python_path or _derive_python_prefix(
+        args.output_dir
+    )
+
     generate_stubs(
         endpoints=args.endpoints,
         prefixes=args.prefix,
@@ -110,9 +115,23 @@ def main() -> None:
         include_javadoc=args.with_javadoc,
         add_runtime_imports=args.runtime_imports,
         remove_namespace_only_stubs=args.remove_namespace_only_stubs,
+        python_package_prefix=python_package_prefix,
     )
 
 
+def _derive_python_prefix(output_dir: str | None) -> str:
+    """Derive the Python package prefix from the output directory.
+
+    If output_dir is None, defaults to 'scyjava.types'.
+    """
+    if output_dir:
+        # For a filesystem path, we can't reliably derive the Python prefix
+        # Return empty string to skip import rewriting
+        return ""
+    # Default case: stubs go to scyjava.types
+    return "scyjava.types"
+
+
 def _get_ouput_dir(output_dir: str | None, python_path: str | None) -> Path:
     if out_dir := output_dir:
         return Path(out_dir)

src/scyjava/_stubs/_genstubs.py

@@ -41,6 +41,7 @@ def generate_stubs(
     include_javadoc: bool = True,
     add_runtime_imports: bool = True,
     remove_namespace_only_stubs: bool = False,
+    python_package_prefix: str = "",
 ) -> None:
     """Generate stubs for the given maven endpoints.
 
@@ -82,6 +83,11 @@ def generate_stubs(
         merge the generated stubs with other stubs in the same namespace.  Without this,
         the `__init__.pyi` for any given module will be whatever whatever the *last*
         stub generator wrote to it (and therefore inaccurate).
+    python_package_prefix : str, optional
+        The Python package prefix under which stubs are being installed. For example,
+        if stubs are being installed to `scyjava.types.org.scijava...`, this should be
+        "scyjava.types". This is used to rewrite imports in the stub files so that
+        type checkers can properly resolve cross-references. Defaults to "".
     """
     try:
         import stubgenj
@@ -136,10 +142,20 @@ def _patched_start(*args: Any, **kwargs: Any) -> None:
     )
 
     output_dir = Path(output_dir)
+    if python_package_prefix:
+        logger.info(
+            "Rewriting stub imports with Python package prefix: %s",
+            python_package_prefix,
+        )
+
     if add_runtime_imports:
         logger.info("Adding runtime imports to generated stubs")
 
     for stub in output_dir.rglob("*.pyi"):
+        # Rewrite imports if a Python package prefix was specified
+        if python_package_prefix:
+            _rewrite_stub_imports(stub, python_package_prefix)
+
         stub_ast = ast.parse(stub.read_text())
         members = {node.name for node in stub_ast.body if hasattr(node, "name")}
         if members == {"__module_protocol__"}:
@@ -178,6 +194,113 @@ def _patched_start(*args: Any, **kwargs: Any) -> None:
 """
 
 
+def _rewrite_stub_imports(stub_path: Path, python_package_prefix: str) -> None:
+    """Rewrite imports in a stub file to use the full Python package path.
+
+    When stubs are generated into a subdirectory like scyjava/types, they need to have
+    their imports rewritten so that type checkers can resolve cross-references. This
+    function transforms imports like:
+
+        import org.scijava.object
+
+    into:
+
+        import scyjava.types.org.scijava.object
+
+    and transforms type references like:
+
+        org.scijava.object.ObjectIndex
+
+    into:
+
+        scyjava.types.org.scijava.object.ObjectIndex
+    """
+    import re
+
+    content = stub_path.read_text()
+
+    # Split into lines for import processing
+    lines = content.split("\n")
+    new_lines = []
+    import_patterns = []  # Patterns to replace in annotations
+
+    i = 0
+    while i < len(lines):
+        line = lines[i]
+        stripped = line.strip()
+
+        # Handle import statements
+        if stripped.startswith("import ") and (
+            "org." in stripped or "java." in stripped
+        ):
+            # Parse "import org.scijava.service"
+            match = stripped.split()
+            if len(match) >= 2:
+                module_name = match[1]
+                if (
+                    not module_name.startswith(".")
+                    and "scyjava.types" not in module_name
+                ):
+                    # Only rewrite org.* imports (not java.*)
+                    if module_name.startswith("org."):
+                        new_module = f"{python_package_prefix}.{module_name}"
+                        # Preserve indentation
+                        indent = line[: len(line) - len(line.lstrip())]
+                        new_lines.append(f"{indent}import {new_module}")
+                        # Record this pattern for later annotation rewriting
+                        import_patterns.append((module_name, new_module))
+                        i += 1
+                        continue
+
+        # Handle "from X import Y" statements
+        elif stripped.startswith("from ") and (" import " in stripped):
+            if "org." in stripped or "java." in stripped:
+                parts = stripped.split(" import ")
+                if len(parts) == 2:
+                    module_part = parts[0].replace("from ", "").strip()
+                    imports_part = parts[1].strip()
+
+                    if (
+                        not module_part.startswith(".")
+                        and "scyjava.types" not in module_part
+                    ):
+                        if module_part.startswith("org."):
+                            new_module = f"{python_package_prefix}.{module_part}"
+                            indent = line[: len(line) - len(line.lstrip())]
+                            new_lines.append(
+                                f"{indent}from {new_module} import {imports_part}"
+                            )
+                            import_patterns.append((module_part, new_module))
+                            i += 1
+                            continue
+
+        new_lines.append(line)
+        i += 1
+
+    # Reconstruct content with rewritten imports
+    new_content = "\n".join(new_lines)
+
+    # Now rewrite type annotations that reference org.* packages
+    # Only replace in type hints, not in already-rewritten import statements
+    # We do this by replacing the old module names with new ones, but being careful
+    # to not double-replace
+    for old_prefix, new_prefix in import_patterns:
+        # Replace qualified names like "org.scijava.service.ServiceIndex"
+        # but NOT names that already have the prefix
+        # Pattern: old_prefix followed by a dot and word characters
+        # Use negative lookbehind to avoid replacing if already prefixed
+        pattern = (
+            r"(?<!"
+            + re.escape(python_package_prefix)
+            + r"\.)"
+            + re.escape(old_prefix)
+            + r"(?=\.\w)"
+        )
+        new_content = re.sub(pattern, new_prefix, new_content)
+
+    stub_path.write_text(new_content)
+
+
 def ruff_check(output: Path, select: str = "E,W,F,I,UP,C4,B,RUF,TC,TID") -> None:
     """Run ruff check and format on the generated stubs."""
     if not shutil.which("ruff"):

tests/test_stubgen.py

@@ -1,7 +1,8 @@
 from __future__ import annotations
 
+import ast
 import sys
-from typing import TYPE_CHECKING
+from pathlib import Path
 from unittest.mock import patch
 
 import jpype
@@ -10,9 +11,6 @@
 import scyjava
 from scyjava._stubs import _cli
 
-if TYPE_CHECKING:
-    from pathlib import Path
-
 
 @pytest.mark.skipif(
     scyjava.config.mode != scyjava.config.Mode.JPYPE,
@@ -56,3 +54,81 @@ def test_stubgen(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None:
         func = Function(1)
         mock_start_jvm.assert_called_once()
         assert isinstance(func, jpype.JObject)
+
+
+@pytest.mark.skipif(
+    scyjava.config.mode != scyjava.config.Mode.JPYPE,
+    reason="Stubgen not supported in JEP",
+)
+def test_stubgen_type_references(
+    monkeypatch: pytest.MonkeyPatch, tmp_path: Path
+) -> None:
+    """Test that generated stubs have properly qualified type references.
+
+    This validates that when stubs are generated with a Python package prefix,
+    all type references are properly rewritten so type checkers can resolve them.
+    """
+    import tempfile
+
+    # Generate stubs with --output-python-path so a prefix is used
+    # (rather than --output-dir which doesn't imply a Python module path)
+    stubs_module = "test_stubs"
+
+    # Create a temporary directory and add it to sys.path
+    with tempfile.TemporaryDirectory() as tmpdir_str:
+        tmpdir = Path(tmpdir_str)
+        original_path = sys.path.copy()
+
+        try:
+            sys.path.insert(0, str(tmpdir))
+
+            # Create the parent module package
+            stubs_pkg = tmpdir / stubs_module
+            stubs_pkg.mkdir()
+            (stubs_pkg / "__init__.py").touch()
+
+            monkeypatch.setattr(
+                sys,
+                "argv",
+                [
+                    "scyjava-stubgen",
+                    "org.scijava:parsington:3.1.0",
+                    "--output-python-path",
+                    stubs_module,
+                ],
+            )
+            _cli.main()
+
+            # Check that import statements were rewritten with the prefix
+            init_stub = stubs_pkg / "org" / "scijava" / "parsington" / "__init__.pyi"
+            assert init_stub.exists(), f"Expected stub file {init_stub} not found"
+
+            content = init_stub.read_text()
+            stub_ast = ast.parse(content)
+
+            # Find all Import and ImportFrom nodes
+            imports = [
+                node
+                for node in ast.walk(stub_ast)
+                if isinstance(node, (ast.Import, ast.ImportFrom))
+            ]
+
+            # Collect imported module names
+            imported_modules = set()
+            for imp in imports:
+                if isinstance(imp, ast.Import):
+                    for alias in imp.names:
+                        imported_modules.add(alias.name)
+                elif isinstance(imp, ast.ImportFrom) and imp.module:
+                    imported_modules.add(imp.module)
+
+            # Verify that bare org.scijava.* imports don't exist (they should be prefixed)
+            org_imports = {
+                m for m in imported_modules if m and m.startswith("org.scijava.")
+            }
+            assert not org_imports, (
+                f"Found unrewritten org.scijava imports in {init_stub}: {org_imports}. "
+                f"These should have been prefixed with '{stubs_module}.'"
+            )
+        finally:
+            sys.path = original_path