experimental-inspect: emit __all__ in the generated stubs

guide/src/type-stub.md

@@ -55,6 +55,8 @@ will generate the following stub file:
 ```python
 import typing
 
+__all__ = ["CONSTANT", "Class", "list_of_int_identity"]
+
 CONSTANT: typing.Final = "FOO"
 
 
@@ -89,3 +91,4 @@ PyO3 also provides the smaller `pyo3-introspection` binary that allows to genera
 - `FromPyObject::INPUT_TYPE` and `IntoPyObject::OUTPUT_TYPE` must be implemented for PyO3 to get the proper input/output type annotations to use.
 - PyO3 is not able to introspect the content of `#[pymodule]` and `#[pymodule_init]` functions.
   If they are present, the module is tagged as incomplete using a fake `def __getattr__(name: str) -> Incomplete: ...` function [following best practices](https://typing.python.org/en/latest/guides/writing_stubs.html#incomplete-stubs).
+  Such a module also gets no `__all__` declaration, because the set of names it exports at runtime is not known.

newsfragments/6242.fixed.md

@@ -0,0 +1,2 @@
+`experimental-inspect`: generate the `__all__` declaration of complete modules in the type stubs, matching the `__all__` that `#[pymodule]` builds at runtime.
+It is placed directly after the imports, following the conventions used by `typeshed`, and lists the members in the order the stub declares them.

noxfile.py

@@ -1700,6 +1700,8 @@ def test_introspection(session: nox.Session):
     with tempfile.TemporaryDirectory() as stub_dir:
         session.install("maturin")
         session.install("ruff")
+        # `pytest-benchmark` backs the `--benchmark-disable` in `pytests/pyproject.toml`
+        session.install("pytest", "pytest-benchmark")
         options = []
         target = os.environ.get("CARGO_BUILD_TARGET")
         if target is not None:
@@ -1735,6 +1737,9 @@ def test_introspection(session: nox.Session):
         )
         _run(session, "ruff", "format", stub_dir)
         _ensure_directory_equals(Path(stub_dir), Path("pytests/stubs"))
+        # The stubs only match a module built with the features used above, so this test
+        # is skipped by the regular `pytests` session and runs here instead.
+        _run(session, "pytest", "pytests/tests/test_stubs.py")
 
 
 def _ensure_directory_equals(expected_dir: Path, actual_dir: Path):

pyo3-introspection/src/stubs.rs

@@ -92,24 +92,28 @@ fn module_stubs(module: &Module, parents: &[&str]) -> String {
         ));
     }
 
+    let dunder_all = dunder_all_stubs(module, &imports);
+
     let mut final_elements = Vec::new();
     if let Some(docstring) = &module.docstring {
         final_elements.push(format!("\"\"\"\n{docstring}\n\"\"\""));
     }
     final_elements.extend(imports.imports);
+    final_elements.extend(dunder_all);
     final_elements.extend(elements);
 
     let mut output = String::new();
 
-    // We insert two line jumps (i.e. empty strings) only above and below multiple line elements (classes with methods, functions with decorators)
+    // We insert two line jumps (i.e. empty strings) only above and below multiple line elements
+    // (classes with methods, functions with decorators) and the `__all__` declaration
     for element in final_elements {
-        let is_multiline = element.contains('\n');
-        if is_multiline && !output.is_empty() && !output.ends_with("\n\n") {
+        let needs_empty_lines = element.contains('\n') || element.starts_with("__all__");
+        if needs_empty_lines && !output.is_empty() && !output.ends_with("\n\n") {
             output.push('\n');
         }
         output.push_str(&element);
         output.push('\n');
-        if is_multiline {
+        if needs_empty_lines {
             output.push('\n');
         }
     }
@@ -121,6 +125,54 @@ fn module_stubs(module: &Module, parents: &[&str]) -> String {
     output
 }
 
+/// Generates the `__all__` declaration of a module, if we are able to write an accurate one.
+///
+/// [`PyModuleMethods::add`] appends every name it adds to the module `__all__`, so any `#[pymodule]`
+/// with at least one member has an `__all__` at runtime and the stub must declare it too.
+///
+/// We only emit it for complete modules: for an incomplete one we do not know the full list of
+/// members, and an `__all__` missing some of them would hide from type checkers names that do exist
+/// at runtime.
+///
+/// [`PyModuleMethods::add`]: https://docs.rs/pyo3/latest/pyo3/types/trait.PyModuleMethods.html#tymethod.add
+fn dunder_all_stubs(module: &Module, imports: &Imports) -> Option<String> {
+    if module.incomplete {
+        return None;
+    }
+    if module.attributes.iter().any(|a| a.name == "__all__") {
+        // The introspection data carries an explicit `__all__`, it is more accurate than ours
+        return None;
+    }
+    let member_count = module.attributes.len()
+        + module.classes.len()
+        + module.functions.len()
+        + module.modules.len();
+    if member_count == 0 {
+        // Nothing was ever added to the module, so it has no `__all__` at runtime either
+        return None;
+    }
+
+    // Each of these lists is already sorted by name, so we just list the members in the order the
+    // stub declares them, with the submodules (declared in their own file) last.
+    let mut elts = Vec::with_capacity(member_count);
+    elts.extend(
+        module
+            .attributes
+            .iter()
+            .map(|a| &a.name)
+            .chain(module.classes.iter().map(|c| &c.name))
+            .chain(module.functions.iter().map(|f| &f.name))
+            .chain(module.modules.iter().map(|m| &m.name))
+            .map(|name| Expr::Constant {
+                value: Constant::Str(name.clone()),
+            }),
+    );
+
+    let mut buffer = "__all__ = ".to_string();
+    imports.serialize_expr(&Expr::List { elts }, &mut buffer);
+    Some(buffer)
+}
+
 fn class_stubs(class: &Class, imports: &Imports) -> String {
     let mut buffer = String::new();
     for decorator in &class.decorators {
@@ -905,6 +957,67 @@ mod tests {
         assert_eq!(output, "dict[A, (A3.C, A3.D, B, A2, int, int2, float)]");
     }
 
+    fn empty_module(name: &str) -> Module {
+        Module {
+            name: name.into(),
+            modules: Vec::new(),
+            classes: Vec::new(),
+            functions: Vec::new(),
+            attributes: Vec::new(),
+            incomplete: false,
+            docstring: None,
+        }
+    }
+
+    /// The `pytests` stubs cover the common cases, we only test the ones they don't have:
+    /// a module with a submodule, and an empty module.
+    #[test]
+    fn test_dunder_all() {
+        let module = Module {
+            modules: vec![empty_module("sub")],
+            classes: vec![Class {
+                name: "Zulu".into(),
+                bases: Vec::new(),
+                methods: Vec::new(),
+                attributes: Vec::new(),
+                decorators: Vec::new(),
+                inner_classes: Vec::new(),
+                docstring: None,
+            }],
+            functions: vec![Function {
+                name: "func".into(),
+                decorators: Vec::new(),
+                arguments: Arguments {
+                    positional_only_arguments: Vec::new(),
+                    arguments: Vec::new(),
+                    vararg: None,
+                    keyword_only_arguments: Vec::new(),
+                    kwarg: None,
+                },
+                returns: None,
+                is_async: false,
+                docstring: None,
+            }],
+            attributes: vec![Attribute {
+                name: "CONST".into(),
+                value: Some(Expr::Constant {
+                    value: Constant::Int("1".into()),
+                }),
+                annotation: None,
+                docstring: None,
+            }],
+            ..empty_module("bar")
+        };
+        // The names are the ones `PyModuleMethods::add` puts in `__all__` at runtime, including
+        // the submodule which is declared in its own stub file.
+        assert_eq!(
+            module_stubs(&module, &["foo"]),
+            "__all__ = [\"CONST\", \"Zulu\", \"func\", \"sub\"]\n\nCONST = 1\nclass Zulu: ...\ndef func(): ...\n"
+        );
+        // Nothing was added to an empty module, so it has no `__all__` at runtime either
+        assert_eq!(module_stubs(&empty_module("bar"), &["foo"]), "");
+    }
+
     #[test]
     fn test_make_module_path_relative() {
         assert_eq!(

pytests/noxfile.py

@@ -26,7 +26,10 @@ def try_install_binary(package: str, constraint: str):
     # - is a dependency of gevent
     try_install_binary("zope.interface", "<7")
     try_install_binary("gevent", ">=22.10.2")
-    ignored_paths = []
+    # The stubs are generated from a build with `experimental-async,experimental-inspect`,
+    # so only such a build exports the members they declare. The `test-introspection`
+    # session runs this test against one.
+    ignored_paths = ["tests/test_stubs.py"]
     if sys.version_info < (3, 10):
         # Match syntax is only available in Python >= 3.10
         ignored_paths.append("tests/test_enums_match.py")

pytests/stubs/annotations.pyi

@@ -1,5 +1,7 @@
 from .pyclasses import EmptyClass
 
+__all__ = ["cross_module_imports", "with_custom_type_annotations"]
+
 def cross_module_imports(_a: EmptyClass) -> None: ...
 def with_custom_type_annotations(
     a: "list[int]", *_args: "str", _b: "int | None" = None, **_kwargs: "bool"

pytests/stubs/awaitable.pyi

@@ -1,5 +1,7 @@
 from typing import Any, final
 
+__all__ = ["FutureAwaitable", "IterAwaitable"]
+
 @final
 class FutureAwaitable:
     def __await__(self, /) -> FutureAwaitable: ...

pytests/stubs/buf_and_str.pyi

@@ -5,6 +5,14 @@ Objects related to PyBuffer and PyStr
 from collections.abc import Sequence
 from typing import Any, final
 
+__all__ = [
+    "BytesExtractor",
+    "map_byte_cow",
+    "map_byte_slice",
+    "map_byte_vec",
+    "return_memoryview",
+]
+
 @final
 class BytesExtractor:
     """

pytests/stubs/comparisons.pyi

@@ -1,5 +1,15 @@
 from typing import final
 
+__all__ = [
+    "Eq",
+    "EqDefaultNe",
+    "EqDerived",
+    "Ordered",
+    "OrderedDefaultNe",
+    "OrderedDerived",
+    "OrderedRichCmp",
+]
+
 @final
 class Eq:
     def __eq__(self, /, other: object) -> bool: ...

pytests/stubs/consts.pyi

@@ -1,5 +1,7 @@
 from typing import Final, final
 
+__all__ = ["ESCAPING", "PI", "ClassWithConst"]
+
 ESCAPING: Final = "S\0\x01\t\n\r\"'\\"
 """
 We experiment with "escaping"

pytests/stubs/datetime.pyi

@@ -1,6 +1,25 @@
 from datetime import date, datetime, time, timedelta, tzinfo
 from typing import final
 
+__all__ = [
+    "TzClass",
+    "date_from_timestamp",
+    "datetime_from_timestamp",
+    "get_date_tuple",
+    "get_datetime_tuple",
+    "get_datetime_tuple_fold",
+    "get_datetime_tzinfo",
+    "get_delta_tuple",
+    "get_time_tuple",
+    "get_time_tuple_fold",
+    "get_time_tzinfo",
+    "make_date",
+    "make_datetime",
+    "make_delta",
+    "make_time",
+    "time_with_fold",
+]
+
 @final
 class TzClass(tzinfo):
     def __new__(cls, /) -> TzClass: ...

pytests/stubs/dict_iter.pyi

@@ -1,5 +1,7 @@
 from typing import final
 
+__all__ = ["DictSize"]
+
 @final
 class DictSize:
     def __new__(cls, /, expected: int) -> DictSize: ...

pytests/stubs/enums.pyi

@@ -1,5 +1,18 @@
 from typing import Any, Final, final
 
+__all__ = [
+    "ComplexEnum",
+    "MixedComplexEnum",
+    "SimpleEnum",
+    "SimpleEnumWithoutDerive",
+    "SimpleTupleEnum",
+    "TupleEnum",
+    "do_complex_stuff",
+    "do_mixed_complex_stuff",
+    "do_simple_stuff",
+    "do_tuple_stuff",
+]
+
 class ComplexEnum:
     @final
     class EmptyStruct(ComplexEnum):

pytests/stubs/misc.pyi

@@ -1,5 +1,14 @@
 from typing import Any
 
+__all__ = [
+    "accepts_bool",
+    "detach_during_finalization",
+    "get_item_and_run_callback",
+    "get_type_fully_qualified_name",
+    "hammer_attaching_in_thread",
+    "issue_219",
+]
+
 def accepts_bool(val: bool) -> bool: ...
 def detach_during_finalization() -> Any: ...
 def get_item_and_run_callback(dict: dict, callback: Any) -> None: ...

pytests/stubs/objstore.pyi

@@ -1,5 +1,7 @@
 from typing import Any, final
 
+__all__ = ["ObjStore"]
+
 @final
 class ObjStore:
     def __new__(cls, /) -> ObjStore: ...

pytests/stubs/othermod.pyi

@@ -1,5 +1,7 @@
 from typing import Final, final
 
+__all__ = ["USIZE_MAX", "USIZE_MIN", "ModClass", "double"]
+
 USIZE_MAX: Final[int]
 USIZE_MIN: Final[int]
 

pytests/stubs/path.pyi

@@ -1,5 +1,7 @@
 from os import PathLike
 from pathlib import Path
 
+__all__ = ["make_path", "take_pathbuf"]
+
 def make_path() -> Path: ...
 def take_pathbuf(path: str | PathLike[str]) -> Path: ...

pytests/stubs/pyclasses.pyi

@@ -1,6 +1,20 @@
 from _typeshed import Incomplete
 from typing import Final, final
 
+__all__ = [
+    "AssertingBaseClass",
+    "ClassWithDecorators",
+    "ClassWithDict",
+    "ClassWithoutConstructor",
+    "EmptyClass",
+    "Number",
+    "PlainObject",
+    "PyClassIter",
+    "PyClassThreadIter",
+    "SubClassWithInit",
+    "map_a_class",
+]
+
 class AssertingBaseClass:
     """
     Demonstrates a base class which can operate on the relevant subclass in its constructor.

pytests/stubs/pyfunctions.pyi

@@ -1,5 +1,18 @@
 from typing import Any
 
+__all__ = [
+    "args_kwargs",
+    "many_keyword_arguments",
+    "none",
+    "positional_only",
+    "simple",
+    "simple_args",
+    "simple_args_kwargs",
+    "simple_kwargs",
+    "with_async",
+    "with_typed_args",
+]
+
 def args_kwargs(*args, **kwargs) -> tuple[tuple, dict | None]: ...
 def many_keyword_arguments(
     *,

pytests/stubs/sequence.pyi

@@ -1,5 +1,7 @@
 from collections.abc import Sequence
 
+__all__ = ["array_to_array_i32", "vec_to_vec_i32", "vec_to_vec_pystring"]
+
 def array_to_array_i32(arr: Sequence[int]) -> list[int]: ...
 def vec_to_vec_i32(vec: Sequence[int]) -> list[int]: ...
 def vec_to_vec_pystring(vec: Sequence[str]) -> list[str]: ...