Skip to content

Commit

Permalink
stubgen.py: handle method aliases similarly to function aliases
Browse files Browse the repository at this point in the history
  • Loading branch information
wojdyr committed Sep 26, 2024
1 parent 8ce0dee commit bb4de49
Show file tree
Hide file tree
Showing 3 changed files with 13 additions and 2 deletions.
10 changes: 9 additions & 1 deletion src/stubgen.py
Original file line number Diff line number Diff line change
Expand Up @@ -376,6 +376,14 @@ def put_nb_func(self, fn: NbFunction, name: Optional[str] = None) -> None:
self.write_ln(f"@{overload}")
self.put_nb_overload(fn, s, name)

def put_nb_method(self, fn: NbFunction, name: Optional[str] = None) -> None:
fn_name = getattr(fn, "__name__", None)
# Check if this function is an alias
if name and fn_name and name != fn_name:
self.write_ln(f"{name} = {fn_name}\n")
return
self.put_nb_func(fn, name)

def put_function(self, fn: Callable[..., Any], name: Optional[str] = None, parent: Optional[object] = None):
"""Append a function of an arbitrary type to the stub"""
# Don't generate a constructor for nanobind classes that aren't constructible
Expand Down Expand Up @@ -848,7 +856,7 @@ def put(self, value: object, name: Optional[str] = None, parent: Optional[object
elif tp_mod == "nanobind":
if tp_name == "nb_method":
value = cast(NbFunction, value)
self.put_nb_func(value, name)
self.put_nb_method(value, name)
elif tp_name == "nb_static_property":
value = cast(NbStaticProperty, value)
self.put_nb_static_property(name, value)
Expand Down
3 changes: 2 additions & 1 deletion tests/test_typing.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -57,9 +57,10 @@ NB_MODULE(test_typing_ext, m) {

m.def("makeNestedClass", [] { return NestedClass(); });

// Aliases to local functoins and types
// Aliases to functions and types
m.attr("FooAlias") = m.attr("Foo");
m.attr("f_alias") = m.attr("f");
nb::type<Foo>().attr("lt_alias") = nb::type<Foo>().attr("__lt__");

// Custom signature generation for classes and methods
struct CustomSignature { int value; };
Expand Down
2 changes: 2 additions & 0 deletions tests/test_typing_ext.pyi.ref
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,8 @@ class Foo:

def __ge__(self, arg: Foo, /) -> bool: ...

lt_alias = __lt__

FooAlias: TypeAlias = Foo

T = TypeVar("T", contravariant=True)
Expand Down

0 comments on commit bb4de49

Please sign in to comment.