feat(stubgen): allow __suffix__ and __prefix__ for classes (#1235)

It is currently not possible to *add* functions to a class using
patterns. This may be necessary if functions are added to a class
through mechanisms other than nanobind. This PR adds the
possibility to match on `__suffix__` and `__prefix__` within a
class, like it is possible to do within modules already, which
can be used to add anything to the beginning or end of a class using
patterns.

Signed-off-by: Ingo Müller <[email protected]>

Author: ingomueller-net

docs/typing.rst

@@ -714,6 +714,6 @@ you may use the special ``\from`` escape code to import them:
        def lookup(array: Array[T], index: Literal[0] = 0) -> _Opt[T]:
            \doc
 
-You may also add free-form text the beginning or the end of the generated stub.
-To do so, add an entry that matches on ``module_name.__prefix__`` or
-``module_name.__suffix__``.
+You may also add free-form text the beginning or the end of the generated stub
+module or of a class. To do so, add an entry that matches on ``name.__prefix__``
+or ``name.__suffix__`` where ``name`` is the name of the module or class.

src/stubgen.py

@@ -567,8 +567,10 @@ def put_type(self, tp: NbType, name: Optional[str]):
                 self.put_docstr(docstr)
                 if len(tp_dict):
                     self.write("\n")
+            self.apply_pattern(self.prefix + ".__prefix__", None)
             for k, v in tp_dict.items():
                 self.put(v, k, tp)
+            self.apply_pattern(self.prefix + ".__suffix__", None)
             if output_len == len(self.output):
                 self.write_ln("pass\n")
             self.depth -= 1

tests/pattern_file.nb

@@ -17,3 +17,9 @@ test_typing_ext.__prefix__:
 
 test_typing_ext.__suffix__:
     # a suffix
+
+test_typing_ext.Foo.__prefix__:
+    # a class prefix
+
+test_typing_ext.Foo.__suffix__:
+    # a class suffix

tests/test_typing_ext.pyi.ref

@@ -9,6 +9,8 @@ from .submodule import F as F, f as f2
 # a prefix
 
 class Foo:
+    # a class prefix
+
     def __lt__(self, arg: int, /) -> bool: ...
 
     def __gt__(self, arg: int, /) -> bool: ...
@@ -19,6 +21,8 @@ class Foo:
 
     lt_alias = __lt__
 
+    # a class suffix
+
 def f() -> None: ...
 
 def makeNestedClass() -> py_stub_test.AClass.NestedClass: ...