use itertools::Itertools; use serde::{Deserialize, Serialize}; use std::fmt::{Display, Formatter}; use std::path::Path; use std::str::FromStr; use std::{fmt, iter, mem}; use thiserror::Error; use toml_edit::{ Array, ArrayOfTables, DocumentMut, Formatted, Item, RawString, Table, TomlError, Value, }; use uv_cache_key::CanonicalUrl; use uv_distribution_types::Index; use uv_fs::PortablePath; use uv_normalize::GroupName; use uv_pep440::{Version, VersionParseError, VersionSpecifier, VersionSpecifiers}; use uv_pep508::{ExtraName, MarkerTree, PackageName, Requirement, VersionOrUrl}; use uv_redacted::DisplaySafeUrl; use crate::pyproject::{DependencyType, Source}; /// Raw and mutable representation of a `pyproject.toml`. /// /// This is useful for operations that require editing an existing `pyproject.toml` while /// preserving comments and other structure, such as `uv add` and `uv remove`. pub struct PyProjectTomlMut { doc: DocumentMut, target: DependencyTarget, } #[derive(Error, Debug)] pub enum Error { #[error("Failed to parse `pyproject.toml`")] Parse(#[from] Box), #[error("Failed to serialize `pyproject.toml`")] Serialize(#[from] Box), #[error("Failed to deserialize `pyproject.toml`")] Deserialize(#[from] Box), #[error("Dependencies in `pyproject.toml` are malformed")] MalformedDependencies, #[error("Sources in `pyproject.toml` are malformed")] MalformedSources, #[error("Workspace in `pyproject.toml` is malformed")] MalformedWorkspace, #[error("Expected a dependency at index {0}")] MissingDependency(usize), #[error("Failed to parse `version` field of `pyproject.toml`")] VersionParse(#[from] VersionParseError), #[error("Cannot perform ambiguous update; found multiple entries for `{}`:\n{}", package_name, requirements.iter().map(|requirement| format!("- `{requirement}`")).join("\n"))] Ambiguous { package_name: PackageName, requirements: Vec, }, #[error("Unknown bound king {0}")] UnknownBoundKind(String), } /// The result of editing an array in a TOML document. #[derive(Debug, Copy, Clone, PartialEq, Eq)] pub enum ArrayEdit { /// An existing entry (at the given index) was updated. Update(usize), /// A new entry was added at the given index (typically, the end of the array). Add(usize), } #[derive(Debug, Copy, Clone, PartialEq, Eq)] enum CommentType { /// A comment that appears on its own line. OwnLine, /// A comment that appears at the end of a line. EndOfLine, } #[derive(Debug, Clone)] struct Comment { text: String, comment_type: CommentType, } impl ArrayEdit { pub fn index(&self) -> usize { match self { Self::Update(i) | Self::Add(i) => *i, } } } /// The default version specifier when adding a dependency. // While PEP 440 allows an arbitrary number of version digits, the `major` and `minor` build on // most projects sticking to two or three components and a SemVer-ish versioning system, so can // bump the major or minor version of a major.minor or major.minor.patch input version. #[derive(Clone, Copy, Debug, Default, Deserialize, PartialEq, Eq, Serialize)] #[serde(rename_all = "kebab-case")] #[cfg_attr(feature = "clap", derive(clap::ValueEnum))] #[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))] pub enum AddBoundsKind { /// Only a lower bound, e.g., `>=1.2.3`. #[default] Lower, /// Allow the same major version, similar to the semver caret, e.g., `>=1.2.3, <2.0.0`. /// /// Leading zeroes are skipped, e.g. `>=0.1.2, <0.2.0`. Major, /// Allow the same minor version, similar to the semver tilde, e.g., `>=1.2.3, <1.3.0`. /// /// Leading zeroes are skipped, e.g. `>=0.1.2, <0.1.3`. Minor, /// Pin the exact version, e.g., `==1.2.3`. /// /// This option is not recommended, as versions are already pinned in the uv lockfile. Exact, } impl Display for AddBoundsKind { fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { match self { Self::Lower => write!(f, "lower"), Self::Major => write!(f, "major"), Self::Minor => write!(f, "minor"), Self::Exact => write!(f, "exact"), } } } impl AddBoundsKind { fn specifiers(self, version: Version) -> VersionSpecifiers { // Nomenclature: "major" is the most significant component of the version, "minor" is the // second most significant component, so most versions are either major.minor.patch or // 0.major.minor. match self { AddBoundsKind::Lower => { VersionSpecifiers::from(VersionSpecifier::greater_than_equal_version(version)) } AddBoundsKind::Major => { let leading_zeroes = version .release() .iter() .take_while(|digit| **digit == 0) .count(); // Special case: The version is 0. if leading_zeroes == version.release().len() { let upper_bound = Version::new( [0, 1] .into_iter() .chain(iter::repeat_n(0, version.release().iter().skip(2).len())), ); return VersionSpecifiers::from_iter([ VersionSpecifier::greater_than_equal_version(version), VersionSpecifier::less_than_version(upper_bound), ]); } // Compute the new major version and pad it to the same length: // 1.2.3 -> 2.0.0 // 1.2 -> 2.0 // 1 -> 2 // We ignore leading zeroes, adding Semver-style semantics to 0.x versions, too: // 0.1.2 -> 0.2.0 // 0.0.1 -> 0.0.2 let major = version.release().get(leading_zeroes).copied().unwrap_or(0); // The length of the lower bound minus the leading zero and bumped component. let trailing_zeros = version.release().iter().skip(leading_zeroes + 1).len(); let upper_bound = Version::new( iter::repeat_n(0, leading_zeroes) .chain(iter::once(major + 1)) .chain(iter::repeat_n(0, trailing_zeros)), ); VersionSpecifiers::from_iter([ VersionSpecifier::greater_than_equal_version(version), VersionSpecifier::less_than_version(upper_bound), ]) } AddBoundsKind::Minor => { let leading_zeroes = version .release() .iter() .take_while(|digit| **digit == 0) .count(); // Special case: The version is 0. if leading_zeroes == version.release().len() { let upper_bound = [0, 0, 1] .into_iter() .chain(iter::repeat_n(0, version.release().iter().skip(3).len())); return VersionSpecifiers::from_iter([ VersionSpecifier::greater_than_equal_version(version), VersionSpecifier::less_than_version(Version::new(upper_bound)), ]); } // If both major and minor version are 0, the concept of bumping the minor version // instead of the major version is not useful. Instead, we bump the next // non-zero part of the version. This avoids extending the three components of 0.0.1 // to the four components of 0.0.1.1. if leading_zeroes >= 2 { let most_significant = version.release().get(leading_zeroes).copied().unwrap_or(0); // The length of the lower bound minus the leading zero and bumped component. let trailing_zeros = version.release().iter().skip(leading_zeroes + 1).len(); let upper_bound = Version::new( iter::repeat_n(0, leading_zeroes) .chain(iter::once(most_significant + 1)) .chain(iter::repeat_n(0, trailing_zeros)), ); return VersionSpecifiers::from_iter([ VersionSpecifier::greater_than_equal_version(version), VersionSpecifier::less_than_version(upper_bound), ]); } // Compute the new minor version and pad it to the same length where possible: // 1.2.3 -> 1.3.0 // 1.2 -> 1.3 // 1 -> 1.1 // We ignore leading zero, adding Semver-style semantics to 0.x versions, too: // 0.1.2 -> 0.1.3 // 0.0.1 -> 0.0.2 // If the version has only one digit, say `1`, or if there are only leading zeroes, // pad with zeroes. let major = version.release().get(leading_zeroes).copied().unwrap_or(0); let minor = version .release() .get(leading_zeroes + 1) .copied() .unwrap_or(0); let upper_bound = Version::new( iter::repeat_n(0, leading_zeroes) .chain(iter::once(major)) .chain(iter::once(minor + 1)) .chain(iter::repeat_n( 0, version.release().iter().skip(leading_zeroes + 2).len(), )), ); VersionSpecifiers::from_iter([ VersionSpecifier::greater_than_equal_version(version), VersionSpecifier::less_than_version(upper_bound), ]) } AddBoundsKind::Exact => { VersionSpecifiers::from_iter([VersionSpecifier::equals_version(version)]) } } } } /// Specifies whether dependencies are added to a script file or a `pyproject.toml` file. #[derive(Debug, Copy, Clone, PartialEq, Eq)] pub enum DependencyTarget { /// A PEP 723 script, with inline metadata. Script, /// A project with a `pyproject.toml`. PyProjectToml, } impl PyProjectTomlMut { /// Initialize a [`PyProjectTomlMut`] from a [`str`]. pub fn from_toml(raw: &str, target: DependencyTarget) -> Result { Ok(Self { doc: raw.parse().map_err(Box::new)?, target, }) } /// Adds a project to the workspace. pub fn add_workspace(&mut self, path: impl AsRef) -> Result<(), Error> { // Get or create `tool.uv.workspace.members`. let members = self .doc .entry("tool") .or_insert(implicit()) .as_table_mut() .ok_or(Error::MalformedWorkspace)? .entry("uv") .or_insert(implicit()) .as_table_mut() .ok_or(Error::MalformedWorkspace)? .entry("workspace") .or_insert(Item::Table(Table::new())) .as_table_mut() .ok_or(Error::MalformedWorkspace)? .entry("members") .or_insert(Item::Value(Value::Array(Array::new()))) .as_array_mut() .ok_or(Error::MalformedWorkspace)?; // Add the path to the workspace. members.push(PortablePath::from(path.as_ref()).to_string()); reformat_array_multiline(members); Ok(()) } /// Retrieves a mutable reference to the `project` [`Table`] of the TOML document, creating the /// table if necessary. /// /// For a script, this returns the root table. fn project(&mut self) -> Result<&mut Table, Error> { let doc = match self.target { DependencyTarget::Script => self.doc.as_table_mut(), DependencyTarget::PyProjectToml => self .doc .entry("project") .or_insert(Item::Table(Table::new())) .as_table_mut() .ok_or(Error::MalformedDependencies)?, }; Ok(doc) } /// Retrieves an optional mutable reference to the `project` [`Table`], returning `None` if it /// doesn't exist. /// /// For a script, this returns the root table. fn project_mut(&mut self) -> Result, Error> { let doc = match self.target { DependencyTarget::Script => Some(self.doc.as_table_mut()), DependencyTarget::PyProjectToml => self .doc .get_mut("project") .map(|project| project.as_table_mut().ok_or(Error::MalformedSources)) .transpose()?, }; Ok(doc) } /// Adds a dependency to `project.dependencies`. /// /// Returns `true` if the dependency was added, `false` if it was updated. pub fn add_dependency( &mut self, req: &Requirement, source: Option<&Source>, raw: bool, ) -> Result { // Get or create `project.dependencies`. let dependencies = self .project()? .entry("dependencies") .or_insert(Item::Value(Value::Array(Array::new()))) .as_array_mut() .ok_or(Error::MalformedDependencies)?; let edit = add_dependency(req, dependencies, source.is_some(), raw)?; if let Some(source) = source { self.add_source(&req.name, source)?; } Ok(edit) } /// Adds a development dependency to `tool.uv.dev-dependencies`. /// /// Returns `true` if the dependency was added, `false` if it was updated. pub fn add_dev_dependency( &mut self, req: &Requirement, source: Option<&Source>, raw: bool, ) -> Result { // Get or create `tool.uv.dev-dependencies`. let dev_dependencies = self .doc .entry("tool") .or_insert(implicit()) .as_table_mut() .ok_or(Error::MalformedSources)? .entry("uv") .or_insert(Item::Table(Table::new())) .as_table_mut() .ok_or(Error::MalformedSources)? .entry("dev-dependencies") .or_insert(Item::Value(Value::Array(Array::new()))) .as_array_mut() .ok_or(Error::MalformedDependencies)?; let edit = add_dependency(req, dev_dependencies, source.is_some(), raw)?; if let Some(source) = source { self.add_source(&req.name, source)?; } Ok(edit) } /// Add an [`Index`] to `tool.uv.index`. pub fn add_index(&mut self, index: &Index) -> Result<(), Error> { let size = self.doc.len(); let existing = self .doc .entry("tool") .or_insert(implicit()) .as_table_mut() .ok_or(Error::MalformedSources)? .entry("uv") .or_insert(implicit()) .as_table_mut() .ok_or(Error::MalformedSources)? .entry("index") .or_insert(Item::ArrayOfTables(ArrayOfTables::new())) .as_array_of_tables_mut() .ok_or(Error::MalformedSources)?; // If there's already an index with the same name or URL, update it (and move it to the top). let mut table = existing .iter() .find(|table| { // If the index has the same name, reuse it. if let Some(index) = index.name.as_deref() { if table .get("name") .and_then(|name| name.as_str()) .is_some_and(|name| name == index) { return true; } } // If the index is the default, and there's another default index, reuse it. if index.default && table .get("default") .is_some_and(|default| default.as_bool() == Some(true)) { return true; } // If there's another index with the same URL, reuse it. if table .get("url") .and_then(|item| item.as_str()) .and_then(|url| DisplaySafeUrl::parse(url).ok()) .is_some_and(|url| { CanonicalUrl::new(&url) == CanonicalUrl::new(index.url.url()) }) { return true; } false }) .cloned() .unwrap_or_default(); // If necessary, update the name. if let Some(index) = index.name.as_deref() { if table .get("name") .and_then(|name| name.as_str()) .is_none_or(|name| name != index) { let mut formatted = Formatted::new(index.to_string()); if let Some(value) = table.get("name").and_then(Item::as_value) { if let Some(prefix) = value.decor().prefix() { formatted.decor_mut().set_prefix(prefix.clone()); } if let Some(suffix) = value.decor().suffix() { formatted.decor_mut().set_suffix(suffix.clone()); } } table.insert("name", Value::String(formatted).into()); } } // If necessary, update the URL. if table .get("url") .and_then(|item| item.as_str()) .is_none_or(|url| url != index.url.without_credentials().as_str()) { let mut formatted = Formatted::new(index.url.without_credentials().to_string()); if let Some(value) = table.get("url").and_then(Item::as_value) { if let Some(prefix) = value.decor().prefix() { formatted.decor_mut().set_prefix(prefix.clone()); } if let Some(suffix) = value.decor().suffix() { formatted.decor_mut().set_suffix(suffix.clone()); } } table.insert("url", Value::String(formatted).into()); } // If necessary, update the default. if index.default { if !table .get("default") .and_then(Item::as_bool) .is_some_and(|default| default) { let mut formatted = Formatted::new(true); if let Some(value) = table.get("default").and_then(Item::as_value) { if let Some(prefix) = value.decor().prefix() { formatted.decor_mut().set_prefix(prefix.clone()); } if let Some(suffix) = value.decor().suffix() { formatted.decor_mut().set_suffix(suffix.clone()); } } table.insert("default", Value::Boolean(formatted).into()); } } // Remove any replaced tables. existing.retain(|table| { // If the index has the same name, skip it. if let Some(index) = index.name.as_deref() { if table .get("name") .and_then(|name| name.as_str()) .is_some_and(|name| name == index) { return false; } } // If there's another default index, skip it. if index.default && table .get("default") .is_some_and(|default| default.as_bool() == Some(true)) { return false; } // If there's another index with the same URL, skip it. if table .get("url") .and_then(|item| item.as_str()) .and_then(|url| DisplaySafeUrl::parse(url).ok()) .is_some_and(|url| CanonicalUrl::new(&url) == CanonicalUrl::new(index.url.url())) { return false; } true }); // Set the position to the minimum, if it's not already the first element. if let Some(min) = existing.iter().filter_map(Table::position).min() { table.set_position(min); // Increment the position of all existing elements. for table in existing.iter_mut() { if let Some(position) = table.position() { table.set_position(position + 1); } } } else { let position = isize::try_from(size).expect("TOML table size fits in `isize`"); table.set_position(position); } // Push the item to the table. existing.push(table); Ok(()) } /// Adds a dependency to `project.optional-dependencies`. /// /// Returns `true` if the dependency was added, `false` if it was updated. pub fn add_optional_dependency( &mut self, group: &ExtraName, req: &Requirement, source: Option<&Source>, raw: bool, ) -> Result { // Get or create `project.optional-dependencies`. let optional_dependencies = self .project()? .entry("optional-dependencies") .or_insert(Item::Table(Table::new())) .as_table_like_mut() .ok_or(Error::MalformedDependencies)?; // Try to find the existing group. let existing_group = optional_dependencies.iter_mut().find_map(|(key, value)| { if ExtraName::from_str(key.get()).is_ok_and(|g| g == *group) { Some(value) } else { None } }); // If the group doesn't exist, create it. let group = match existing_group { Some(value) => value, None => optional_dependencies .entry(group.as_ref()) .or_insert(Item::Value(Value::Array(Array::new()))), } .as_array_mut() .ok_or(Error::MalformedDependencies)?; let added = add_dependency(req, group, source.is_some(), raw)?; // If `project.optional-dependencies` is an inline table, reformat it. // // Reformatting can drop comments between keys, but you can't put comments // between items in an inline table anyway. if let Some(optional_dependencies) = self .project()? .get_mut("optional-dependencies") .and_then(Item::as_inline_table_mut) { optional_dependencies.fmt(); } if let Some(source) = source { self.add_source(&req.name, source)?; } Ok(added) } /// Adds a dependency to `dependency-groups`. /// /// Returns `true` if the dependency was added, `false` if it was updated. pub fn add_dependency_group_requirement( &mut self, group: &GroupName, req: &Requirement, source: Option<&Source>, raw: bool, ) -> Result { // Get or create `dependency-groups`. let dependency_groups = self .doc .entry("dependency-groups") .or_insert(Item::Table(Table::new())) .as_table_like_mut() .ok_or(Error::MalformedDependencies)?; let was_sorted = dependency_groups .get_values() .iter() .filter_map(|(dotted_ks, _)| dotted_ks.first()) .map(|k| k.get()) .is_sorted(); // Try to find the existing group. let existing_group = dependency_groups.iter_mut().find_map(|(key, value)| { if GroupName::from_str(key.get()).is_ok_and(|g| g == *group) { Some(value) } else { None } }); // If the group doesn't exist, create it. let group = match existing_group { Some(value) => value, None => dependency_groups .entry(group.as_ref()) .or_insert(Item::Value(Value::Array(Array::new()))), } .as_array_mut() .ok_or(Error::MalformedDependencies)?; let added = add_dependency(req, group, source.is_some(), raw)?; // To avoid churn in pyproject.toml, we only sort new group keys if the // existing keys were sorted. if was_sorted { dependency_groups.sort_values(); } // If `dependency-groups` is an inline table, reformat it. // // Reformatting can drop comments between keys, but you can't put comments // between items in an inline table anyway. if let Some(dependency_groups) = self .doc .get_mut("dependency-groups") .and_then(Item::as_inline_table_mut) { dependency_groups.fmt(); } if let Some(source) = source { self.add_source(&req.name, source)?; } Ok(added) } /// Set the constraint for a requirement for an existing dependency. pub fn set_dependency_bound( &mut self, dependency_type: &DependencyType, index: usize, version: Version, bound_kind: AddBoundsKind, ) -> Result<(), Error> { let group = match dependency_type { DependencyType::Production => self.dependencies_array()?, DependencyType::Dev => self.dev_dependencies_array()?, DependencyType::Optional(extra) => self.optional_dependencies_array(extra)?, DependencyType::Group(group) => self.dependency_groups_array(group)?, }; let Some(req) = group.get(index) else { return Err(Error::MissingDependency(index)); }; let mut req = req .as_str() .and_then(try_parse_requirement) .ok_or(Error::MalformedDependencies)?; req.version_or_url = Some(VersionOrUrl::VersionSpecifier( bound_kind.specifiers(version), )); group.replace(index, req.to_string()); Ok(()) } /// Get the TOML array for `project.dependencies`. fn dependencies_array(&mut self) -> Result<&mut Array, Error> { // Get or create `project.dependencies`. let dependencies = self .project()? .entry("dependencies") .or_insert(Item::Value(Value::Array(Array::new()))) .as_array_mut() .ok_or(Error::MalformedDependencies)?; Ok(dependencies) } /// Get the TOML array for `tool.uv.dev-dependencies`. fn dev_dependencies_array(&mut self) -> Result<&mut Array, Error> { // Get or create `tool.uv.dev-dependencies`. let dev_dependencies = self .doc .entry("tool") .or_insert(implicit()) .as_table_mut() .ok_or(Error::MalformedSources)? .entry("uv") .or_insert(Item::Table(Table::new())) .as_table_mut() .ok_or(Error::MalformedSources)? .entry("dev-dependencies") .or_insert(Item::Value(Value::Array(Array::new()))) .as_array_mut() .ok_or(Error::MalformedDependencies)?; Ok(dev_dependencies) } /// Get the TOML array for a `project.optional-dependencies` entry. fn optional_dependencies_array(&mut self, group: &ExtraName) -> Result<&mut Array, Error> { // Get or create `project.optional-dependencies`. let optional_dependencies = self .project()? .entry("optional-dependencies") .or_insert(Item::Table(Table::new())) .as_table_like_mut() .ok_or(Error::MalformedDependencies)?; // Try to find the existing extra. let existing_key = optional_dependencies.iter().find_map(|(key, _value)| { if ExtraName::from_str(key).is_ok_and(|g| g == *group) { Some(key.to_string()) } else { None } }); // If the group doesn't exist, create it. let group = optional_dependencies .entry(existing_key.as_deref().unwrap_or(group.as_ref())) .or_insert(Item::Value(Value::Array(Array::new()))) .as_array_mut() .ok_or(Error::MalformedDependencies)?; Ok(group) } /// Get the TOML array for a `dependency-groups` entry. fn dependency_groups_array(&mut self, group: &GroupName) -> Result<&mut Array, Error> { // Get or create `dependency-groups`. let dependency_groups = self .doc .entry("dependency-groups") .or_insert(Item::Table(Table::new())) .as_table_like_mut() .ok_or(Error::MalformedDependencies)?; // Try to find the existing group. let existing_key = dependency_groups.iter().find_map(|(key, _value)| { if GroupName::from_str(key).is_ok_and(|g| g == *group) { Some(key.to_string()) } else { None } }); // If the group doesn't exist, create it. let group = dependency_groups .entry(existing_key.as_deref().unwrap_or(group.as_ref())) .or_insert(Item::Value(Value::Array(Array::new()))) .as_array_mut() .ok_or(Error::MalformedDependencies)?; Ok(group) } /// Adds a source to `tool.uv.sources`. fn add_source(&mut self, name: &PackageName, source: &Source) -> Result<(), Error> { // Get or create `tool.uv.sources`. let sources = self .doc .entry("tool") .or_insert(implicit()) .as_table_mut() .ok_or(Error::MalformedSources)? .entry("uv") .or_insert(implicit()) .as_table_mut() .ok_or(Error::MalformedSources)? .entry("sources") .or_insert(Item::Table(Table::new())) .as_table_mut() .ok_or(Error::MalformedSources)?; if let Some(key) = find_source(name, sources) { sources.remove(&key); } add_source(name, source, sources)?; Ok(()) } /// Removes all occurrences of dependencies with the given name. pub fn remove_dependency(&mut self, name: &PackageName) -> Result, Error> { // Try to get `project.dependencies`. let Some(dependencies) = self .project_mut()? .and_then(|project| project.get_mut("dependencies")) .map(|dependencies| { dependencies .as_array_mut() .ok_or(Error::MalformedDependencies) }) .transpose()? else { return Ok(Vec::new()); }; let requirements = remove_dependency(name, dependencies); self.remove_source(name)?; Ok(requirements) } /// Removes all occurrences of development dependencies with the given name. pub fn remove_dev_dependency(&mut self, name: &PackageName) -> Result, Error> { // Try to get `tool.uv.dev-dependencies`. let Some(dev_dependencies) = self .doc .get_mut("tool") .map(|tool| tool.as_table_mut().ok_or(Error::MalformedDependencies)) .transpose()? .and_then(|tool| tool.get_mut("uv")) .map(|tool_uv| tool_uv.as_table_mut().ok_or(Error::MalformedDependencies)) .transpose()? .and_then(|tool_uv| tool_uv.get_mut("dev-dependencies")) .map(|dependencies| { dependencies .as_array_mut() .ok_or(Error::MalformedDependencies) }) .transpose()? else { return Ok(Vec::new()); }; let requirements = remove_dependency(name, dev_dependencies); self.remove_source(name)?; Ok(requirements) } /// Removes all occurrences of optional dependencies in the group with the given name. pub fn remove_optional_dependency( &mut self, name: &PackageName, group: &ExtraName, ) -> Result, Error> { // Try to get `project.optional-dependencies.`. let Some(optional_dependencies) = self .project_mut()? .and_then(|project| project.get_mut("optional-dependencies")) .map(|extras| { extras .as_table_like_mut() .ok_or(Error::MalformedDependencies) }) .transpose()? .and_then(|extras| { extras.iter_mut().find_map(|(key, value)| { if ExtraName::from_str(key.get()).is_ok_and(|g| g == *group) { Some(value) } else { None } }) }) .map(|dependencies| { dependencies .as_array_mut() .ok_or(Error::MalformedDependencies) }) .transpose()? else { return Ok(Vec::new()); }; let requirements = remove_dependency(name, optional_dependencies); self.remove_source(name)?; Ok(requirements) } /// Removes all occurrences of the dependency in the group with the given name. pub fn remove_dependency_group_requirement( &mut self, name: &PackageName, group: &GroupName, ) -> Result, Error> { // Try to get `project.optional-dependencies.`. let Some(group_dependencies) = self .doc .get_mut("dependency-groups") .map(|groups| { groups .as_table_like_mut() .ok_or(Error::MalformedDependencies) }) .transpose()? .and_then(|groups| { groups.iter_mut().find_map(|(key, value)| { if GroupName::from_str(key.get()).is_ok_and(|g| g == *group) { Some(value) } else { None } }) }) .map(|dependencies| { dependencies .as_array_mut() .ok_or(Error::MalformedDependencies) }) .transpose()? else { return Ok(Vec::new()); }; let requirements = remove_dependency(name, group_dependencies); self.remove_source(name)?; Ok(requirements) } /// Remove a matching source from `tool.uv.sources`, if it exists. fn remove_source(&mut self, name: &PackageName) -> Result<(), Error> { // If the dependency is still in use, don't remove the source. if !self.find_dependency(name, None).is_empty() { return Ok(()); } if let Some(sources) = self .doc .get_mut("tool") .map(|tool| tool.as_table_mut().ok_or(Error::MalformedSources)) .transpose()? .and_then(|tool| tool.get_mut("uv")) .map(|tool_uv| tool_uv.as_table_mut().ok_or(Error::MalformedSources)) .transpose()? .and_then(|tool_uv| tool_uv.get_mut("sources")) .map(|sources| sources.as_table_mut().ok_or(Error::MalformedSources)) .transpose()? { if let Some(key) = find_source(name, sources) { sources.remove(&key); // Remove the `tool.uv.sources` table if it is empty. if sources.is_empty() { self.doc .entry("tool") .or_insert(implicit()) .as_table_mut() .ok_or(Error::MalformedSources)? .entry("uv") .or_insert(implicit()) .as_table_mut() .ok_or(Error::MalformedSources)? .remove("sources"); } } } Ok(()) } /// Returns `true` if the `tool.uv.dev-dependencies` table is present. pub fn has_dev_dependencies(&self) -> bool { self.doc .get("tool") .and_then(Item::as_table) .and_then(|tool| tool.get("uv")) .and_then(Item::as_table) .and_then(|uv| uv.get("dev-dependencies")) .is_some() } /// Returns `true` if the `dependency-groups` table is present and contains the given group. pub fn has_dependency_group(&self, group: &GroupName) -> bool { self.doc .get("dependency-groups") .and_then(Item::as_table) .and_then(|groups| groups.get(group.as_ref())) .is_some() } /// Returns all the places in this `pyproject.toml` that contain a dependency with the given /// name. /// /// This method searches `project.dependencies`, `tool.uv.dev-dependencies`, and /// `tool.uv.optional-dependencies`. pub fn find_dependency( &self, name: &PackageName, marker: Option<&MarkerTree>, ) -> Vec { let mut types = Vec::new(); if let Some(project) = self.doc.get("project").and_then(Item::as_table) { // Check `project.dependencies`. if let Some(dependencies) = project.get("dependencies").and_then(Item::as_array) { if !find_dependencies(name, marker, dependencies).is_empty() { types.push(DependencyType::Production); } } // Check `project.optional-dependencies`. if let Some(extras) = project .get("optional-dependencies") .and_then(Item::as_table) { for (extra, dependencies) in extras { let Some(dependencies) = dependencies.as_array() else { continue; }; let Ok(extra) = ExtraName::from_str(extra) else { continue; }; if !find_dependencies(name, marker, dependencies).is_empty() { types.push(DependencyType::Optional(extra)); } } } } // Check `dependency-groups`. if let Some(groups) = self.doc.get("dependency-groups").and_then(Item::as_table) { for (group, dependencies) in groups { let Some(dependencies) = dependencies.as_array() else { continue; }; let Ok(group) = GroupName::from_str(group) else { continue; }; if !find_dependencies(name, marker, dependencies).is_empty() { types.push(DependencyType::Group(group)); } } } // Check `tool.uv.dev-dependencies`. if let Some(dev_dependencies) = self .doc .get("tool") .and_then(Item::as_table) .and_then(|tool| tool.get("uv")) .and_then(Item::as_table) .and_then(|uv| uv.get("dev-dependencies")) .and_then(Item::as_array) { if !find_dependencies(name, marker, dev_dependencies).is_empty() { types.push(DependencyType::Dev); } } types } pub fn version(&mut self) -> Result { let version = self .doc .get("project") .and_then(Item::as_table) .and_then(|project| project.get("version")) .and_then(Item::as_str) .ok_or(Error::MalformedWorkspace)?; Ok(Version::from_str(version)?) } pub fn has_dynamic_version(&mut self) -> bool { let Some(dynamic) = self .doc .get("project") .and_then(Item::as_table) .and_then(|project| project.get("dynamic")) .and_then(Item::as_array) else { return false; }; dynamic.iter().any(|val| val.as_str() == Some("version")) } pub fn set_version(&mut self, version: &Version) -> Result<(), Error> { let project = self .doc .get_mut("project") .and_then(Item::as_table_mut) .ok_or(Error::MalformedWorkspace)?; project.insert( "version", Item::Value(Value::String(Formatted::new(version.to_string()))), ); Ok(()) } } /// Returns an implicit table. fn implicit() -> Item { let mut table = Table::new(); table.set_implicit(true); Item::Table(table) } /// Adds a dependency to the given `deps` array. /// /// Returns `true` if the dependency was added, `false` if it was updated. pub fn add_dependency( req: &Requirement, deps: &mut Array, has_source: bool, raw: bool, ) -> Result { let mut to_replace = find_dependencies(&req.name, Some(&req.marker), deps); match to_replace.as_slice() { [] => { #[derive(Debug, Copy, Clone)] enum Sort { /// The list is sorted in a case-insensitive manner. CaseInsensitive, /// The list is sorted naively in a case-insensitive manner. CaseInsensitiveNaive, /// The list is sorted in a case-sensitive manner. CaseSensitive, /// The list is sorted naively in a case-sensitive manner. CaseSensitiveNaive, /// The list is unsorted. Unsorted, } fn is_sorted(items: I) -> bool where I: IntoIterator, T: PartialOrd + Copy, { items.into_iter().tuple_windows().all(|(a, b)| a <= b) } // `deps` are either requirements (strings) or include groups (inline tables). // Here we pull out just the requirements for determining the sort. let reqs: Vec<_> = deps.iter().filter_map(Value::as_str).collect(); let reqs_lowercase: Vec<_> = reqs.iter().copied().map(str::to_lowercase).collect(); // Determine if the dependency list is sorted prior to // adding the new dependency; the new dependency list // will be sorted only when the original list is sorted // so that user's custom dependency ordering is preserved. // // Any items which aren't strings are ignored, e.g. // `{ include-group = "..." }` in dependency-groups. // // We account for both case-sensitive and case-insensitive sorting. let sort = if is_sorted( reqs_lowercase .iter() .map(String::as_str) .map(split_specifiers), ) { Sort::CaseInsensitive } else if is_sorted(reqs.iter().copied().map(split_specifiers)) { Sort::CaseSensitive } else if is_sorted(reqs_lowercase.iter().map(String::as_str)) { Sort::CaseInsensitiveNaive } else if is_sorted(reqs) { Sort::CaseSensitiveNaive } else { Sort::Unsorted }; let req_string = if raw { req.displayable_with_credentials().to_string() } else { req.to_string() }; let index = match sort { Sort::CaseInsensitive => deps.iter().position(|dep| { dep.as_str().is_some_and(|dep| { split_specifiers(&dep.to_lowercase()) > split_specifiers(&req_string.to_lowercase()) }) }), Sort::CaseInsensitiveNaive => deps.iter().position(|dep| { dep.as_str() .is_some_and(|dep| dep.to_lowercase() > req_string.to_lowercase()) }), Sort::CaseSensitive => deps.iter().position(|dep| { dep.as_str() .is_some_and(|dep| split_specifiers(dep) > split_specifiers(&req_string)) }), Sort::CaseSensitiveNaive => deps .iter() .position(|dep| dep.as_str().is_some_and(|dep| *dep > *req_string)), Sort::Unsorted => None, }; let index = index.unwrap_or_else(|| { // The dependency should be added to the end, ignoring any // `include-group` items. This preserves the order for users who // keep their `include-groups` at the bottom. deps.iter() .enumerate() .filter_map(|(i, dep)| if dep.is_str() { Some(i + 1) } else { None }) .last() .unwrap_or(deps.len()) }); let mut value = Value::from(req_string.as_str()); let decor = value.decor_mut(); // Ensure comments remain on the correct line, post-insertion match index { val if val == deps.len() => { // If we're adding to the end of the list, treat trailing comments as leading comments // on the added dependency. // // For example, given: // ```toml // dependencies = [ // "anyio", # trailing comment // ] // ``` // // If we add `flask` to the end, we want to retain the comment on `anyio`: // ```toml // dependencies = [ // "anyio", # trailing comment // "flask", // ] // ``` decor.set_prefix(deps.trailing().clone()); deps.set_trailing(""); } 0 => { // If the dependency is prepended to a non-empty list, do nothing } val => { // Retain position of end-of-line comments when a dependency is inserted right below it. // // For example, given: // ```toml // dependencies = [ // "anyio", # end-of-line comment // "flask", // ] // ``` // // If we add `pydantic` (between `anyio` and `flask`), we want to retain the comment on `anyio`: // ```toml // dependencies = [ // "anyio", # end-of-line comment // "pydantic", // "flask", // ] // ``` let targeted_decor = deps.get_mut(val).unwrap().decor_mut(); decor.set_prefix(targeted_decor.prefix().unwrap().clone()); targeted_decor.set_prefix(""); // Re-formatted later by `reformat_array_multiline` } } deps.insert_formatted(index, value); // `reformat_array_multiline` uses the indentation of the first dependency entry. // Therefore, we retrieve the indentation of the first dependency entry and apply it to // the new entry. Note that it is only necessary if the newly added dependency is going // to be the first in the list _and_ the dependency list was not empty prior to adding // the new dependency. if deps.len() > 1 && index == 0 { let prefix = deps .clone() .get(index + 1) .unwrap() .decor() .prefix() .unwrap() .clone(); // However, if the prefix includes a comment, we don't want to duplicate it. // Depending on the location of the comment, we either want to leave it as-is, or // attach it to the entry that's being moved to the next line. // // For example, given: // ```toml // dependencies = [ # comment // "flask", // ] // ``` // // If we add `anyio` to the beginning, we want to retain the comment on the open // bracket: // ```toml // dependencies = [ # comment // "anyio", // "flask", // ] // ``` // // However, given: // ```toml // dependencies = [ // # comment // "flask", // ] // ``` // // If we add `anyio` to the beginning, we want the comment to move down with the // existing entry: // entry: // ```toml // dependencies = [ // "anyio", // # comment // "flask", // ] if let Some(prefix) = prefix.as_str() { // Treat anything before the first own-line comment as a prefix on the new // entry; anything after the first own-line comment is a prefix on the existing // entry. // // This is equivalent to using the first and last line content as the prefix for // the new entry, and the rest as the prefix for the existing entry. if let Some((first_line, rest)) = prefix.split_once(['\r', '\n']) { // Determine the appropriate newline character. let newline = { let mut chars = prefix[first_line.len()..].chars(); match (chars.next(), chars.next()) { (Some('\r'), Some('\n')) => "\r\n", (Some('\r'), _) => "\r", (Some('\n'), _) => "\n", _ => "\n", } }; let last_line = rest.lines().last().unwrap_or_default(); let prefix = format!("{first_line}{newline}{last_line}"); deps.get_mut(index).unwrap().decor_mut().set_prefix(prefix); let prefix = format!("{newline}{rest}"); deps.get_mut(index + 1) .unwrap() .decor_mut() .set_prefix(prefix); } else { deps.get_mut(index).unwrap().decor_mut().set_prefix(prefix); } } else { deps.get_mut(index).unwrap().decor_mut().set_prefix(prefix); } } reformat_array_multiline(deps); Ok(ArrayEdit::Add(index)) } [_] => { let (i, mut old_req) = to_replace.remove(0); update_requirement(&mut old_req, req, has_source); deps.replace(i, old_req.to_string()); reformat_array_multiline(deps); Ok(ArrayEdit::Update(i)) } // Cannot perform ambiguous updates. _ => Err(Error::Ambiguous { package_name: req.name.clone(), requirements: to_replace .into_iter() .map(|(_, requirement)| requirement) .collect(), }), } } /// Update an existing requirement. fn update_requirement(old: &mut Requirement, new: &Requirement, has_source: bool) { // Add any new extras. let mut extras = old.extras.to_vec(); extras.extend(new.extras.iter().cloned()); extras.sort_unstable(); extras.dedup(); old.extras = extras.into_boxed_slice(); // Clear the requirement source if we are going to add to `tool.uv.sources`. if has_source { old.clear_url(); } // Update the source if a new one was specified. match &new.version_or_url { None => {} Some(VersionOrUrl::VersionSpecifier(specifier)) if specifier.is_empty() => {} Some(version_or_url) => old.version_or_url = Some(version_or_url.clone()), } // Update the marker expression. if new.marker.contents().is_some() { old.marker = new.marker; } } /// Removes all occurrences of dependencies with the given name from the given `deps` array. fn remove_dependency(name: &PackageName, deps: &mut Array) -> Vec { // Remove matching dependencies. let removed = find_dependencies(name, None, deps) .into_iter() .rev() // Reverse to preserve indices as we remove them. .filter_map(|(i, _)| { deps.remove(i) .as_str() .and_then(|req| Requirement::from_str(req).ok()) }) .collect::>(); if !removed.is_empty() { reformat_array_multiline(deps); } removed } /// Returns a `Vec` containing the all dependencies with the given name, along with their positions /// in the array. fn find_dependencies( name: &PackageName, marker: Option<&MarkerTree>, deps: &Array, ) -> Vec<(usize, Requirement)> { let mut to_replace = Vec::new(); for (i, dep) in deps.iter().enumerate() { if let Some(req) = dep.as_str().and_then(try_parse_requirement) { if marker.is_none_or(|m| *m == req.marker) && *name == req.name { to_replace.push((i, req)); } } } to_replace } /// Returns the key in `tool.uv.sources` that matches the given package name. fn find_source(name: &PackageName, sources: &Table) -> Option { for (key, _) in sources { if PackageName::from_str(key).is_ok_and(|ref key| key == name) { return Some(key.to_string()); } } None } // Add a source to `tool.uv.sources`. fn add_source(req: &PackageName, source: &Source, sources: &mut Table) -> Result<(), Error> { // Serialize as an inline table. let mut doc = toml::to_string(&source) .map_err(Box::new)? .parse::() .unwrap(); let table = mem::take(doc.as_table_mut()).into_inline_table(); sources.insert(req.as_ref(), Item::Value(Value::InlineTable(table))); Ok(()) } impl fmt::Display for PyProjectTomlMut { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { self.doc.fmt(f) } } fn try_parse_requirement(req: &str) -> Option { Requirement::from_str(req).ok() } /// Reformats a TOML array to multi line while trying to preserve all comments /// and move them around. This also formats the array to have a trailing comma. fn reformat_array_multiline(deps: &mut Array) { fn find_comments(s: Option<&RawString>) -> Box + '_> { let iter = s .and_then(|x| x.as_str()) .unwrap_or("") .lines() .scan( (false, false), |(prev_line_was_empty, prev_line_was_comment), line| { let trimmed_line = line.trim(); if let Some(index) = trimmed_line.find('#') { let comment_text = trimmed_line[index..].trim().to_string(); let comment_type = if (*prev_line_was_empty) || (*prev_line_was_comment) { CommentType::OwnLine } else { CommentType::EndOfLine }; *prev_line_was_empty = trimmed_line.is_empty(); *prev_line_was_comment = true; Some(Some(Comment { text: comment_text, comment_type, })) } else { *prev_line_was_empty = trimmed_line.is_empty(); *prev_line_was_comment = false; Some(None) } }, ) .flatten(); Box::new(iter) } let mut indentation_prefix = None; for item in deps.iter_mut() { let decor = item.decor_mut(); let mut prefix = String::new(); // Calculate the indentation prefix based on the indentation of the first dependency entry. if indentation_prefix.is_none() { let decor_prefix = decor .prefix() .and_then(|s| s.as_str()) .and_then(|s| s.lines().last()) .unwrap_or_default(); let decor_prefix = decor_prefix .split_once('#') .map(|(s, _)| s) .unwrap_or(decor_prefix); indentation_prefix = (!decor_prefix.is_empty()).then_some(decor_prefix.to_string()); } let indentation_prefix_str = format!("\n{}", indentation_prefix.as_deref().unwrap_or(" ")); for comment in find_comments(decor.prefix()).chain(find_comments(decor.suffix())) { match comment.comment_type { CommentType::OwnLine => { prefix.push_str(&indentation_prefix_str); } CommentType::EndOfLine => { prefix.push(' '); } } prefix.push_str(&comment.text); } prefix.push_str(&indentation_prefix_str); decor.set_prefix(prefix); decor.set_suffix(""); } deps.set_trailing(&{ let mut comments = find_comments(Some(deps.trailing())).peekable(); let mut rv = String::new(); if comments.peek().is_some() { for comment in comments { match comment.comment_type { CommentType::OwnLine => { let indentation_prefix_str = format!("\n{}", indentation_prefix.as_deref().unwrap_or(" ")); rv.push_str(&indentation_prefix_str); } CommentType::EndOfLine => { rv.push(' '); } } rv.push_str(&comment.text); } } if !rv.is_empty() || !deps.is_empty() { rv.push('\n'); } rv }); deps.set_trailing_comma(true); } /// Split a requirement into the package name and its dependency specifiers. /// /// E.g., given `flask>=1.0`, this function returns `("flask", ">=1.0")`. But given /// `Flask>=1.0`, this function returns `("Flask", ">=1.0")`. /// /// Extras are retained, such that `flask[dotenv]>=1.0` returns `("flask[dotenv]", ">=1.0")`. fn split_specifiers(req: &str) -> (&str, &str) { let (name, specifiers) = req .find(['>', '<', '=', '~', '!', '@']) .map_or((req, ""), |pos| { let (name, specifiers) = req.split_at(pos); (name, specifiers) }); (name.trim(), specifiers.trim()) } #[cfg(test)] mod test { use super::{AddBoundsKind, split_specifiers}; use std::str::FromStr; use uv_pep440::Version; #[test] fn split() { assert_eq!(split_specifiers("flask>=1.0"), ("flask", ">=1.0")); assert_eq!(split_specifiers("Flask>=1.0"), ("Flask", ">=1.0")); assert_eq!( split_specifiers("flask[dotenv]>=1.0"), ("flask[dotenv]", ">=1.0") ); assert_eq!(split_specifiers("flask[dotenv]"), ("flask[dotenv]", "")); assert_eq!( split_specifiers( "flask @ https://files.pythonhosted.org/packages/af/47/93213ee66ef8fae3b93b3e29206f6b251e65c97bd91d8e1c5596ef15af0a/flask-3.1.0-py3-none-any.whl" ), ( "flask", "@ https://files.pythonhosted.org/packages/af/47/93213ee66ef8fae3b93b3e29206f6b251e65c97bd91d8e1c5596ef15af0a/flask-3.1.0-py3-none-any.whl" ) ); } #[test] fn bound_kind_to_specifiers_exact() { let tests = [ ("0", "==0"), ("0.0", "==0.0"), ("0.0.0", "==0.0.0"), ("0.1", "==0.1"), ("0.0.1", "==0.0.1"), ("0.0.0.1", "==0.0.0.1"), ("1.0.0", "==1.0.0"), ("1.2", "==1.2"), ("1.2.3", "==1.2.3"), ("1.2.3.4", "==1.2.3.4"), ("1.2.3.4a1.post1", "==1.2.3.4a1.post1"), ]; for (version, expected) in tests { let actual = AddBoundsKind::Exact .specifiers(Version::from_str(version).unwrap()) .to_string(); assert_eq!(actual, expected, "{version}"); } } #[test] fn bound_kind_to_specifiers_lower() { let tests = [ ("0", ">=0"), ("0.0", ">=0.0"), ("0.0.0", ">=0.0.0"), ("0.1", ">=0.1"), ("0.0.1", ">=0.0.1"), ("0.0.0.1", ">=0.0.0.1"), ("1", ">=1"), ("1.0.0", ">=1.0.0"), ("1.2", ">=1.2"), ("1.2.3", ">=1.2.3"), ("1.2.3.4", ">=1.2.3.4"), ("1.2.3.4a1.post1", ">=1.2.3.4a1.post1"), ]; for (version, expected) in tests { let actual = AddBoundsKind::Lower .specifiers(Version::from_str(version).unwrap()) .to_string(); assert_eq!(actual, expected, "{version}"); } } #[test] fn bound_kind_to_specifiers_major() { let tests = [ ("0", ">=0, <0.1"), ("0.0", ">=0.0, <0.1"), ("0.0.0", ">=0.0.0, <0.1.0"), ("0.0.0.0", ">=0.0.0.0, <0.1.0.0"), ("0.1", ">=0.1, <0.2"), ("0.0.1", ">=0.0.1, <0.0.2"), ("0.0.1.1", ">=0.0.1.1, <0.0.2.0"), ("0.0.0.1", ">=0.0.0.1, <0.0.0.2"), ("1", ">=1, <2"), ("1.0.0", ">=1.0.0, <2.0.0"), ("1.2", ">=1.2, <2.0"), ("1.2.3", ">=1.2.3, <2.0.0"), ("1.2.3.4", ">=1.2.3.4, <2.0.0.0"), ("1.2.3.4a1.post1", ">=1.2.3.4a1.post1, <2.0.0.0"), ]; for (version, expected) in tests { let actual = AddBoundsKind::Major .specifiers(Version::from_str(version).unwrap()) .to_string(); assert_eq!(actual, expected, "{version}"); } } #[test] fn bound_kind_to_specifiers_minor() { let tests = [ ("0", ">=0, <0.0.1"), ("0.0", ">=0.0, <0.0.1"), ("0.0.0", ">=0.0.0, <0.0.1"), ("0.0.0.0", ">=0.0.0.0, <0.0.1.0"), ("0.1", ">=0.1, <0.1.1"), ("0.0.1", ">=0.0.1, <0.0.2"), ("0.0.1.1", ">=0.0.1.1, <0.0.2.0"), ("0.0.0.1", ">=0.0.0.1, <0.0.0.2"), ("1", ">=1, <1.1"), ("1.0.0", ">=1.0.0, <1.1.0"), ("1.2", ">=1.2, <1.3"), ("1.2.3", ">=1.2.3, <1.3.0"), ("1.2.3.4", ">=1.2.3.4, <1.3.0.0"), ("1.2.3.4a1.post1", ">=1.2.3.4a1.post1, <1.3.0.0"), ]; for (version, expected) in tests { let actual = AddBoundsKind::Minor .specifiers(Version::from_str(version).unwrap()) .to_string(); assert_eq!(actual, expected, "{version}"); } } }