Skip to content

Commit

Permalink
Refactor StubInfoBuilder to reduce dependence on PyProject
Browse files Browse the repository at this point in the history
To support generating stubs without a pyproject.toml file, it needs to be possible to build a StubInfo without one. The StubInfo only uses it to find the python root, so that value can be determined earlier in the process and passed along instead of the PyProject.

StubInfoBuilder is not exported, so its ::new() can be renamed without causing backwards incompatibility.
  • Loading branch information
tvanbaak committed Nov 1, 2024
1 parent 4721c02 commit c78176a
Showing 1 changed file with 10 additions and 13 deletions.
23 changes: 10 additions & 13 deletions pyo3-stub-gen/src/generate/stub_info.rs
Original file line number Diff line number Diff line change
Expand Up @@ -5,27 +5,22 @@ use std::{collections::BTreeMap, fs, io::Write, path::*};
#[derive(Debug, Clone, PartialEq)]
pub struct StubInfo {
pub modules: BTreeMap<String, Module>,
pub pyproject: PyProject,
pub python_root: PathBuf,
}

impl StubInfo {
pub fn from_pyproject_toml(path: impl AsRef<Path>) -> Result<Self> {
let pyproject = PyProject::parse_toml(path)?;
Ok(StubInfoBuilder::new(pyproject).build())
Ok(StubInfoBuilder::from_pyproject_toml(pyproject).build())
}

pub fn generate(&self) -> Result<()> {
let python_root = self
.pyproject
.python_source()
.unwrap_or(PathBuf::from(std::env::var("CARGO_MANIFEST_DIR").unwrap()));

for (name, module) in self.modules.iter() {
let path = name.replace(".", "/");
let dest = if module.submodules.is_empty() {
python_root.join(format!("{path}.pyi"))
self.python_root.join(format!("{path}.pyi"))
} else {
python_root.join(path).join("__init__.pyi")
self.python_root.join(path).join("__init__.pyi")
};

let dir = dest.parent().context("Cannot get parent directory")?;
Expand All @@ -47,15 +42,17 @@ impl StubInfo {
struct StubInfoBuilder {
modules: BTreeMap<String, Module>,
default_module_name: String,
pyproject: PyProject,
python_root: PathBuf,
}

impl StubInfoBuilder {
fn new(pyproject: PyProject) -> Self {
fn from_pyproject_toml(pyproject: PyProject) -> Self {
Self {
modules: BTreeMap::new(),
default_module_name: pyproject.module_name().to_string(),
pyproject,
python_root: pyproject
.python_source()
.unwrap_or(PathBuf::from(std::env::var("CARGO_MANIFEST_DIR").unwrap())),
}
}

Expand Down Expand Up @@ -158,7 +155,7 @@ impl StubInfoBuilder {
self.register_submodules();
StubInfo {
modules: self.modules,
pyproject: self.pyproject,
python_root: self.python_root,
}
}
}

0 comments on commit c78176a

Please sign in to comment.