Stub: allow custom imports

These imports are written in a small DSL like:
```rust
#[pymodule(stubs = {
    from datetime import datetime as dt, time;
    from uuid import UUID;
})]
```
Then parsed, sent as an AST inside the introspection data (following the same AST format as the type hints) and serialized by the introspection crate that merges these imports with the auto generated ones

The `#[pymodule]` parameter is named `stub` because we might include some other features in the future like protocols

Author: Tpt

newsfragments/5877.added.md

@@ -0,0 +1 @@
+Introspection: allow to set custom stub imports in `#[pymodule]` macro

pyo3-introspection/src/introspection.rs

@@ -1,6 +1,6 @@
 use crate::model::{
-    Argument, Arguments, Attribute, Class, Constant, Expr, Function, Module, Operator,
-    VariableLengthArgument,
+    Argument, Arguments, Attribute, Class, Constant, Expr, Function, ImportAlias, Module, Operator,
+    Statement, VariableLengthArgument,
 };
 use anyhow::{anyhow, bail, ensure, Context, Result};
 use goblin::elf::section_header::SHN_XINDEX;
@@ -52,6 +52,7 @@ fn parse_chunks(chunks: &[Chunk], main_module_name: &str) -> Result<Module> {
             members,
             doc,
             incomplete,
+            stubs,
         } = chunk
         {
             if name == main_module_name {
@@ -65,6 +66,7 @@ fn parse_chunks(chunks: &[Chunk], main_module_name: &str) -> Result<Module> {
                     name,
                     members,
                     *incomplete,
+                    stubs,
                     doc.as_deref(),
                     &chunks_by_id,
                     &chunks_by_parent,
@@ -82,6 +84,7 @@ fn convert_module(
     name: &str,
     members: &[String],
     mut incomplete: bool,
+    stubs: &[ChunkStatement],
     docstring: Option<&str>,
     chunks_by_id: &HashMap<&str, &Chunk>,
     chunks_by_parent: &HashMap<&str, Vec<&Chunk>>,
@@ -114,6 +117,35 @@ fn convert_module(
         functions,
         attributes,
         incomplete,
+        stubs: stubs
+            .iter()
+            .map(|statement| match statement {
+                ChunkStatement::ImportFrom {
+                    module,
+                    names,
+                    level,
+                } => Statement::ImportFrom {
+                    module: module.clone(),
+                    names: names
+                        .iter()
+                        .map(|alias| ImportAlias {
+                            name: alias.name.clone(),
+                            asname: alias.asname.clone(),
+                        })
+                        .collect(),
+                    level: *level,
+                },
+                ChunkStatement::Import { names } => Statement::Import {
+                    names: names
+                        .iter()
+                        .map(|alias| ImportAlias {
+                            name: alias.name.clone(),
+                            asname: alias.asname.clone(),
+                        })
+                        .collect(),
+                },
+            })
+            .collect(),
         docstring: docstring.map(Into::into),
     })
 }
@@ -139,12 +171,14 @@ fn convert_members<'a>(
                 members,
                 incomplete,
                 doc,
+                stubs,
             } => {
                 modules.push(convert_module(
                     id,
                     name,
                     members,
                     *incomplete,
+                    stubs,
                     doc.as_deref(),
                     chunks_by_id,
                     chunks_by_parent,
@@ -672,6 +706,8 @@ enum Chunk {
         #[serde(default)]
         doc: Option<String>,
         incomplete: bool,
+        #[serde(default)]
+        stubs: Vec<ChunkStatement>,
     },
     Class {
         id: String,
@@ -739,6 +775,25 @@ struct ChunkArgument {
     annotation: Option<ChunkExpr>,
 }
 
+#[derive(Deserialize)]
+#[serde(tag = "type", rename_all = "lowercase")]
+enum ChunkStatement {
+    ImportFrom {
+        module: String,
+        names: Vec<ChunkAlias>,
+        level: usize,
+    },
+    Import {
+        names: Vec<ChunkAlias>,
+    },
+}
+
+#[derive(Deserialize)]
+struct ChunkAlias {
+    name: String,
+    asname: Option<String>,
+}
+
 #[derive(Deserialize)]
 #[serde(tag = "type", rename_all = "lowercase")]
 enum ChunkExpr {

pyo3-introspection/src/model.rs

@@ -6,6 +6,7 @@ pub struct Module {
     pub functions: Vec<Function>,
     pub attributes: Vec<Attribute>,
     pub incomplete: bool,
+    pub stubs: Vec<Statement>,
     pub docstring: Option<String>,
 }
 
@@ -74,6 +75,30 @@ pub struct VariableLengthArgument {
     pub annotation: Option<Expr>,
 }
 
+/// A python statement
+///
+/// This is the `stmt` production of the [Python `ast` module grammar](https://docs.python.org/3/library/ast.html#abstract-grammar)
+#[derive(Debug, Eq, PartialEq, Clone, Hash)]
+pub enum Statement {
+    /// `from {module} import {names}`
+    ImportFrom {
+        module: String,
+        names: Vec<ImportAlias>,
+        level: usize,
+    },
+    /// `import {names}`
+    Import { names: Vec<ImportAlias> },
+}
+
+/// A python import alias `{name} as {asname}`
+///
+/// This is the `alias` production of the [Python `ast` module grammar](https://docs.python.org/3/library/ast.html#abstract-grammar)
+#[derive(Debug, Eq, PartialEq, Clone, Hash)]
+pub struct ImportAlias {
+    pub name: String,
+    pub asname: Option<String>,
+}
+
 /// A python expression
 ///
 /// This is the `expr` production of the [Python `ast` module grammar](https://docs.python.org/3/library/ast.html#abstract-grammar)