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"
 
 
@@ -92,5 +94,6 @@ PyO3 also provides the smaller `pyo3-introspection` binary that allows to genera
   If you commit generated stubs to your repository, generate them with the oldest Python version you support (or with the matching `abi3-pyXY` feature) so that they are valid for every version.
 - 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).
-  A `#[pymodule_init]` function [declared without the module argument](module.md#procedural-initialization) is exempt: it is not handed the module, so the module is taken to be complete.
+  Such a module also gets no `__all__` declaration, because the set of names it exports at runtime is not known.
+  A `#[pymodule_init]` function [declared without the module argument](module.md#procedural-initialization) is exempt: it is not handed the module, so the module is taken to be complete and keeps its `__all__`.
   A `Python<'_>` marker on its own does not count as the module argument.

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

@@ -1766,6 +1766,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:
@@ -1801,6 +1803,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

@@ -95,26 +95,30 @@ 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 {
         let mut buffer = String::new();
         push_docstring(&mut buffer, "", docstring);
         final_elements.push(buffer);
     }
     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');
         }
     }
@@ -126,6 +130,49 @@ 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 one at runtime. Incomplete modules are skipped: we do not know all
+/// of their members, and a partial `__all__` would hide names that do exist at runtime.
+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 listing them in the order the stub declares
+    // them keeps the output stable.
+    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 {
@@ -992,21 +1039,24 @@ mod tests {
         assert_eq!(output, "dict[A, (A3.C, A3.D, B, A2, int, int2, float)]");
     }
 
-    /// A submodule is an attribute of its parent at runtime, so the parent stub must re-export it.
-    #[test]
-    fn submodules_are_re_exported_by_their_parent() {
-        let submodule = |name: &str| Module {
+    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,
-        };
+        }
+    }
+
+    /// A submodule is an attribute of its parent at runtime, so the parent stub must re-export it.
+    #[test]
+    fn submodules_are_re_exported_by_their_parent() {
         let module = Module {
             name: "foo".into(),
-            modules: vec![submodule("zulu"), submodule("alpha")],
+            modules: vec![empty_module("zulu"), empty_module("alpha")],
             classes: Vec::new(),
             functions: Vec::new(),
             attributes: Vec::new(),
@@ -1015,10 +1065,81 @@ mod tests {
         };
         assert_eq!(
             module_stubs(&module, &[]),
-            "from . import alpha as alpha, zulu as zulu\n"
+            "from . import alpha as alpha, zulu as zulu\n\n__all__ = [\"zulu\", \"alpha\"]\n"
         );
         // A module without submodules gets no such line.
-        assert_eq!(module_stubs(&submodule("alpha"), &["foo"]), "");
+        assert_eq!(module_stubs(&empty_module("alpha"), &["foo"]), "");
+    }
+
+    /// A module with one member of every kind that lands in `__all__` at runtime.
+    fn populated_module() -> 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 `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() {
+        // 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(&populated_module(), &["foo"]),
+            "from . import sub as sub\n\n__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"]), "");
+    }
+
+    /// A `#[pymodule_init]` that is handed the module can add anything to it, so the module is
+    /// tagged incomplete and the members we know about are not the whole of `__all__`. An
+    /// initialiser declared without the module argument leaves the module complete and keeps its
+    /// `__all__`, which `test_dunder_all` and the `othermod` stubs of `pytests` cover.
+    #[test]
+    fn test_dunder_all_omitted_for_incomplete_module() {
+        let module = Module {
+            incomplete: true,
+            ..populated_module()
+        };
+        let stubs = module_stubs(&module, &["foo"]);
+        assert!(
+            !stubs.contains("__all__"),
+            "an incomplete module must not declare `__all__`:\n{stubs}"
+        );
+        assert!(stubs.contains("def __getattr__(name: str) -> Incomplete: ...\n"));
     }
 
     #[test]
@@ -1269,7 +1390,7 @@ mod tests {
         };
         assert_eq!(
             module_stubs(&module, &[]),
-            "from math import inf, nan\nX = inf\ndef func(a=nan): ...\n"
+            "from math import inf, nan\n\n__all__ = [\"X\", \"func\"]\n\nX = inf\ndef func(a=nan): ...\n"
         );
     }
 
@@ -1315,7 +1436,7 @@ mod tests {
         };
         assert_eq!(
             module_stubs(&module, &[]),
-            "from builtins import int as int2\nclass int: ...\ndef make(a: int2) -> int: ...\n"
+            "from builtins import int as int2\n\n__all__ = [\"int\", \"make\"]\n\nclass int: ...\ndef make(a: int2) -> int: ...\n"
         );
     }
 }

pytests/noxfile.py

@@ -33,7 +33,10 @@ def try_install_binary(package: str, constraint: str):
     # hypothesis itself depends on PyO3 so newer Python versions may fail
     # to build
     try_install_binary("hypothesis", ">=6.156.1")
-    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/src/othermod.rs

@@ -38,7 +38,7 @@ pub mod othermod {
     pub const USIZE_MAX: usize = usize::MAX;
 
     // An initialiser without the module argument leaves the module complete for introspection:
-    // `stubs/othermod.pyi` has no `__getattr__` catch-all.
+    // `stubs/othermod.pyi` has no `__getattr__` catch-all and keeps its `__all__`.
     #[pymodule_init]
     fn init() {}
 }

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

@@ -6,6 +6,14 @@ from _typeshed import SupportsGetItem
 from collections.abc import Buffer
 from typing import SupportsIndex, 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 SupportsIndex, 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 SupportsFloat, SupportsIndex, 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 SupportsIndex, final
 
+__all__ = ["DictSize"]
+
 @final
 class DictSize:
     def __new__(cls, /, expected: SupportsIndex) -> DictSize: ...

pytests/stubs/enums.pyi

@@ -1,6 +1,19 @@
 from typing import Any, Final, SupportsFloat, SupportsIndex, final
 from typing_extensions import disjoint_base
 
+__all__ = [
+    "ComplexEnum",
+    "MixedComplexEnum",
+    "SimpleEnum",
+    "SimpleEnumWithoutDerive",
+    "SimpleTupleEnum",
+    "TupleEnum",
+    "do_complex_stuff",
+    "do_mixed_complex_stuff",
+    "do_simple_stuff",
+    "do_tuple_stuff",
+]
+
 @disjoint_base
 class ComplexEnum:
     @final