diff --git a/Cargo.lock b/Cargo.lock index 59ce99529..886baa622 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2028,9 +2028,11 @@ dependencies = [ "puffin-package", "puffin-resolver", "puffin-workspace", + "pyproject-toml", "tempfile", "thiserror", "tokio", + "toml 0.8.2", "tracing", "tracing-subscriber", "tracing-tree", diff --git a/crates/puffin-cli/Cargo.toml b/crates/puffin-cli/Cargo.toml index c5a2e7a7d..c64deaa78 100644 --- a/crates/puffin-cli/Cargo.toml +++ b/crates/puffin-cli/Cargo.toml @@ -33,9 +33,11 @@ indicatif = { workspace = true } itertools = { workspace = true } miette = { workspace = true, features = ["fancy"] } owo-colors = { workspace = true } +pyproject-toml = { workspace = true } tempfile = { workspace = true } thiserror = { workspace = true } tokio = { workspace = true } +toml = { workspace = true } tracing = { workspace = true } tracing-subscriber = { workspace = true } tracing-tree = { workspace = true } diff --git a/crates/puffin-cli/src/requirements.rs b/crates/puffin-cli/src/requirements.rs index a7442eb09..db1f79f7f 100644 --- a/crates/puffin-cli/src/requirements.rs +++ b/crates/puffin-cli/src/requirements.rs @@ -2,7 +2,7 @@ use std::path::PathBuf; use std::str::FromStr; use anyhow::{bail, Result}; -use itertools::Either; +use fs_err as fs; use pep508_rs::Requirement; use puffin_package::requirements_txt::RequirementsTxt; @@ -12,7 +12,9 @@ pub(crate) enum RequirementsSource { /// A dependency was provided on the command line (e.g., `pip install flask`). Name(String), /// Dependencies were provided via a `requirements.txt` file (e.g., `pip install -r requirements.txt`). - Path(PathBuf), + RequirementsTxt(PathBuf), + /// Dependencies were provided via a `pyproject.toml` file (e.g., `pip-compile pyproject.toml`). + PyprojectToml(PathBuf), } impl From for RequirementsSource { @@ -23,30 +25,56 @@ impl From for RequirementsSource { impl From for RequirementsSource { fn from(path: PathBuf) -> Self { - Self::Path(path) + if path.ends_with("pyproject.toml") { + Self::PyprojectToml(path) + } else { + Self::RequirementsTxt(path) + } } } impl RequirementsSource { /// Return an iterator over the requirements in this source. pub(crate) fn requirements(&self) -> Result> { - match self { - Self::Name(name) => { - let requirement = Requirement::from_str(name)?; - Ok(Either::Left(std::iter::once(requirement))) + let iter_name = if let Self::Name(name) = self { + let requirement = Requirement::from_str(name)?; + Some(std::iter::once(requirement)) + } else { + None + }; + + let iter_requirements_txt = if let Self::RequirementsTxt(path) = self { + let requirements_txt = RequirementsTxt::parse(path, std::env::current_dir()?)?; + if !requirements_txt.constraints.is_empty() { + bail!("Constraints in requirements files are not supported"); } - Self::Path(path) => { - let requirements_txt = RequirementsTxt::parse(path, std::env::current_dir()?)?; - if !requirements_txt.constraints.is_empty() { - bail!("Constraints in requirements files are not supported"); - } - Ok(Either::Right( - requirements_txt - .requirements - .into_iter() - .map(|entry| entry.requirement), - )) - } - } + Some( + requirements_txt + .requirements + .into_iter() + .map(|entry| entry.requirement), + ) + } else { + None + }; + + let iter_pyproject_toml = if let Self::PyprojectToml(path) = self { + let pyproject_toml = + toml::from_str::(&fs::read_to_string(path)?)?; + Some( + pyproject_toml + .project + .into_iter() + .flat_map(|project| project.dependencies.into_iter().flatten()), + ) + } else { + None + }; + + Ok(iter_name + .into_iter() + .flatten() + .chain(iter_requirements_txt.into_iter().flatten()) + .chain(iter_pyproject_toml.into_iter().flatten())) } }