diff --git a/cmake/nanobind-config.cmake b/cmake/nanobind-config.cmake index bc680e57..8e204d63 100644 --- a/cmake/nanobind-config.cmake +++ b/cmake/nanobind-config.cmake @@ -590,7 +590,7 @@ endfunction() # --------------------------------------------------------------------------- function (nanobind_add_stub name) - cmake_parse_arguments(PARSE_ARGV 1 ARG "VERBOSE;INCLUDE_PRIVATE;EXCLUDE_DOCSTRINGS;INSTALL_TIME;RECURSIVE;EXCLUDE_FROM_ALL" "MODULE;COMPONENT;PATTERN_FILE;OUTPUT_PATH" "PYTHON_PATH;DEPENDS;MARKER_FILE;OUTPUT") + cmake_parse_arguments(PARSE_ARGV 1 ARG "VERBOSE;INCLUDE_PRIVATE;EXCLUDE_DOCSTRINGS;EXCLUDE_VALUES;INSTALL_TIME;RECURSIVE;EXCLUDE_FROM_ALL" "MODULE;COMPONENT;PATTERN_FILE;OUTPUT_PATH" "PYTHON_PATH;DEPENDS;MARKER_FILE;OUTPUT") if (EXISTS ${NB_DIR}/src/stubgen.py) set(NB_STUBGEN "${NB_DIR}/src/stubgen.py") @@ -614,6 +614,10 @@ function (nanobind_add_stub name) list(APPEND NB_STUBGEN_ARGS -D) endif() + if (ARG_EXCLUDE_VALUES) + list(APPEND NB_STUBGEN_ARGS --exclude-values) + endif() + if (ARG_RECURSIVE) list(APPEND NB_STUBGEN_ARGS -r) endif() diff --git a/docs/typing.rst b/docs/typing.rst index 92b5af55..6c9d8a6f 100644 --- a/docs/typing.rst +++ b/docs/typing.rst @@ -540,7 +540,7 @@ The program has the following command line options: .. code-block:: text usage: python -m nanobind.stubgen [-h] [-o FILE] [-O PATH] [-i PATH] [-m MODULE] - [-r] [-M FILE] [-P] [-D] [-q] + [-r] [-M FILE] [-P] [-D] [--exclude-values] [-q] Generate stubs for nanobind-based extensions. @@ -559,6 +559,7 @@ The program has the following command line options: -P, --include-private include private members (with single leading or trailing underscore) -D, --exclude-docstrings exclude docstrings from the generated stub + --exclude-values force the use of ... for values -q, --quiet do not generate any output in the absence of failures diff --git a/src/stubgen.py b/src/stubgen.py index 61a650f3..1b591cc0 100755 --- a/src/stubgen.py +++ b/src/stubgen.py @@ -1004,12 +1004,14 @@ def expr_str(self, e: Any, abbrev: bool = True) -> Optional[str]: """ tp = type(e) if issubclass(tp, (bool, int, type(None), type(builtins.Ellipsis))): - return repr(e) + s = repr(e) + if len(s) < self.max_expr_length or not abbrev: + return s elif issubclass(tp, float): s = repr(e) if "inf" in s or "nan" in s: - return f"float('{s}')" - else: + s = f"float('{s}')" + if len(s) < self.max_expr_length or not abbrev: return s elif issubclass(tp, type) or typing.get_origin(e): return self.type_str(e) @@ -1025,13 +1027,17 @@ def expr_str(self, e: Any, abbrev: bool = True) -> Optional[str]: tv = self.import_object("typing", "TypeVar") s = f'{tv}("{e.__name__}"' for v in getattr(e, "__constraints__", ()): - v = self.expr_str(v) + v = self.type_str(v) assert v s += ", " + v - for k in ["contravariant", "covariant", "bound", "infer_variance"]: + if v := getattr(e, "__bound__", None): + v = self.type_str(v) + assert v + s += ", bound=" + v + for k in ["contravariant", "covariant", "infer_variance"]: v = getattr(e, f"__{k}__", None) if v: - v = self.expr_str(v) + v = self.expr_str(v, abbrev=False) if v is None: return None s += f", {k}=" + v @@ -1319,6 +1325,14 @@ def parse_options(args: List[str]) -> argparse.Namespace: help="exclude docstrings from the generated stub", ) + parser.add_argument( + "--exclude-values", + dest="exclude_values", + default=False, + action="store_true", + help="force the use of ... for values", + ) + parser.add_argument( "-q", "--quiet", @@ -1463,6 +1477,7 @@ def main(args: Optional[List[str]] = None) -> None: recursive=opt.recursive, include_docstrings=opt.include_docstrings, include_private=opt.include_private, + max_expr_length=0 if opt.exclude_values else 50, patterns=patterns, output_file=file ) diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 44120600..5783b39b 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -104,6 +104,7 @@ foreach (NAME functions classes ndarray jax tensorflow stl enum typing make_iter set(EXTRA MARKER_FILE py.typed PATTERN_FILE "${CMAKE_CURRENT_SOURCE_DIR}/pattern_file.nb" + EXCLUDE_VALUES ) set(EXTRA_DEPENDS "${OUT_DIR}/py_stub_test.py") else() diff --git a/tests/test_typing.cpp b/tests/test_typing.cpp index 0793d054..b6a24324 100644 --- a/tests/test_typing.cpp +++ b/tests/test_typing.cpp @@ -114,6 +114,10 @@ NB_MODULE(test_typing_ext, m) { m.def("list_front", [](nb::list l) { return l[0]; }, nb::sig("def list_front[T](arg: list[T], /) -> T")); + // Type variables with constraints and a bound. + m.attr("T2") = nb::type_var("T2", "bound"_a = nb::type()); + m.attr("T3") = nb::type_var("T3", *nb::make_tuple(nb::type(), nb::type())); + // Some statements that will be modified by the pattern file m.def("remove_me", []{}); m.def("tweak_me", [](nb::object o) { return o; }, "prior docstring\nremains preserved"); diff --git a/tests/test_typing_ext.pyi.ref b/tests/test_typing_ext.pyi.ref index 4b2533b0..d0eb3374 100644 --- a/tests/test_typing_ext.pyi.ref +++ b/tests/test_typing_ext.pyi.ref @@ -44,7 +44,7 @@ class CustomSignature(Iterable[int]): def value(self, value: Optional[int], /) -> None: """docstring for setter""" -pytree: dict = {'a' : ('b', [123])} +pytree: dict = ... T = TypeVar("T", contravariant=True) @@ -63,6 +63,10 @@ class WrapperTypeParam[T]: def list_front[T](arg: list[T], /) -> T: ... +T2 = TypeVar("T2", bound=Foo) + +T3 = TypeVar("T3", Foo, Wrapper) + def tweak_me(arg: int): """ prior docstring