PyO3: Adds generate-stubs command

This builds the package in debug mode then tries to extract the generated stubs in a given target directory

Author: Tpt

src/commands/generate_stubs.rs

@@ -0,0 +1,62 @@
+use anyhow::{Context, Result, ensure};
+use fs_err as fs;
+use maturin::{BuildOptions, BuildOrchestrator, CargoOptions, OutputOptions, PythonOptions};
+use std::io;
+use std::path::PathBuf;
+use tempfile::TempDir;
+use tracing::instrument;
+
+#[instrument(skip_all)]
+pub fn generate_stubs(
+    output: PathBuf,
+    python_options: PythonOptions,
+    cargo_options: CargoOptions,
+) -> Result<()> {
+    let temporary_wheels_dir = TempDir::new()?;
+    let build_context = BuildOptions {
+        python: python_options,
+        cargo: cargo_options,
+        generate_stubs: true,
+        output: OutputOptions {
+            out: Some(temporary_wheels_dir.path().into()),
+            ..Default::default()
+        },
+        ..Default::default()
+    }
+    .into_build_context()
+    .build()?;
+
+    let orchestrator = BuildOrchestrator::new(&build_context);
+    let wheels = orchestrator.build_wheels()?;
+    let mut found_stubs = false;
+    for wheel in wheels {
+        let mut archive = zip::ZipArchive::new(fs::File::open(&wheel.0)?)
+            .with_context(|| format!("Failed to open wheel {}", wheel.0.display()))?;
+        for idx in 0..archive.len() {
+            let mut entry = archive.by_index(idx)?;
+            if entry.name().ends_with(".pyi") {
+                let output_path = output.join(entry.name());
+                if let Some(output_parent) = output_path.parent() {
+                    fs::create_dir_all(output_parent)?;
+                }
+                io::copy(&mut entry, &mut fs::File::create_new(&output_path)?).with_context(
+                    || {
+                        format!(
+                            "Failed to copy {} from {} to {}",
+                            entry.name(),
+                            wheel.0.display(),
+                            output_path.display()
+                        )
+                    },
+                )?;
+                found_stubs = true;
+            }
+        }
+    }
+    ensure!(
+        found_stubs,
+        "No auto-generated stubs found for package {}",
+        build_context.project.module_name
+    );
+    Ok(())
+}

src/commands/mod.rs

@@ -14,6 +14,7 @@ use tracing::instrument;
 
 pub(crate) mod build;
 pub(crate) mod develop;
+pub(crate) mod generate_stubs;
 pub(crate) mod pep517;
 #[cfg(feature = "upload")]
 pub(crate) mod publish;

src/main.rs

@@ -14,7 +14,7 @@ use clap::Parser;
 use maturin::GenerateJsonSchemaOptions;
 #[cfg(feature = "upload")]
 use maturin::PublishOpt;
-use maturin::{BuildOptions, DevelopOptions, TargetTriple};
+use maturin::{BuildOptions, CargoOptions, DevelopOptions, PythonOptions, TargetTriple};
 #[cfg(feature = "scaffolding")]
 use maturin::{GenerateProjectOptions, ci::GenerateCI};
 use std::env;
@@ -163,6 +163,19 @@ enum Command {
         #[arg(value_name = "FILE")]
         files: Vec<PathBuf>,
     },
+    /// Autogenerate type stubs
+    #[command(name = "generate-stubs")]
+    GenerateStub {
+        /// The directory to store the type stubs in
+        #[arg(short, long)]
+        out: PathBuf,
+        /// Python and bindings options
+        #[command(flatten)]
+        python: PythonOptions,
+        /// Cargo build options
+        #[command(flatten)]
+        cargo: CargoOptions,
+    },
     /// Backend for the PEP 517 integration. Not for human consumption
     ///
     /// The commands are meant to be called from the python PEP 517
@@ -254,6 +267,9 @@ fn run() -> Result<()> {
         Command::GenerateCI(generate_ci) => commands::generate_ci(generate_ci)?,
         #[cfg(feature = "upload")]
         Command::Upload { publish, files } => commands::upload(publish, files)?,
+        Command::GenerateStub { out, python, cargo } => {
+            commands::generate_stubs::generate_stubs(out, python, cargo)?
+        }
         #[cfg(feature = "cli-completion")]
         Command::Completions { shell } => {
             commands::completions(shell, &mut Opt::command());

tests/cmd/maturin.stdout

@@ -4,16 +4,17 @@ packages
 Usage: maturin[EXE] [OPTIONS] <COMMAND>
 
 Commands:
-  build        Build the crate into python packages
-  publish      Build and publish the crate as python packages to pypi
-  list-python  Search and list the available python installations
-  develop      Install the crate as module in the current virtualenv
-  sdist        Build only a source distribution (sdist) without compiling
-  init         Create a new cargo project in an existing directory
-  new          Create a new cargo project
-  generate-ci  Generate CI configuration
-  upload       Upload python packages to pypi
-  help         Print this message or the help of the given subcommand(s)
+  build           Build the crate into python packages
+  publish         Build and publish the crate as python packages to pypi
+  list-python     Search and list the available python installations
+  develop         Install the crate as module in the current virtualenv
+  sdist           Build only a source distribution (sdist) without compiling
+  init            Create a new cargo project in an existing directory
+  new             Create a new cargo project
+  generate-ci     Generate CI configuration
+  upload          Upload python packages to pypi
+  generate-stubs  Autogenerate type stubs
+  help            Print this message or the help of the given subcommand(s)
 
 Options:
   -v, --verbose...

tests/common/other.rs

@@ -12,8 +12,10 @@ use pretty_assertions::assert_eq;
 use std::collections::BTreeSet;
 use std::io::Read;
 use std::path::{Path, PathBuf};
+use std::process::Command;
 use tar::Archive;
 use time::PrimitiveDateTime;
+use walkdir::WalkDir;
 use zip::ZipArchive;
 
 pub fn copy_dir_recursive(src: &Path, dst: &Path) -> std::io::Result<()> {
@@ -626,3 +628,42 @@ pub fn test_build_wheels_from_sdist(package: impl AsRef<Path>, unique_name: &str
 
     Ok(())
 }
+
+pub fn generate_stubs(
+    package: impl AsRef<Path>,
+    unique_name: &str,
+    expected_files: &[&str],
+) -> Result<()> {
+    let package = package.as_ref();
+    let output_dir = tempfile::tempdir()?;
+    assert!(
+        Command::new(env!("CARGO_BIN_EXE_maturin"))
+            .arg("generate-stubs")
+            .arg("-m")
+            .arg(package.join("Cargo.toml"))
+            .arg("--target-dir")
+            .arg(format!("test-crates/targets/{unique_name}"))
+            .arg("--out")
+            .arg(output_dir.path())
+            .status()?
+            .success()
+    );
+    let found_files = WalkDir::new(output_dir.path())
+        .sort_by_file_name()
+        .into_iter()
+        .filter(|e| match e {
+            Ok(e) => e.file_type().is_file(),
+            Err(_) => true,
+        })
+        .map(|e| {
+            Ok(e?
+                .path()
+                .strip_prefix(output_dir.path())?
+                .to_str()
+                .unwrap()
+                .replace('\\', '/'))
+        })
+        .collect::<Result<Vec<_>>>()?;
+    assert_eq!(found_files, expected_files);
+    Ok(())
+}

tests/run/integration.rs

@@ -164,3 +164,15 @@ fn abi3_without_version() {
 fn abi3_python_interpreter_args() {
     handle_result(other::abi3_python_interpreter_args());
 }
+
+#[test]
+fn abi3_generate_stubs() {
+    handle_result(other::generate_stubs(
+        "test-crates/pyo3-stub-generation",
+        "integration-pyo3-stub-generation-generate-stubs",
+        &[
+            "pyo3_stub_generation/__init__.pyi",
+            "pyo3_stub_generation/submodule.pyi",
+        ],
+    ));
+}