Skip to content

Commit aeae256

Browse files
authored
A few minor improvements to the stubgen (wjakob#1179)
* Generalized support for `types.ModuleType` to a few other interpreter internal types. * Fixed rendering of ellipsis in expression context, e.g. `Alias = tuple[T, ...]`.
1 parent db637da commit aeae256

File tree

3 files changed

+23
-8
lines changed

3 files changed

+23
-8
lines changed

src/stubgen.py

Lines changed: 17 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -93,6 +93,16 @@ class and repeatedly call ``.put()`` to register modules or contents within the
9393
"__firstlineno__", "__static_attributes__", "__annotations__", "__annotate__",
9494
"__annotate_func__"
9595
]
96+
97+
# Interpreter-internal types.
98+
TYPES_TYPES = {
99+
getattr(types, name): name for name in [
100+
"MethodDescriptorType",
101+
"MemberDescriptorType",
102+
"ModuleType",
103+
]
104+
}
105+
96106
# fmt: on
97107

98108
# This type is used to track per-module imports (``import name as desired_name``)
@@ -988,10 +998,9 @@ def expr_str(self, e: Any, abbrev: bool = True) -> Optional[str]:
988998
complicated.
989999
"""
9901000
tp = type(e)
991-
for t in [bool, int, type(None), type(builtins.Ellipsis)]:
992-
if issubclass(tp, t):
993-
return repr(e)
994-
if issubclass(tp, float):
1001+
if issubclass(tp, (bool, int, type(None), type(builtins.Ellipsis))):
1002+
return repr(e)
1003+
elif issubclass(tp, float):
9951004
s = repr(e)
9961005
if "inf" in s or "nan" in s:
9971006
return f"float('{s}')"
@@ -1140,8 +1149,10 @@ def type_str(self, tp: Union[List[Any], Tuple[Any, ...], Dict[Any, Any], Any]) -
11401149
+ ", ".join(args_gen)
11411150
+ "]"
11421151
)
1143-
elif tp is types.ModuleType:
1144-
result = "types.ModuleType"
1152+
elif tp in TYPES_TYPES:
1153+
result = f"types.{TYPES_TYPES[tp]}"
1154+
elif tp is Ellipsis:
1155+
result = "..."
11451156
elif isinstance(tp, type):
11461157
result = tp.__module__ + "." + tp.__qualname__
11471158
else:

tests/test_typing.cpp

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -57,7 +57,9 @@ NB_MODULE(test_typing_ext, m) {
5757

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

60-
// Aliases to local functoins and types
60+
m.attr("AnyTuple") = nb::typing().attr("Tuple")[nb::make_tuple(nb::any_type(), nb::ellipsis())];
61+
62+
// Aliases to local functions and types
6163
m.attr("FooAlias") = m.attr("Foo");
6264
m.attr("f_alias") = m.attr("f");
6365
nb::type<Foo>().attr("lt_alias") = nb::type<Foo>().attr("__lt__");

tests/test_typing_ext.pyi.ref

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
from collections.abc import Iterable
22
import py_stub_test
3-
from typing import Generic, Optional, Self, TypeAlias, TypeVar
3+
from typing import Any, Generic, Optional, Self, TypeAlias, TypeVar
44

55
from . import submodule as submodule
66
from .submodule import F as F, f as f2
@@ -23,6 +23,8 @@ def f() -> None: ...
2323

2424
def makeNestedClass() -> py_stub_test.AClass.NestedClass: ...
2525

26+
AnyTuple: TypeAlias = tuple[Any, ...]
27+
2628
FooAlias: TypeAlias = Foo
2729

2830
f_alias = f

0 commit comments

Comments
 (0)