Stub: allow custom imports

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)