Fix: current-module prefix stripping (#301)

`fix.py`:
Current-module prefix stripping now also walks nested annotation
parameters instead of only the outer type. The fix covers the exact
shapes that broke in my downstream projects (pyAMReX, ImpactX):
`list[...]`, `typing.Optional[...]`, and `dict[..., ...]` for the type
`...`.

Concrete example seen downstream in ImpactX for
`impactx.MADXParser.*` types: LHS broken -> RHS correct now.
- `list[impactx.MADXParser.Token]` -> `list[Token]`
- `typing.Optional[impactx.MADXParser.Expression]` ->
`typing.Optional[Expression]`
- `dict[str, impactx.MADXParser.Expression]` -> `dict[str, Expression]`

This originated in a pure Python file that is used on top of pybind11
modules. I have seen similar issues for a while though and [patch
around](https://github.com/AMReX-Codes/pyamrex/blob/26.04/.github/update_stub.sh#L20-L26)
them
[awkwardly](https://github.com/BLAST-ImpactX/impactx/blob/26.03/.github/update_stub.sh#L19-L21).

`parse.py`:
Runtime generic annotations are parsed before falling back to opaque
values, and `_is_generic_alias()` now also recognizes typing generics
like `Optional[...]`

<details>
<summary>Reproducer (PyTest)</summary>

```py
"""Regression tests for current-module prefix cleanup in parsed annotations."""

from argparse import Namespace
from types import ModuleType

from pybind11_stubgen import stub_parser_from_args
from pybind11_stubgen.printer import Printer
from pybind11_stubgen.structs import QualifiedName


def make_args() -> Namespace:
    """Build the same parser configuration shape used by the CLI entrypoint."""

    return Namespace(
        output_dir=".",
        root_suffix=None,
        ignore_invalid_expressions=None,
        ignore_invalid_identifiers=None,
        ignore_unresolved_names=None,
        ignore_all_errors=False,
        enum_class_locations=[],
        numpy_array_wrap_with_annotated=False,
        numpy_array_use_type_var=False,
        numpy_array_remove_parameters=False,
        print_invalid_expressions_as_is=False,
        print_safe_value_reprs=None,
        print_value_comments=False,
        exit_code=False,
        dry_run=True,
        stub_extension="pyi",
        module_names=[],
    )


def make_module() -> ModuleType:
    """
    Create an in-memory module that mirrors the ImpactX regression.

    The module intentionally does not use ``from __future__ import annotations``.
    That forces Python to evaluate local annotations such as ``list[Token]`` into
    runtime generic objects whose string form contains the fully-qualified module
    name, e.g. ``list[impactx.MADXParser.Token]``.
    """

    module = ModuleType("impactx.MADXParser")
    module.__dict__["__name__"] = "impactx.MADXParser"

    exec(
        """
from typing import Optional


class Token:
    pass


class Expression:
    pass


def tokenize(tokens: list[Token], expr: Optional[Expression] = None) -> dict[str, Expression]:
    pass
""",
        module.__dict__,
    )

    return module


def test_current_module_prefix_is_stripped_from_runtime_generic_annotations():
    """
    Parse a module through the normal stub parser stack and verify that nested
    references to names from the current module are rendered as local names.
    """

    parser = stub_parser_from_args(make_args())
    module = parser.handle_module(QualifiedName.from_str("impactx.MADXParser"), make_module())

    assert module is not None

    functions = {function.name: function for function in module.functions}
    tokenize = functions["tokenize"]
    rendered = Printer(invalid_expr_as_ellipses=False).print_function(tokenize)

    assert str(tokenize.args[0].annotation) == "list[Token]"
    assert str(tokenize.returns) == "dict[str, Expression]"
    assert rendered[0] == (
        "def tokenize(tokens: list[Token], expr: Expression | None = None) "
        "-> dict[str, Expression]:"
    )
```
</details>

Author: ax3l

pybind11_stubgen/parser/mixins/fix.py

@@ -487,8 +487,8 @@ def handle_attribute(self, path: QualifiedName, attr: Any) -> Attribute | None:
         result = super().handle_attribute(path, attr)
         if result is None:
             return None
-        if isinstance(result.annotation, ResolvedType):
-            result.annotation.name = self._strip_current_module(result.annotation.name)
+        if result.annotation is not None:
+            result.annotation = self._strip_current_module_prefix(result.annotation)
         return result
 
     def handle_module(
@@ -516,9 +516,31 @@ def parse_annotation_str(
         self, annotation_str: str
     ) -> ResolvedType | InvalidExpression | Value:
         result = super().parse_annotation_str(annotation_str)
-        if isinstance(result, ResolvedType):
-            result.name = self._strip_current_module(result.name)
-        return result
+        return self._strip_current_module_prefix(result)
+
+    def _strip_current_module_prefix(
+        self, annotation: ResolvedType | InvalidExpression | Value
+    ) -> ResolvedType | InvalidExpression | Value:
+        """
+        Strip the current module prefix from all resolved type names in an
+        annotation tree.
+
+        Python may evaluate local annotations such as ``list[Token]`` into
+        runtime generics like ``list[mymodule.MyClass.Token]``. The outer
+        container type (``list`` / ``typing.Optional`` / ``dict``) is valid,
+        but nested parameters that point back into the current module should be
+        rendered as local names in the generated stub.
+        """
+        if not isinstance(annotation, ResolvedType):
+            return annotation
+
+        annotation.name = self._strip_current_module(annotation.name)
+        if annotation.parameters is not None:
+            annotation.parameters = [
+                self._strip_current_module_prefix(parameter)
+                for parameter in annotation.parameters
+            ]
+        return annotation
 
     def _strip_current_module(self, name: QualifiedName) -> QualifiedName:
         if name[: len(self.__current_module)] == self.__current_module:

pybind11_stubgen/parser/mixins/parse.py

@@ -5,6 +5,7 @@
 import re
 import sys
 import types
+import typing
 from typing import Any, Callable, TypeVar
 
 from pybind11_stubgen.parser.errors import (
@@ -296,12 +297,12 @@ def handle_function(self, path: QualifiedName, func: Any) -> list[Function]:
                     func_args[arg_name].annotation = self.parse_annotation_str(
                         annotation
                     )
-                elif not isinstance(annotation, type):
-                    func_args[arg_name].annotation = self.handle_value(annotation)
                 elif self._is_generic_alias(annotation):
                     func_args[arg_name].annotation = self.parse_annotation_str(
                         str(annotation)
                     )
+                elif not isinstance(annotation, type):
+                    func_args[arg_name].annotation = self.handle_value(annotation)
                 else:
                     func_args[arg_name].annotation = ResolvedType(
                         name=self.handle_type(annotation),
@@ -335,9 +336,11 @@ def handle_function(self, path: QualifiedName, func: Any) -> list[Function]:
 
     def _is_generic_alias(self, annotation: type) -> bool:
         generic_alias_t: type | None = getattr(types, "GenericAlias", None)
-        if generic_alias_t is None:
-            return False
-        return isinstance(annotation, generic_alias_t)
+        return (
+            generic_alias_t is not None
+            and isinstance(annotation, generic_alias_t)
+            or typing.get_origin(annotation) is not None
+        )
 
     def handle_import(self, path: QualifiedName, origin: Any) -> Import | None:
         full_name = self._get_full_name(path, origin)

tests/py-demo/demo/pure_python/functions_3_8_plus.py

@@ -1,6 +1,14 @@
 import typing
 
 
+class Token:
+    pass
+
+
+class Expression:
+    pass
+
+
 def args_mix(
     a: int,
     b: float = 0.5,
@@ -11,3 +19,8 @@ def args_mix(
     y=int,
     **kwargs: typing.Dict[int, str],
 ): ...
+
+
+def nested_current_module_annotations(
+    tokens: list[Token], expr: typing.Optional[Expression] = None
+) -> dict[str, Expression]: ...

tests/stubs/python-3.11/pybind11-v2.11/numpy-array-wrap-with-annotated/demo/pure_python/functions_3_8_plus.pyi

@@ -2,7 +2,19 @@ from __future__ import annotations
 
 import typing as typing
 
-__all__: list[str] = ["args_mix", "typing"]
+__all__: list[str] = [
+    "Expression",
+    "Token",
+    "args_mix",
+    "nested_current_module_annotations",
+    "typing",
+]
+
+class Token:
+    pass
+
+class Expression:
+    pass
 
 def args_mix(
     a: int,
@@ -11,5 +23,8 @@ def args_mix(
     *args: int,
     x: int = 1,
     y=int,
-    **kwargs: typing.Dict[int, str],
+    **kwargs: dict[int, str],
 ): ...
+def nested_current_module_annotations(
+    tokens: list[Token], expr: Expression | None = None
+) -> dict[str, Expression]: ...

tests/stubs/python-3.11/pybind11-v2.12/numpy-array-wrap-with-annotated/demo/pure_python/functions_3_8_plus.pyi

@@ -2,7 +2,19 @@ from __future__ import annotations
 
 import typing as typing
 
-__all__: list[str] = ["args_mix", "typing"]
+__all__: list[str] = [
+    "Expression",
+    "Token",
+    "args_mix",
+    "nested_current_module_annotations",
+    "typing",
+]
+
+class Token:
+    pass
+
+class Expression:
+    pass
 
 def args_mix(
     a: int,
@@ -11,5 +23,8 @@ def args_mix(
     *args: int,
     x: int = 1,
     y=int,
-    **kwargs: typing.Dict[int, str],
+    **kwargs: dict[int, str],
 ): ...
+def nested_current_module_annotations(
+    tokens: list[Token], expr: Expression | None = None
+) -> dict[str, Expression]: ...

tests/stubs/python-3.11/pybind11-v2.13/numpy-array-use-type-var/demo/pure_python/functions_3_8_plus.pyi

@@ -2,7 +2,19 @@ from __future__ import annotations
 
 import typing as typing
 
-__all__: list[str] = ["args_mix", "typing"]
+__all__: list[str] = [
+    "Expression",
+    "Token",
+    "args_mix",
+    "nested_current_module_annotations",
+    "typing",
+]
+
+class Token:
+    pass
+
+class Expression:
+    pass
 
 def args_mix(
     a: int,
@@ -11,5 +23,8 @@ def args_mix(
     *args: int,
     x: int = 1,
     y=int,
-    **kwargs: typing.Dict[int, str],
+    **kwargs: dict[int, str],
 ): ...
+def nested_current_module_annotations(
+    tokens: list[Token], expr: Expression | None = None
+) -> dict[str, Expression]: ...

tests/stubs/python-3.11/pybind11-v2.13/numpy-array-wrap-with-annotated/demo/pure_python/functions_3_8_plus.pyi

@@ -2,7 +2,19 @@ from __future__ import annotations
 
 import typing as typing
 
-__all__: list[str] = ["args_mix", "typing"]
+__all__: list[str] = [
+    "Expression",
+    "Token",
+    "args_mix",
+    "nested_current_module_annotations",
+    "typing",
+]
+
+class Token:
+    pass
+
+class Expression:
+    pass
 
 def args_mix(
     a: int,
@@ -11,5 +23,8 @@ def args_mix(
     *args: int,
     x: int = 1,
     y=int,
-    **kwargs: typing.Dict[int, str],
+    **kwargs: dict[int, str],
 ): ...
+def nested_current_module_annotations(
+    tokens: list[Token], expr: Expression | None = None
+) -> dict[str, Expression]: ...

tests/stubs/python-3.11/pybind11-v2.9/numpy-array-wrap-with-annotated/demo/pure_python/functions_3_8_plus.pyi

@@ -2,7 +2,19 @@ from __future__ import annotations
 
 import typing as typing
 
-__all__: list[str] = ["args_mix", "typing"]
+__all__: list[str] = [
+    "Expression",
+    "Token",
+    "args_mix",
+    "nested_current_module_annotations",
+    "typing",
+]
+
+class Token:
+    pass
+
+class Expression:
+    pass
 
 def args_mix(
     a: int,
@@ -11,5 +23,8 @@ def args_mix(
     *args: int,
     x: int = 1,
     y=int,
-    **kwargs: typing.Dict[int, str],
+    **kwargs: dict[int, str],
 ): ...
+def nested_current_module_annotations(
+    tokens: list[Token], expr: Expression | None = None
+) -> dict[str, Expression]: ...

tests/stubs/python-3.11/pybind11-v3.0/numpy-array-use-type-var/demo/pure_python/functions_3_8_plus.pyi

@@ -2,7 +2,19 @@ from __future__ import annotations
 
 import typing as typing
 
-__all__: list[str] = ["args_mix", "typing"]
+__all__: list[str] = [
+    "Expression",
+    "Token",
+    "args_mix",
+    "nested_current_module_annotations",
+    "typing",
+]
+
+class Token:
+    pass
+
+class Expression:
+    pass
 
 def args_mix(
     a: int,
@@ -11,5 +23,8 @@ def args_mix(
     *args: int,
     x: int = 1,
     y=int,
-    **kwargs: typing.Dict[int, str],
+    **kwargs: dict[int, str],
 ): ...
+def nested_current_module_annotations(
+    tokens: list[Token], expr: Expression | None = None
+) -> dict[str, Expression]: ...

tests/stubs/python-3.11/pybind11-v3.0/numpy-array-wrap-with-annotated/demo/pure_python/functions_3_8_plus.pyi

@@ -2,7 +2,19 @@ from __future__ import annotations
 
 import typing as typing
 
-__all__: list[str] = ["args_mix", "typing"]
+__all__: list[str] = [
+    "Expression",
+    "Token",
+    "args_mix",
+    "nested_current_module_annotations",
+    "typing",
+]
+
+class Token:
+    pass
+
+class Expression:
+    pass
 
 def args_mix(
     a: int,
@@ -11,5 +23,8 @@ def args_mix(
     *args: int,
     x: int = 1,
     y=int,
-    **kwargs: typing.Dict[int, str],
+    **kwargs: dict[int, str],
 ): ...
+def nested_current_module_annotations(
+    tokens: list[Token], expr: Expression | None = None
+) -> dict[str, Expression]: ...