use anyhow::Result;
use fs_err as fs;
use maturin::{BuildOptions, BuildOrchestrator, CargoOptions, OutputOptions, PythonOptions};
use std::collections::HashMap;
use std::path::{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 mut stubs = orchestrator.generate_stubs()?;
let project_layout = &build_context.project.project_layout;
let extension_name = &project_layout.extension_name;
let module_dir = output.join(project_layout.module_dir());
if project_layout.python_module.is_some() {
if stubs.len() == 1
&& let Some(stub) = stubs.remove(Path::new("__init__.pyi"))
{
write_stub(&module_dir.join(format!("{extension_name}.pyi")), &stub)?;
} else {
write_stub_dir(&module_dir.join(extension_name), &stubs)?;
}
} else {
write_stub_dir(&module_dir, &stubs)?;
}
Ok(())
}
fn write_stub(path: &Path, content: &str) -> Result<()> {
if let Some(parent) = path.parent() {
fs::create_dir_all(parent)?;
}
fs::write(path, content)?;
Ok(())
}
fn write_stub_dir(dir: &Path, stubs: &HashMap<PathBuf, String>) -> Result<()> {
if dir.exists() {
fs::remove_dir_all(dir)?;
}
fs::create_dir_all(dir)?;
for (path, content) in stubs {
write_stub(&dir.join(path), content)?;
}
Ok(())
}