Add a puffin add command (#117)

This needs far better error handling and user-facing feedback, but it
does the basic operation (and includes discovery of the `pyproject.toml`
file, etc.).

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"

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" }

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 }

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<ExitStatus> {
+    // 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)
+}

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;

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 {

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"] }

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),
+}

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<Path>) -> Option<PathBuf> {
+    for directory in path.as_ref().ancestors() {
+        let pyproject_toml = directory.join("pyproject.toml");
+        if pyproject_toml.is_file() {
+            return Some(pyproject_toml);
+        }
+    }
+    None
+}

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<Item = &str> {
+    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)
+        })
+}

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<BuildSystem>,
+
+    /// Project metadata
+    project: Option<Project>,
+}
+
+#[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<Path>) -> 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<Self, Self::Error> {
+        // 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::<PyProjectToml>(&contents)?;
+
+        // Parse the raw document.
+        let document = contents.parse::<Document>()?;
+
+        Ok(Self {
+            pyproject_toml,
+            document,
+        })
+    }
+}