diff --git a/Cargo.lock b/Cargo.lock index 3b5cc7f4d..f1b35a78a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1832,6 +1832,7 @@ dependencies = [ "puffin-interpreter", "puffin-package", "puffin-resolver", + "puffin-workspace", "tempfile", "tokio", "tracing", @@ -1942,6 +1943,20 @@ dependencies = [ "wheel-filename", ] +[[package]] +name = "puffin-workspace" +version = "0.0.1" +dependencies = [ + "fs-err", + "pep440_rs 0.3.12", + "pep508_rs", + "puffin-package", + "pyproject-toml", + "serde", + "thiserror", + "toml_edit 0.20.2", +] + [[package]] name = "pyo3" version = "0.19.2" diff --git a/Cargo.toml b/Cargo.toml index 16e28d4de..387c0ff21 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -58,6 +58,7 @@ thiserror = { version = "1.0.49" } tokio = { version = "1.16.1", features = ["rt-multi-thread"] } tokio-util = { version = "0.7.9", features = ["compat"] } toml = { version = "0.8.2" } +toml_edit = { version = "0.20.2" } tracing = { version = "0.1.37" } tracing-subscriber = { version = "0.3.17", features = ["env-filter"] } tracing-tree = { version = "0.2.5" } diff --git a/crates/puffin-cli/Cargo.toml b/crates/puffin-cli/Cargo.toml index 163083a23..357bb812d 100644 --- a/crates/puffin-cli/Cargo.toml +++ b/crates/puffin-cli/Cargo.toml @@ -20,6 +20,7 @@ puffin-installer = { path = "../puffin-installer" } puffin-interpreter = { path = "../puffin-interpreter" } puffin-package = { path = "../puffin-package" } puffin-resolver = { path = "../puffin-resolver" } +puffin-workspace = { path = "../puffin-workspace" } anyhow = { workspace = true } bitflags = { workspace = true } diff --git a/crates/puffin-cli/src/commands/add.rs b/crates/puffin-cli/src/commands/add.rs new file mode 100644 index 000000000..cbd6ed22c --- /dev/null +++ b/crates/puffin-cli/src/commands/add.rs @@ -0,0 +1,29 @@ +use anyhow::Result; +use tracing::info; + +use crate::commands::ExitStatus; +use crate::printer::Printer; + +/// Add a dependency to the workspace. +pub(crate) fn add(name: &str, _printer: Printer) -> Result { + // Locate the workspace. + let Some(workspace_root) = puffin_workspace::find_pyproject_toml(std::env::current_dir()?) + else { + return Err(anyhow::anyhow!( + "Could not find a `pyproject.toml` file in the current directory or any of its parents" + )); + }; + + info!("Found workspace at: {}", workspace_root.display()); + + // Parse the manifest. + let mut manifest = puffin_workspace::Workspace::try_from(workspace_root.as_path())?; + + // Add the dependency. + manifest.add_dependency(name)?; + + // Write the manifest back to disk. + manifest.save(&workspace_root)?; + + Ok(ExitStatus::Success) +} diff --git a/crates/puffin-cli/src/commands/mod.rs b/crates/puffin-cli/src/commands/mod.rs index fd39aa1b3..62dfbd755 100644 --- a/crates/puffin-cli/src/commands/mod.rs +++ b/crates/puffin-cli/src/commands/mod.rs @@ -1,6 +1,7 @@ use std::process::ExitCode; use std::time::Duration; +pub(crate) use add::add; pub(crate) use clean::clean; pub(crate) use compile::compile; pub(crate) use freeze::freeze; @@ -8,6 +9,7 @@ pub(crate) use sync::{sync, SyncFlags}; pub(crate) use uninstall::uninstall; pub(crate) use venv::venv; +mod add; mod clean; mod compile; mod freeze; diff --git a/crates/puffin-cli/src/main.rs b/crates/puffin-cli/src/main.rs index f16db7870..43fa091ad 100644 --- a/crates/puffin-cli/src/main.rs +++ b/crates/puffin-cli/src/main.rs @@ -45,6 +45,8 @@ enum Commands { Uninstall(UninstallArgs), /// Create a virtual environment. Venv(VenvArgs), + /// Add a dependency to the workspace. + Add(AddArgs), } #[derive(Args)] @@ -83,6 +85,12 @@ struct VenvArgs { name: PathBuf, } +#[derive(Args)] +struct AddArgs { + /// The name of the package to add. + name: String, +} + #[tokio::main] async fn main() -> ExitCode { let cli = Cli::parse(); @@ -153,6 +161,7 @@ async fn main() -> ExitCode { .await } Commands::Venv(args) => commands::venv(&args.name, args.python.as_deref(), printer).await, + Commands::Add(args) => commands::add(&args.name, printer), }; match result { diff --git a/crates/puffin-workspace/Cargo.toml b/crates/puffin-workspace/Cargo.toml new file mode 100644 index 000000000..a2526a624 --- /dev/null +++ b/crates/puffin-workspace/Cargo.toml @@ -0,0 +1,21 @@ +[package] +name = "puffin-workspace" +version = "0.0.1" +edition = { workspace = true } +rust-version = { workspace = true } +homepage = { workspace = true } +documentation = { workspace = true } +repository = { workspace = true } +authors = { workspace = true } +license = { workspace = true } + +[dependencies] +pep440_rs = { path = "../pep440-rs" } +pep508_rs = { path = "../pep508-rs" } +puffin-package = { path = "../puffin-package" } + +fs-err = { workspace = true } +pyproject-toml = { workspace = true } +serde = { workspace = true } +thiserror = { workspace = true } +toml_edit = { workspace = true, features = ["serde"] } diff --git a/crates/puffin-workspace/src/error.rs b/crates/puffin-workspace/src/error.rs new file mode 100644 index 000000000..a51601380 --- /dev/null +++ b/crates/puffin-workspace/src/error.rs @@ -0,0 +1,17 @@ +use std::io; +use thiserror::Error; + +#[derive(Error, Debug)] +pub enum WorkspaceError { + #[error(transparent)] + IO(#[from] io::Error), + + #[error(transparent)] + InvalidToml(#[from] toml_edit::TomlError), + + #[error(transparent)] + InvalidPyproject(#[from] toml_edit::de::Error), + + #[error(transparent)] + InvalidRequirement(#[from] pep508_rs::Pep508Error), +} diff --git a/crates/puffin-workspace/src/lib.rs b/crates/puffin-workspace/src/lib.rs new file mode 100644 index 000000000..9a34c2aec --- /dev/null +++ b/crates/puffin-workspace/src/lib.rs @@ -0,0 +1,19 @@ +use std::path::{Path, PathBuf}; + +pub use error::WorkspaceError; +pub use workspace::Workspace; + +mod error; +mod toml; +mod workspace; + +/// Find the closest `pyproject.toml` file to the given path. +pub fn find_pyproject_toml(path: impl AsRef) -> Option { + for directory in path.as_ref().ancestors() { + let pyproject_toml = directory.join("pyproject.toml"); + if pyproject_toml.is_file() { + return Some(pyproject_toml); + } + } + None +} diff --git a/crates/puffin-workspace/src/toml.rs b/crates/puffin-workspace/src/toml.rs new file mode 100644 index 000000000..ae11fde8d --- /dev/null +++ b/crates/puffin-workspace/src/toml.rs @@ -0,0 +1,41 @@ +/// Reformat a TOML array to use multiline format. +pub(crate) fn format_multiline_array(dependencies: &mut toml_edit::Array) { + for item in dependencies.iter_mut() { + let decor = item.decor_mut(); + let mut prefix = String::new(); + for comment in find_comments(decor.prefix()).chain(find_comments(decor.suffix())) { + prefix.push_str("\n "); + prefix.push_str(comment); + } + prefix.push_str("\n "); + decor.set_prefix(prefix); + decor.set_suffix(""); + } + + dependencies.set_trailing(&{ + let mut comments = find_comments(Some(dependencies.trailing())).peekable(); + let mut value = String::new(); + if comments.peek().is_some() { + for comment in comments { + value.push_str("\n "); + value.push_str(comment); + } + } + value.push('\n'); + value + }); + + dependencies.set_trailing_comma(true); +} + +/// Return an iterator over the comments in a raw string. +fn find_comments(raw_string: Option<&toml_edit::RawString>) -> impl Iterator { + raw_string + .and_then(toml_edit::RawString::as_str) + .unwrap_or("") + .lines() + .filter_map(|line| { + let line = line.trim(); + line.starts_with('#').then_some(line) + }) +} diff --git a/crates/puffin-workspace/src/workspace.rs b/crates/puffin-workspace/src/workspace.rs new file mode 100644 index 000000000..e267d8369 --- /dev/null +++ b/crates/puffin-workspace/src/workspace.rs @@ -0,0 +1,137 @@ +use std::io; +use std::path::Path; +use std::str::FromStr; + +use fs_err as fs; +use pyproject_toml::{BuildSystem, Project}; +use serde::{Deserialize, Serialize}; +use toml_edit::Document; + +use pep508_rs::Requirement; +use puffin_package::package_name::PackageName; + +use crate::toml::format_multiline_array; +use crate::WorkspaceError; + +/// Unlike [`pyproject_toml::PyProjectToml`], in our case `build_system` is also optional +#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq)] +#[serde(rename_all = "kebab-case")] +struct PyProjectToml { + /// Build-related data + build_system: Option, + + /// Project metadata + project: Option, +} + +#[derive(Debug)] +pub struct Workspace { + /// The parsed `pyproject.toml`. + #[allow(unused)] + pyproject_toml: PyProjectToml, + + /// The raw document. + document: Document, +} + +impl Workspace { + /// Add a dependency to the workspace. + pub fn add_dependency(&mut self, dependency: &str) -> Result<(), WorkspaceError> { + let requirement = Requirement::from_str(dependency)?; + + let Some(project) = self + .document + .get_mut("project") + .map(|project| project.as_table_mut().unwrap()) + else { + // No `project` table. + let mut dependencies = toml_edit::Array::new(); + dependencies.push(dependency); + format_multiline_array(&mut dependencies); + + let mut project = toml_edit::Table::new(); + project.insert( + "dependencies", + toml_edit::Item::Value(toml_edit::Value::Array(dependencies)), + ); + + self.document + .insert("project", toml_edit::Item::Table(project)); + + return Ok(()); + }; + + let Some(dependencies) = project + .get_mut("dependencies") + .map(|dependencies| dependencies.as_array_mut().unwrap()) + else { + // No `dependencies` array. + let mut dependencies = toml_edit::Array::new(); + dependencies.push(dependency); + format_multiline_array(&mut dependencies); + + project.insert( + "dependencies", + toml_edit::Item::Value(toml_edit::Value::Array(dependencies)), + ); + return Ok(()); + }; + + // TODO(charlie): Awkward `drop` pattern required to work around destructors, apparently. + let mut iter = dependencies.iter(); + let index = iter.position(|item| { + let Some(item) = item.as_str() else { + return false; + }; + + let Ok(existing) = Requirement::from_str(item) else { + return false; + }; + + PackageName::normalize(&requirement.name) == PackageName::normalize(existing.name) + }); + drop(iter); + + if let Some(index) = index { + dependencies.replace(index, dependency); + } else { + dependencies.push(dependency); + } + + format_multiline_array(dependencies); + + Ok(()) + } + + /// Save the workspace to disk. + pub fn save(&self, path: impl AsRef) -> Result<(), WorkspaceError> { + let file = fs::File::create(path.as_ref())?; + self.write(file) + } + + /// Write the workspace to a writer. + fn write(&self, mut writer: impl io::Write) -> Result<(), WorkspaceError> { + writer.write_all(self.document.to_string().as_bytes())?; + Ok(()) + } +} + +impl TryFrom<&Path> for Workspace { + type Error = WorkspaceError; + + fn try_from(path: &Path) -> Result { + // Read the `pyproject.toml` from disk. + let contents = fs::read_to_string(path)?; + + // Parse the `pyproject.toml` file. + let pyproject_toml = toml_edit::de::from_str::(&contents)?; + + // Parse the raw document. + let document = contents.parse::()?; + + Ok(Self { + pyproject_toml, + document, + }) + } +}