fix: 🐛 do not remove self parameter annotation when types do not match (#195)

Author: ringohoffman

pybind11_stubgen/parser/mixins/fix.py

@@ -40,25 +40,46 @@
 
 
 class RemoveSelfAnnotation(IParser):
+
+    __any_t_name = QualifiedName.from_str("Any")
+    __typing_any_t_name = QualifiedName.from_str("typing.Any")
+
     def handle_method(self, path: QualifiedName, method: Any) -> list[Method]:
         methods = super().handle_method(path, method)
         for method in methods:
-            self._remove_self_arg_annotation(method.function)
+            self._remove_self_arg_annotation(path, method.function)
         return methods
 
     def handle_property(self, path: QualifiedName, prop: Any) -> Property | None:
         prop = super().handle_property(path, prop)
         if prop is not None:
             if prop.getter is not None:
-                self._remove_self_arg_annotation(prop.getter)
+                self._remove_self_arg_annotation(path, prop.getter)
             if prop.setter is not None:
-                self._remove_self_arg_annotation(prop.setter)
+                self._remove_self_arg_annotation(path, prop.setter)
 
         return prop
 
-    def _remove_self_arg_annotation(self, func: Function):
-        if len(func.args) > 0 and func.args[0].name == "self":
-            func.args[0].annotation = None
+    def _remove_self_arg_annotation(self, path: QualifiedName, func: Function) -> None:
+        if len(func.args) == 0:
+            return
+        fully_qualified_class_name = QualifiedName(path[:-1])  # remove the method name
+        first_arg = func.args[0]
+        if (
+            first_arg.name == "self"
+            and isinstance(first_arg.annotation, ResolvedType)
+            and not first_arg.annotation.parameters
+            and (
+                first_arg.annotation.name
+                in {
+                    self.__any_t_name,
+                    self.__typing_any_t_name,
+                    fully_qualified_class_name,
+                    fully_qualified_class_name[-len(first_arg.annotation.name) :],
+                }
+            )
+        ):
+            first_arg.annotation = None
 
 
 class FixMissingImports(IParser):

tests/py-demo/bindings/src/modules/aliases.cpp

@@ -16,6 +16,8 @@ namespace {
     };
     struct Bar3 {
     };
+    struct Bar4 {
+    };
 } // namespace
 
 void bind_aliases_module(py::module &&m) {
@@ -57,13 +59,18 @@ void bind_aliases_module(py::module &&m) {
     {
         auto &&sub = m.def_submodule("foreign_method_arg");
         auto &&pyBar = py::class_<Bar2>(sub, "Bar2");
-        pyBar.def("set_foo", [](const demo::Foo &) { return 13; });
+        pyBar.def("set_foo", [](const Bar2 &, const demo::Foo &) { return 13; });
     }
     {
         auto &&sub = m.def_submodule("foreign_method_return");
         auto &&pyBar = py::class_<Bar3>(sub, "Bar3");
         pyBar.def("get_foo", []() { return demo::Foo(); });
     }
+    {
+        auto &&sub = m.def_submodule("missing_self_arg");
+        auto &&pyBar = py::class_<Bar4>(sub, "Bar4");
+        pyBar.def("set_foo", [](const demo::Foo &) { return 13; });
+    }
 
     {
         m.attr("foreign_type_alias") = m.attr("foreign_method_arg").attr("Bar2");

tests/stubs/python-3.12/pybind11-master/numpy-array-use-type-var/demo/_bindings/aliases/__init__.pyi

@@ -16,6 +16,7 @@ from . import (
     foreign_method_arg,
     foreign_method_return,
     foreign_return,
+    missing_self_arg,
 )
 
 __all__ = [
@@ -33,6 +34,7 @@ __all__ = [
     "func",
     "local_func_alias",
     "local_type_alias",
+    "missing_self_arg",
     "random",
 ]
 

tests/stubs/python-3.12/pybind11-master/numpy-array-use-type-var/demo/_bindings/aliases/foreign_method_arg.pyi

@@ -5,4 +5,4 @@ import demo._bindings.classes
 __all__ = ["Bar2"]
 
 class Bar2:
-    def set_foo(self) -> int: ...
+    def set_foo(self, arg0: demo._bindings.classes.Foo) -> int: ...

tests/stubs/python-3.12/pybind11-master/numpy-array-use-type-var/demo/_bindings/aliases/missing_self_arg.pyi

@@ -0,0 +1,8 @@
+from __future__ import annotations
+
+import demo._bindings.classes
+
+__all__ = ["Bar4"]
+
+class Bar4:
+    def set_foo(self: demo._bindings.classes.Foo) -> int: ...

tests/stubs/python-3.12/pybind11-master/numpy-array-wrap-with-annotated/demo/_bindings/aliases/__init__.pyi

@@ -16,6 +16,7 @@ from . import (
     foreign_method_arg,
     foreign_method_return,
     foreign_return,
+    missing_self_arg,
 )
 
 __all__ = [
@@ -33,6 +34,7 @@ __all__ = [
     "func",
     "local_func_alias",
     "local_type_alias",
+    "missing_self_arg",
     "random",
 ]
 

tests/stubs/python-3.12/pybind11-master/numpy-array-wrap-with-annotated/demo/_bindings/aliases/foreign_method_arg.pyi

@@ -5,4 +5,4 @@ import demo._bindings.classes
 __all__ = ["Bar2"]
 
 class Bar2:
-    def set_foo(self) -> int: ...
+    def set_foo(self, arg0: demo._bindings.classes.Foo) -> int: ...

tests/stubs/python-3.12/pybind11-master/numpy-array-wrap-with-annotated/demo/_bindings/aliases/missing_self_arg.pyi

@@ -0,0 +1,8 @@
+from __future__ import annotations
+
+import demo._bindings.classes
+
+__all__ = ["Bar4"]
+
+class Bar4:
+    def set_foo(self: demo._bindings.classes.Foo) -> int: ...

tests/stubs/python-3.12/pybind11-v2.11/numpy-array-wrap-with-annotated/demo/_bindings/aliases/__init__.pyi

@@ -16,6 +16,7 @@ from . import (
     foreign_method_arg,
     foreign_method_return,
     foreign_return,
+    missing_self_arg,
 )
 
 __all__ = [
@@ -33,6 +34,7 @@ __all__ = [
     "func",
     "local_func_alias",
     "local_type_alias",
+    "missing_self_arg",
     "random",
 ]
 

tests/stubs/python-3.12/pybind11-v2.11/numpy-array-wrap-with-annotated/demo/_bindings/aliases/foreign_method_arg.pyi

@@ -5,4 +5,4 @@ import demo._bindings.classes
 __all__ = ["Bar2"]
 
 class Bar2:
-    def set_foo(self) -> int: ...
+    def set_foo(self, arg0: demo._bindings.classes.Foo) -> int: ...