Add stub generation support for overloads.

newsfragments/5957.added.md

@@ -0,0 +1 @@
+Add `#[pyo3(overload(...))]` attribute for declaring `@typing.overload` variants in generated stubs.

pyo3-introspection/src/introspection.rs

@@ -1,6 +1,6 @@
 use crate::model::{
     Argument, Arguments, Attribute, Class, Constant, Expr, Function, Module, Operator,
-    VariableLengthArgument,
+    OverloadSignature, VariableLengthArgument,
 };
 use anyhow::{anyhow, bail, ensure, Context, Result};
 use goblin::elf::section_header::SHN_XINDEX;
@@ -177,13 +177,15 @@ fn convert_members<'a>(
                 is_async,
                 returns,
                 doc,
+                overloads,
             } => functions.push(convert_function(
                 name,
                 arguments,
                 decorators,
                 returns,
                 *is_async,
                 doc.as_deref(),
+                overloads,
                 type_hint_for_annotation_id,
             )),
             Chunk::Attribute {
@@ -268,13 +270,15 @@ fn convert_class(
     })
 }
 
+#[expect(clippy::too_many_arguments)]
 fn convert_function(
     name: &str,
     arguments: &ChunkArguments,
     decorators: &[ChunkExpr],
     returns: &Option<ChunkExpr>,
     is_async: bool,
     docstring: Option<&str>,
+    overloads: &[ChunkOverload],
     type_hint_for_annotation_id: &HashMap<String, Expr>,
 ) -> Function {
     Function {
@@ -283,36 +287,53 @@ fn convert_function(
             .iter()
             .map(|e| convert_expr(e, type_hint_for_annotation_id))
             .collect(),
-        arguments: Arguments {
-            positional_only_arguments: arguments
-                .posonlyargs
-                .iter()
-                .map(|a| convert_argument(a, type_hint_for_annotation_id))
-                .collect(),
-            arguments: arguments
-                .args
-                .iter()
-                .map(|a| convert_argument(a, type_hint_for_annotation_id))
-                .collect(),
-            vararg: arguments
-                .vararg
-                .as_ref()
-                .map(|a| convert_variable_length_argument(a, type_hint_for_annotation_id)),
-            keyword_only_arguments: arguments
-                .kwonlyargs
-                .iter()
-                .map(|e| convert_argument(e, type_hint_for_annotation_id))
-                .collect(),
-            kwarg: arguments
-                .kwarg
-                .as_ref()
-                .map(|a| convert_variable_length_argument(a, type_hint_for_annotation_id)),
-        },
+        arguments: convert_arguments(arguments, type_hint_for_annotation_id),
         returns: returns
             .as_ref()
             .map(|a| convert_expr(a, type_hint_for_annotation_id)),
         is_async,
         docstring: docstring.map(Into::into),
+        overloads: overloads
+            .iter()
+            .map(|o| OverloadSignature {
+                arguments: convert_arguments(&o.arguments, type_hint_for_annotation_id),
+                returns: o
+                    .returns
+                    .as_ref()
+                    .map(|r| convert_expr(r, type_hint_for_annotation_id)),
+            })
+            .collect(),
+    }
+}
+
+fn convert_arguments(
+    arguments: &ChunkArguments,
+    type_hint_for_annotation_id: &HashMap<String, Expr>,
+) -> Arguments {
+    Arguments {
+        positional_only_arguments: arguments
+            .posonlyargs
+            .iter()
+            .map(|a| convert_argument(a, type_hint_for_annotation_id))
+            .collect(),
+        arguments: arguments
+            .args
+            .iter()
+            .map(|a| convert_argument(a, type_hint_for_annotation_id))
+            .collect(),
+        vararg: arguments
+            .vararg
+            .as_ref()
+            .map(|a| convert_variable_length_argument(a, type_hint_for_annotation_id)),
+        keyword_only_arguments: arguments
+            .kwonlyargs
+            .iter()
+            .map(|e| convert_argument(e, type_hint_for_annotation_id))
+            .collect(),
+        kwarg: arguments
+            .kwarg
+            .as_ref()
+            .map(|a| convert_variable_length_argument(a, type_hint_for_annotation_id)),
     }
 }
 
@@ -700,6 +721,8 @@ enum Chunk {
         is_async: bool,
         #[serde(default)]
         doc: Option<String>,
+        #[serde(default)]
+        overloads: Vec<ChunkOverload>,
     },
     Attribute {
         #[serde(default)]
@@ -739,6 +762,13 @@ struct ChunkArgument {
     annotation: Option<ChunkExpr>,
 }
 
+#[derive(Deserialize)]
+struct ChunkOverload {
+    arguments: Box<ChunkArguments>,
+    #[serde(default)]
+    returns: Option<ChunkExpr>,
+}
+
 #[derive(Deserialize)]
 #[serde(tag = "type", rename_all = "lowercase")]
 enum ChunkExpr {

pyo3-introspection/src/model.rs

@@ -31,6 +31,14 @@ pub struct Function {
     pub returns: Option<Expr>,
     pub is_async: bool,
     pub docstring: Option<String>,
+    /// `@overload` variants for this function
+    pub overloads: Vec<OverloadSignature>,
+}
+
+#[derive(Debug, Eq, PartialEq, Clone, Hash)]
+pub struct OverloadSignature {
+    pub arguments: Arguments,
+    pub returns: Option<Expr>,
 }
 
 #[derive(Debug, Eq, PartialEq, Clone, Hash)]