Skip to main content

uv_workspace/
pyproject_mut.rs

1use std::fmt::{Display, Formatter};
2use std::path::Path;
3use std::str::FromStr;
4use std::{fmt, iter, mem};
5
6use itertools::Itertools;
7use serde::{Deserialize, Serialize};
8use thiserror::Error;
9use toml_edit::{
10    Array, ArrayOfTables, DocumentMut, Formatted, Item, RawString, Table, TomlError, Value,
11};
12
13use uv_cache_key::CanonicalUrl;
14use uv_distribution_types::{Index, IndexFormat, IndexUrl};
15use uv_fs::{PortablePath, is_same_file_allow_missing, try_relative_to_if};
16use uv_normalize::{ExtraName, GroupName, PackageName};
17use uv_pep440::{Version, VersionParseError, VersionSpecifier, VersionSpecifiers};
18use uv_pep508::{MarkerTree, Requirement, VersionOrUrl};
19
20use crate::pyproject::{DependencyType, Source};
21
22/// Raw and mutable representation of a `pyproject.toml`.
23///
24/// This is useful for operations that require editing an existing `pyproject.toml` while
25/// preserving comments and other structure, such as `uv add` and `uv remove`.
26pub struct PyProjectTomlMut {
27    doc: DocumentMut,
28    target: DependencyTarget,
29}
30
31fn index_locations_equal(existing: &str, incoming: &IndexUrl, root_dir: &Path) -> bool {
32    let Ok(existing) = IndexUrl::parse(existing, Some(root_dir)) else {
33        return false;
34    };
35
36    if let (IndexUrl::Path(existing), IndexUrl::Path(incoming)) = (&existing, incoming)
37        && let (Ok(existing), Ok(incoming)) = (existing.to_file_path(), incoming.to_file_path())
38        && let Some(equal) = is_same_file_allow_missing(&existing, &incoming)
39    {
40        return equal;
41    }
42
43    CanonicalUrl::new(existing.url().clone()) == CanonicalUrl::new(incoming.url().clone())
44}
45
46#[derive(Error, Debug)]
47pub enum Error {
48    #[error("Failed to parse `pyproject.toml`")]
49    Parse(#[from] Box<TomlError>),
50    #[error("Failed to serialize `pyproject.toml`")]
51    Serialize(#[from] Box<toml::ser::Error>),
52    #[error("Failed to deserialize `pyproject.toml`")]
53    Deserialize(#[from] Box<toml::de::Error>),
54    #[error("Dependencies in `pyproject.toml` are malformed")]
55    MalformedDependencies,
56    #[error("Sources in `pyproject.toml` are malformed")]
57    MalformedSources,
58    #[error("Workspace in `pyproject.toml` is malformed")]
59    MalformedWorkspace,
60    #[error("Expected a dependency at index {0}")]
61    MissingDependency(usize),
62    #[error("Failed to parse `version` field of `pyproject.toml`")]
63    VersionParse(#[from] VersionParseError),
64    #[error("Cannot perform ambiguous update; found multiple entries for `{}`:\n{}", package_name, requirements.iter().map(|requirement| format!("- `{requirement}`")).join("\n"))]
65    Ambiguous {
66        package_name: PackageName,
67        requirements: Vec<Requirement>,
68    },
69    #[error("Unknown bound king {0}")]
70    UnknownBoundKind(String),
71}
72
73/// The result of editing an array in a TOML document.
74#[derive(Debug, Copy, Clone, PartialEq, Eq)]
75pub enum ArrayEdit {
76    /// An existing entry (at the given index) was updated.
77    Update(usize),
78    /// A new entry was added at the given index (typically, the end of the array).
79    Add(usize),
80}
81
82#[derive(Debug, Clone, PartialEq, Eq)]
83enum CommentType {
84    /// A comment that appears on its own line.
85    OwnLine,
86    /// A comment that appears at the end of a line.
87    EndOfLine { leading_whitespace: String },
88}
89
90#[derive(Debug, Clone)]
91struct Comment {
92    text: String,
93    kind: CommentType,
94}
95
96/// The default version specifier when adding a dependency.
97// While PEP 440 allows an arbitrary number of version digits, the `major` and `minor` build on
98// most projects sticking to two or three components and a SemVer-ish versioning system, so can
99// bump the major or minor version of a major.minor or major.minor.patch input version.
100#[derive(Clone, Copy, Debug, Default, Deserialize, PartialEq, Eq, Serialize)]
101#[serde(rename_all = "kebab-case")]
102#[cfg_attr(feature = "clap", derive(clap::ValueEnum))]
103#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
104pub enum AddBoundsKind {
105    /// Only a lower bound, e.g., `>=1.2.3`.
106    #[default]
107    Lower,
108    /// Allow the same major version, similar to the semver caret, e.g., `>=1.2.3, <2.0.0`.
109    ///
110    /// Leading zeroes are skipped, e.g. `>=0.1.2, <0.2.0`.
111    Major,
112    /// Allow the same minor version, similar to the semver tilde, e.g., `>=1.2.3, <1.3.0`.
113    ///
114    /// Leading zeroes are skipped, e.g. `>=0.1.2, <0.1.3`.
115    Minor,
116    /// Pin the exact version, e.g., `==1.2.3`.
117    ///
118    /// This option is not recommended, as versions are already pinned in the uv lockfile.
119    Exact,
120}
121
122impl Display for AddBoundsKind {
123    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
124        match self {
125            Self::Lower => write!(f, "lower"),
126            Self::Major => write!(f, "major"),
127            Self::Minor => write!(f, "minor"),
128            Self::Exact => write!(f, "exact"),
129        }
130    }
131}
132
133impl AddBoundsKind {
134    fn specifiers(self, version: Version) -> VersionSpecifiers {
135        // Nomenclature: "major" is the most significant component of the version, "minor" is the
136        // second most significant component, so most versions are either major.minor.patch or
137        // 0.major.minor.
138        match self {
139            Self::Lower => {
140                VersionSpecifiers::from(VersionSpecifier::greater_than_equal_version(version))
141            }
142            Self::Major => {
143                let leading_zeroes = version
144                    .release()
145                    .iter()
146                    .take_while(|digit| **digit == 0)
147                    .count();
148
149                // Special case: The version is 0.
150                if leading_zeroes == version.release().len() {
151                    let upper_bound = Version::new(
152                        [0, 1]
153                            .into_iter()
154                            .chain(iter::repeat_n(0, version.release().iter().skip(2).len())),
155                    );
156                    return VersionSpecifiers::from_iter([
157                        VersionSpecifier::greater_than_equal_version(version),
158                        VersionSpecifier::less_than_version(upper_bound),
159                    ]);
160                }
161
162                // Compute the new major version and pad it to the same length:
163                // 1.2.3 -> 2.0.0
164                // 1.2 -> 2.0
165                // 1 -> 2
166                // We ignore leading zeroes, adding Semver-style semantics to 0.x versions, too:
167                // 0.1.2 -> 0.2.0
168                // 0.0.1 -> 0.0.2
169                let major = version.release().get(leading_zeroes).copied().unwrap_or(0);
170                // The length of the lower bound minus the leading zero and bumped component.
171                let trailing_zeros = version.release().iter().skip(leading_zeroes + 1).len();
172                let upper_bound = Version::new(
173                    iter::repeat_n(0, leading_zeroes)
174                        .chain(iter::once(major + 1))
175                        .chain(iter::repeat_n(0, trailing_zeros)),
176                );
177
178                VersionSpecifiers::from_iter([
179                    VersionSpecifier::greater_than_equal_version(version),
180                    VersionSpecifier::less_than_version(upper_bound),
181                ])
182            }
183            Self::Minor => {
184                let leading_zeroes = version
185                    .release()
186                    .iter()
187                    .take_while(|digit| **digit == 0)
188                    .count();
189
190                // Special case: The version is 0.
191                if leading_zeroes == version.release().len() {
192                    let upper_bound = [0, 0, 1]
193                        .into_iter()
194                        .chain(iter::repeat_n(0, version.release().iter().skip(3).len()));
195                    return VersionSpecifiers::from_iter([
196                        VersionSpecifier::greater_than_equal_version(version),
197                        VersionSpecifier::less_than_version(Version::new(upper_bound)),
198                    ]);
199                }
200
201                // If both major and minor version are 0, the concept of bumping the minor version
202                // instead of the major version is not useful. Instead, we bump the next
203                // non-zero part of the version. This avoids extending the three components of 0.0.1
204                // to the four components of 0.0.1.1.
205                if leading_zeroes >= 2 {
206                    let most_significant =
207                        version.release().get(leading_zeroes).copied().unwrap_or(0);
208                    // The length of the lower bound minus the leading zero and bumped component.
209                    let trailing_zeros = version.release().iter().skip(leading_zeroes + 1).len();
210                    let upper_bound = Version::new(
211                        iter::repeat_n(0, leading_zeroes)
212                            .chain(iter::once(most_significant + 1))
213                            .chain(iter::repeat_n(0, trailing_zeros)),
214                    );
215                    return VersionSpecifiers::from_iter([
216                        VersionSpecifier::greater_than_equal_version(version),
217                        VersionSpecifier::less_than_version(upper_bound),
218                    ]);
219                }
220
221                // Compute the new minor version and pad it to the same length where possible:
222                // 1.2.3 -> 1.3.0
223                // 1.2 -> 1.3
224                // 1 -> 1.1
225                // We ignore leading zero, adding Semver-style semantics to 0.x versions, too:
226                // 0.1.2 -> 0.1.3
227                // 0.0.1 -> 0.0.2
228
229                // If the version has only one digit, say `1`, or if there are only leading zeroes,
230                // pad with zeroes.
231                let major = version.release().get(leading_zeroes).copied().unwrap_or(0);
232                let minor = version
233                    .release()
234                    .get(leading_zeroes + 1)
235                    .copied()
236                    .unwrap_or(0);
237                let upper_bound = Version::new(
238                    iter::repeat_n(0, leading_zeroes)
239                        .chain(iter::once(major))
240                        .chain(iter::once(minor + 1))
241                        .chain(iter::repeat_n(
242                            0,
243                            version.release().iter().skip(leading_zeroes + 2).len(),
244                        )),
245                );
246
247                VersionSpecifiers::from_iter([
248                    VersionSpecifier::greater_than_equal_version(version),
249                    VersionSpecifier::less_than_version(upper_bound),
250                ])
251            }
252            Self::Exact => {
253                VersionSpecifiers::from_iter([VersionSpecifier::equals_version(version)])
254            }
255        }
256    }
257}
258
259/// Specifies whether dependencies are added to a script file or a `pyproject.toml` file.
260#[derive(Debug, Copy, Clone, PartialEq, Eq)]
261pub enum DependencyTarget {
262    /// A PEP 723 script, with inline metadata.
263    Script,
264    /// A project with a `pyproject.toml`.
265    PyProjectToml,
266}
267
268impl PyProjectTomlMut {
269    /// Initialize a [`PyProjectTomlMut`] from a [`str`].
270    pub fn from_toml(raw: &str, target: DependencyTarget) -> Result<Self, Error> {
271        Ok(Self {
272            doc: raw.parse().map_err(Box::new)?,
273            target,
274        })
275    }
276
277    /// Adds a project to the workspace.
278    pub fn add_workspace(&mut self, path: impl AsRef<Path>) -> Result<(), Error> {
279        // Get or create `tool.uv.workspace.members`.
280        let members = self
281            .doc
282            .entry("tool")
283            .or_insert(implicit())
284            .as_table_mut()
285            .ok_or(Error::MalformedWorkspace)?
286            .entry("uv")
287            .or_insert(implicit())
288            .as_table_mut()
289            .ok_or(Error::MalformedWorkspace)?
290            .entry("workspace")
291            .or_insert(Item::Table(Table::new()))
292            .as_table_mut()
293            .ok_or(Error::MalformedWorkspace)?
294            .entry("members")
295            .or_insert(Item::Value(Value::Array(Array::new())))
296            .as_array_mut()
297            .ok_or(Error::MalformedWorkspace)?;
298
299        // Add the path to the workspace.
300        members.push(PortablePath::from(path.as_ref()).to_string());
301
302        reformat_array_multiline(members);
303
304        Ok(())
305    }
306
307    /// Retrieves a mutable reference to the `project` [`Table`] of the TOML document, creating the
308    /// table if necessary.
309    ///
310    /// For a script, this returns the root table.
311    fn project(&mut self) -> Result<&mut Table, Error> {
312        let doc = match self.target {
313            DependencyTarget::Script => self.doc.as_table_mut(),
314            DependencyTarget::PyProjectToml => self
315                .doc
316                .entry("project")
317                .or_insert(Item::Table(Table::new()))
318                .as_table_mut()
319                .ok_or(Error::MalformedDependencies)?,
320        };
321        Ok(doc)
322    }
323
324    /// Retrieves an optional mutable reference to the `project` [`Table`], returning `None` if it
325    /// doesn't exist.
326    ///
327    /// For a script, this returns the root table.
328    fn project_mut(&mut self) -> Result<Option<&mut Table>, Error> {
329        let doc = match self.target {
330            DependencyTarget::Script => Some(self.doc.as_table_mut()),
331            DependencyTarget::PyProjectToml => self
332                .doc
333                .get_mut("project")
334                .map(|project| project.as_table_mut().ok_or(Error::MalformedSources))
335                .transpose()?,
336        };
337        Ok(doc)
338    }
339
340    /// Adds a dependency to `project.dependencies`.
341    ///
342    /// Returns `true` if the dependency was added, `false` if it was updated.
343    pub fn add_dependency(
344        &mut self,
345        req: &Requirement,
346        source: Option<&Source>,
347        raw: bool,
348    ) -> Result<ArrayEdit, Error> {
349        // Get or create `project.dependencies`.
350        let dependencies = self
351            .project()?
352            .entry("dependencies")
353            .or_insert(Item::Value(Value::Array(Array::new())))
354            .as_array_mut()
355            .ok_or(Error::MalformedDependencies)?;
356
357        let edit = add_dependency(req, dependencies, source.is_some(), raw)?;
358
359        if let Some(source) = source {
360            self.add_source(&req.name, source)?;
361        }
362
363        Ok(edit)
364    }
365
366    /// Replaces every exact match for a dependency declaration without modifying its source.
367    ///
368    /// Returns the position of every dependency that was replaced.
369    pub fn replace_dependency_declaration(
370        &mut self,
371        dependency_type: &DependencyType,
372        existing: &Requirement,
373        replacement: &Requirement,
374    ) -> Result<Vec<ArrayEdit>, Error> {
375        let Some(dependencies) = self.dependency_type_array_mut(dependency_type)? else {
376            return Ok(Vec::new());
377        };
378
379        let replacement = replacement.to_string();
380        let mut edits = Vec::new();
381        for (index, requirement) in
382            find_dependencies(&existing.name, Some(&existing.marker), dependencies)
383        {
384            if same_requirement_declaration(&requirement, existing) {
385                dependencies.replace(index, replacement.clone());
386                edits.push(ArrayEdit::Update(index));
387            }
388        }
389        Ok(edits)
390    }
391
392    /// Removes every exact string match for a dependency declaration without modifying its source.
393    ///
394    /// Returns the position of every dependency that was removed.
395    pub fn remove_dependency_declaration_text(
396        &mut self,
397        dependency_type: &DependencyType,
398        existing: &str,
399    ) -> Result<Vec<ArrayEdit>, Error> {
400        let Some(dependencies) = self.dependency_type_array_mut(dependency_type)? else {
401            return Ok(Vec::new());
402        };
403
404        let mut edits = Vec::new();
405        for index in dependencies
406            .iter()
407            .enumerate()
408            .filter_map(|(index, dependency)| {
409                (dependency.as_str() == Some(existing)).then_some(index)
410            })
411            .collect::<Vec<_>>()
412            .into_iter()
413            .rev()
414        {
415            remove_dependency_at(index, dependencies);
416            edits.push(ArrayEdit::Update(index));
417        }
418        if !edits.is_empty() {
419            reformat_array_multiline(dependencies);
420        }
421        edits.reverse();
422        Ok(edits)
423    }
424
425    /// Adds a development dependency to `tool.uv.dev-dependencies`.
426    ///
427    /// Returns `true` if the dependency was added, `false` if it was updated.
428    pub fn add_dev_dependency(
429        &mut self,
430        req: &Requirement,
431        source: Option<&Source>,
432        raw: bool,
433    ) -> Result<ArrayEdit, Error> {
434        // Get or create `tool.uv.dev-dependencies`.
435        let dev_dependencies = self
436            .doc
437            .entry("tool")
438            .or_insert(implicit())
439            .as_table_mut()
440            .ok_or(Error::MalformedSources)?
441            .entry("uv")
442            .or_insert(Item::Table(Table::new()))
443            .as_table_mut()
444            .ok_or(Error::MalformedSources)?
445            .entry("dev-dependencies")
446            .or_insert(Item::Value(Value::Array(Array::new())))
447            .as_array_mut()
448            .ok_or(Error::MalformedDependencies)?;
449
450        let edit = add_dependency(req, dev_dependencies, source.is_some(), raw)?;
451
452        if let Some(source) = source {
453            self.add_source(&req.name, source)?;
454        }
455
456        Ok(edit)
457    }
458
459    /// Add an [`Index`] to `tool.uv.index`.
460    pub fn add_index(&mut self, index: &Index, root_dir: &Path) -> Result<(), Error> {
461        let size = self.doc.len();
462        let existing = self
463            .doc
464            .entry("tool")
465            .or_insert(implicit())
466            .as_table_mut()
467            .ok_or(Error::MalformedSources)?
468            .entry("uv")
469            .or_insert(implicit())
470            .as_table_mut()
471            .ok_or(Error::MalformedSources)?
472            .entry("index")
473            .or_insert(Item::ArrayOfTables(ArrayOfTables::new()))
474            .as_array_of_tables_mut()
475            .ok_or(Error::MalformedSources)?;
476
477        // If there's already an index with the same name or URL, update it (and move it to the top).
478        let mut table = existing
479            .iter()
480            .find(|table| {
481                // If the index has the same name, reuse it.
482                if let Some(index) = index.name.as_deref()
483                    && table
484                        .get("name")
485                        .and_then(|name| name.as_str())
486                        .is_some_and(|name| name == index)
487                {
488                    return true;
489                }
490
491                // If the index is the default, and there's another default index, reuse it.
492                if index.default
493                    && table
494                        .get("default")
495                        .is_some_and(|default| default.as_bool() == Some(true))
496                {
497                    return true;
498                }
499
500                // If there's another index with the same URL, reuse it.
501                if table
502                    .get("url")
503                    .and_then(|item| item.as_str())
504                    .is_some_and(|url| index_locations_equal(url, &index.url, root_dir))
505                {
506                    return true;
507                }
508
509                false
510            })
511            .cloned()
512            .unwrap_or_default();
513
514        // If necessary, update the name.
515        if let Some(index) = index.name.as_deref()
516            && table
517                .get("name")
518                .and_then(|name| name.as_str())
519                .is_none_or(|name| name != index)
520        {
521            let mut formatted = Formatted::new(index.to_string());
522            if let Some(value) = table.get("name").and_then(Item::as_value) {
523                if let Some(prefix) = value.decor().prefix() {
524                    formatted.decor_mut().set_prefix(prefix.clone());
525                }
526                if let Some(suffix) = value.decor().suffix() {
527                    formatted.decor_mut().set_suffix(suffix.clone());
528                }
529            }
530            table.insert("name", Value::String(formatted).into());
531        }
532
533        let url = if let IndexUrl::Path(url) = &index.url
534            && let Ok(path) = url.to_file_path()
535            && let Ok(path) = try_relative_to_if(path, root_dir, !url.was_given_absolute())
536        {
537            PortablePath::from(&path).to_string()
538        } else {
539            index.url.without_credentials().to_string()
540        };
541        let existing_url = table.get("url").and_then(|item| item.as_str());
542
543        // Update the stored URL independently of whether the index location changed.
544        let url_needs_update = existing_url.is_none_or(|existing| existing != url);
545        let index_location_changed = existing_url
546            .is_none_or(|existing| !index_locations_equal(existing, &index.url, root_dir));
547
548        // If necessary, update the URL.
549        if url_needs_update {
550            let mut formatted = Formatted::new(url);
551            if let Some(value) = table.get("url").and_then(Item::as_value) {
552                if let Some(prefix) = value.decor().prefix() {
553                    formatted.decor_mut().set_prefix(prefix.clone());
554                }
555                if let Some(suffix) = value.decor().suffix() {
556                    formatted.decor_mut().set_suffix(suffix.clone());
557                }
558            }
559            table.insert("url", Value::String(formatted).into());
560        }
561
562        // If necessary, update the default.
563        if index.default {
564            if !table
565                .get("default")
566                .and_then(Item::as_bool)
567                .is_some_and(|default| default)
568            {
569                let mut formatted = Formatted::new(true);
570                if let Some(value) = table.get("default").and_then(Item::as_value) {
571                    if let Some(prefix) = value.decor().prefix() {
572                        formatted.decor_mut().set_prefix(prefix.clone());
573                    }
574                    if let Some(suffix) = value.decor().suffix() {
575                        formatted.decor_mut().set_suffix(suffix.clone());
576                    }
577                }
578                table.insert("default", Value::Boolean(formatted).into());
579            }
580        }
581
582        // If the index location changed, sync the format to match the incoming index.
583        if index_location_changed {
584            match index.format {
585                IndexFormat::Flat => {
586                    if table
587                        .get("format")
588                        .and_then(Item::as_str)
589                        .is_none_or(|format| format != "flat")
590                    {
591                        let mut formatted = Formatted::new("flat".to_string());
592                        if let Some(value) = table.get("format").and_then(Item::as_value) {
593                            if let Some(prefix) = value.decor().prefix() {
594                                formatted.decor_mut().set_prefix(prefix.clone());
595                            }
596                            if let Some(suffix) = value.decor().suffix() {
597                                formatted.decor_mut().set_suffix(suffix.clone());
598                            }
599                        }
600                        table.insert("format", Value::String(formatted).into());
601                    }
602                }
603                IndexFormat::Simple => {
604                    // Remove the format key if it exists (Simple is the default).
605                    table.remove("format");
606                }
607            }
608        }
609
610        // Remove any replaced tables.
611        existing.retain(|table| {
612            // If the index has the same name, skip it.
613            if let Some(index) = index.name.as_deref()
614                && table
615                    .get("name")
616                    .and_then(|name| name.as_str())
617                    .is_some_and(|name| name == index)
618            {
619                return false;
620            }
621
622            // If there's another default index, skip it.
623            if index.default
624                && table
625                    .get("default")
626                    .is_some_and(|default| default.as_bool() == Some(true))
627            {
628                return false;
629            }
630
631            // If there's another index with the same URL, skip it.
632            if table
633                .get("url")
634                .and_then(|item| item.as_str())
635                .is_some_and(|url| index_locations_equal(url, &index.url, root_dir))
636            {
637                return false;
638            }
639
640            true
641        });
642
643        // Set the position to the minimum, if it's not already the first element.
644        if let Some(min) = existing.iter().filter_map(Table::position).min() {
645            table.set_position(Some(min));
646
647            // Increment the position of all existing elements.
648            for table in existing.iter_mut() {
649                if let Some(position) = table.position() {
650                    table.set_position(Some(position + 1));
651                }
652            }
653        } else {
654            let position = isize::try_from(size).expect("TOML table size fits in `isize`");
655            table.set_position(Some(position));
656        }
657
658        // Push the item to the table.
659        existing.push(table);
660
661        Ok(())
662    }
663
664    /// Adds a dependency to `project.optional-dependencies`.
665    ///
666    /// Returns `true` if the dependency was added, `false` if it was updated.
667    pub fn add_optional_dependency(
668        &mut self,
669        group: &ExtraName,
670        req: &Requirement,
671        source: Option<&Source>,
672        raw: bool,
673    ) -> Result<ArrayEdit, Error> {
674        // Get or create `project.optional-dependencies`.
675        let optional_dependencies = self
676            .project()?
677            .entry("optional-dependencies")
678            .or_insert(Item::Table(Table::new()))
679            .as_table_like_mut()
680            .ok_or(Error::MalformedDependencies)?;
681
682        // Try to find the existing group.
683        let existing_group = optional_dependencies.iter_mut().find_map(|(key, value)| {
684            if ExtraName::from_str(key.get()).is_ok_and(|g| g == *group) {
685                Some(value)
686            } else {
687                None
688            }
689        });
690
691        // If the group doesn't exist, create it.
692        let group = match existing_group {
693            Some(value) => value,
694            None => optional_dependencies
695                .entry(group.as_ref())
696                .or_insert(Item::Value(Value::Array(Array::new()))),
697        }
698        .as_array_mut()
699        .ok_or(Error::MalformedDependencies)?;
700
701        let added = add_dependency(req, group, source.is_some(), raw)?;
702
703        // If `project.optional-dependencies` is an inline table, reformat it.
704        //
705        // Reformatting can drop comments between keys, but you can't put comments
706        // between items in an inline table anyway.
707        if let Some(optional_dependencies) = self
708            .project()?
709            .get_mut("optional-dependencies")
710            .and_then(Item::as_inline_table_mut)
711        {
712            optional_dependencies.fmt();
713        }
714
715        if let Some(source) = source {
716            self.add_source(&req.name, source)?;
717        }
718
719        Ok(added)
720    }
721
722    /// Ensure that an optional dependency group exists, creating an empty group if it doesn't.
723    pub fn ensure_optional_dependency(&mut self, extra: &ExtraName) -> Result<(), Error> {
724        // Get or create `project.optional-dependencies`.
725        let optional_dependencies = self
726            .project()?
727            .entry("optional-dependencies")
728            .or_insert(Item::Table(Table::new()))
729            .as_table_like_mut()
730            .ok_or(Error::MalformedDependencies)?;
731
732        // Check if the extra already exists.
733        let extra_exists = optional_dependencies
734            .iter()
735            .any(|(key, _value)| ExtraName::from_str(key).is_ok_and(|e| e == *extra));
736
737        // If the extra doesn't exist, create it.
738        if !extra_exists {
739            optional_dependencies.insert(extra.as_ref(), Item::Value(Value::Array(Array::new())));
740        }
741
742        // If `project.optional-dependencies` is an inline table, reformat it.
743        //
744        // Reformatting can drop comments between keys, but you can't put comments
745        // between items in an inline table anyway.
746        if let Some(optional_dependencies) = self
747            .project()?
748            .get_mut("optional-dependencies")
749            .and_then(Item::as_inline_table_mut)
750        {
751            optional_dependencies.fmt();
752        }
753
754        Ok(())
755    }
756
757    /// Adds a dependency to `dependency-groups`.
758    ///
759    /// Returns `true` if the dependency was added, `false` if it was updated.
760    pub fn add_dependency_group_requirement(
761        &mut self,
762        group: &GroupName,
763        req: &Requirement,
764        source: Option<&Source>,
765        raw: bool,
766    ) -> Result<ArrayEdit, Error> {
767        // Get or create `dependency-groups`.
768        let dependency_groups = self
769            .doc
770            .entry("dependency-groups")
771            .or_insert(Item::Table(Table::new()))
772            .as_table_like_mut()
773            .ok_or(Error::MalformedDependencies)?;
774
775        let was_sorted = dependency_groups
776            .get_values()
777            .iter()
778            .filter_map(|(dotted_ks, _)| dotted_ks.first())
779            .map(|k| k.get())
780            .is_sorted();
781
782        // Try to find the existing group.
783        let existing_group = dependency_groups.iter_mut().find_map(|(key, value)| {
784            if GroupName::from_str(key.get()).is_ok_and(|g| g == *group) {
785                Some(value)
786            } else {
787                None
788            }
789        });
790
791        // If the group doesn't exist, create it.
792        let group = match existing_group {
793            Some(value) => value,
794            None => dependency_groups
795                .entry(group.as_ref())
796                .or_insert(Item::Value(Value::Array(Array::new()))),
797        }
798        .as_array_mut()
799        .ok_or(Error::MalformedDependencies)?;
800
801        let added = add_dependency(req, group, source.is_some(), raw)?;
802
803        // To avoid churn in pyproject.toml, we only sort new group keys if the
804        // existing keys were sorted.
805        if was_sorted {
806            dependency_groups.sort_values();
807        }
808
809        // If `dependency-groups` is an inline table, reformat it.
810        //
811        // Reformatting can drop comments between keys, but you can't put comments
812        // between items in an inline table anyway.
813        if let Some(dependency_groups) = self
814            .doc
815            .get_mut("dependency-groups")
816            .and_then(Item::as_inline_table_mut)
817        {
818            dependency_groups.fmt();
819        }
820
821        if let Some(source) = source {
822            self.add_source(&req.name, source)?;
823        }
824
825        Ok(added)
826    }
827
828    /// Ensure that a dependency group exists, creating an empty group if it doesn't.
829    pub fn ensure_dependency_group(&mut self, group: &GroupName) -> Result<(), Error> {
830        // Get or create `dependency-groups`.
831        let dependency_groups = self
832            .doc
833            .entry("dependency-groups")
834            .or_insert(Item::Table(Table::new()))
835            .as_table_like_mut()
836            .ok_or(Error::MalformedDependencies)?;
837
838        let was_sorted = dependency_groups
839            .get_values()
840            .iter()
841            .filter_map(|(dotted_ks, _)| dotted_ks.first())
842            .map(|k| k.get())
843            .is_sorted();
844
845        // Check if the group already exists.
846        let group_exists = dependency_groups
847            .iter()
848            .any(|(key, _value)| GroupName::from_str(key).is_ok_and(|g| g == *group));
849
850        // If the group doesn't exist, create it.
851        if !group_exists {
852            dependency_groups.insert(group.as_ref(), Item::Value(Value::Array(Array::new())));
853
854            // To avoid churn in pyproject.toml, we only sort new group keys if the
855            // existing keys were sorted.
856            if was_sorted {
857                dependency_groups.sort_values();
858            }
859        }
860
861        // If `dependency-groups` is an inline table, reformat it.
862        //
863        // Reformatting can drop comments between keys, but you can't put comments
864        // between items in an inline table anyway.
865        if let Some(dependency_groups) = self
866            .doc
867            .get_mut("dependency-groups")
868            .and_then(Item::as_inline_table_mut)
869        {
870            dependency_groups.fmt();
871        }
872
873        Ok(())
874    }
875
876    /// Set the constraint for a requirement for an existing dependency.
877    pub fn set_dependency_bound(
878        &mut self,
879        dependency_type: &DependencyType,
880        index: usize,
881        version: Version,
882        bound_kind: AddBoundsKind,
883    ) -> Result<(), Error> {
884        let group = match dependency_type {
885            DependencyType::Production => self.dependencies_array()?,
886            DependencyType::Dev => self.dev_dependencies_array()?,
887            DependencyType::Optional(extra) => self.optional_dependencies_array(extra)?,
888            DependencyType::Group(group) => self.dependency_groups_array(group)?,
889        };
890
891        let Some(req) = group.get(index) else {
892            return Err(Error::MissingDependency(index));
893        };
894
895        let mut req = req
896            .as_str()
897            .and_then(try_parse_requirement)
898            .ok_or(Error::MalformedDependencies)?;
899        req.version_or_url = Some(VersionOrUrl::VersionSpecifier(
900            bound_kind.specifiers(version),
901        ));
902        group.replace(index, req.to_string());
903
904        Ok(())
905    }
906
907    /// Get the TOML array for `project.dependencies`.
908    fn dependencies_array(&mut self) -> Result<&mut Array, Error> {
909        // Get or create `project.dependencies`.
910        let dependencies = self
911            .project()?
912            .entry("dependencies")
913            .or_insert(Item::Value(Value::Array(Array::new())))
914            .as_array_mut()
915            .ok_or(Error::MalformedDependencies)?;
916
917        Ok(dependencies)
918    }
919
920    /// Get the TOML array for `tool.uv.dev-dependencies`.
921    fn dev_dependencies_array(&mut self) -> Result<&mut Array, Error> {
922        // Get or create `tool.uv.dev-dependencies`.
923        let dev_dependencies = self
924            .doc
925            .entry("tool")
926            .or_insert(implicit())
927            .as_table_mut()
928            .ok_or(Error::MalformedSources)?
929            .entry("uv")
930            .or_insert(Item::Table(Table::new()))
931            .as_table_mut()
932            .ok_or(Error::MalformedSources)?
933            .entry("dev-dependencies")
934            .or_insert(Item::Value(Value::Array(Array::new())))
935            .as_array_mut()
936            .ok_or(Error::MalformedDependencies)?;
937
938        Ok(dev_dependencies)
939    }
940
941    /// Get the TOML array for a `project.optional-dependencies` entry.
942    fn optional_dependencies_array(&mut self, group: &ExtraName) -> Result<&mut Array, Error> {
943        // Get or create `project.optional-dependencies`.
944        let optional_dependencies = self
945            .project()?
946            .entry("optional-dependencies")
947            .or_insert(Item::Table(Table::new()))
948            .as_table_like_mut()
949            .ok_or(Error::MalformedDependencies)?;
950
951        // Try to find the existing extra.
952        let existing_key = optional_dependencies.iter().find_map(|(key, _value)| {
953            if ExtraName::from_str(key).is_ok_and(|g| g == *group) {
954                Some(key.to_string())
955            } else {
956                None
957            }
958        });
959
960        // If the group doesn't exist, create it.
961        let group = optional_dependencies
962            .entry(existing_key.as_deref().unwrap_or(group.as_ref()))
963            .or_insert(Item::Value(Value::Array(Array::new())))
964            .as_array_mut()
965            .ok_or(Error::MalformedDependencies)?;
966
967        Ok(group)
968    }
969
970    /// Get the TOML array for a `dependency-groups` entry.
971    fn dependency_groups_array(&mut self, group: &GroupName) -> Result<&mut Array, Error> {
972        // Get or create `dependency-groups`.
973        let dependency_groups = self
974            .doc
975            .entry("dependency-groups")
976            .or_insert(Item::Table(Table::new()))
977            .as_table_like_mut()
978            .ok_or(Error::MalformedDependencies)?;
979
980        // Try to find the existing group.
981        let existing_key = dependency_groups.iter().find_map(|(key, _value)| {
982            if GroupName::from_str(key).is_ok_and(|g| g == *group) {
983                Some(key.to_string())
984            } else {
985                None
986            }
987        });
988
989        // If the group doesn't exist, create it.
990        let group = dependency_groups
991            .entry(existing_key.as_deref().unwrap_or(group.as_ref()))
992            .or_insert(Item::Value(Value::Array(Array::new())))
993            .as_array_mut()
994            .ok_or(Error::MalformedDependencies)?;
995
996        Ok(group)
997    }
998
999    /// Get an existing TOML array for a dependency type.
1000    fn dependency_type_array_mut(
1001        &mut self,
1002        dependency_type: &DependencyType,
1003    ) -> Result<Option<&mut Array>, Error> {
1004        let dependencies = match dependency_type {
1005            DependencyType::Production => self
1006                .project_mut()?
1007                .and_then(|project| project.get_mut("dependencies"))
1008                .map(|dependencies| {
1009                    dependencies
1010                        .as_array_mut()
1011                        .ok_or(Error::MalformedDependencies)
1012                })
1013                .transpose()?,
1014            DependencyType::Dev => self
1015                .doc
1016                .get_mut("tool")
1017                .map(|tool| tool.as_table_mut().ok_or(Error::MalformedDependencies))
1018                .transpose()?
1019                .and_then(|tool| tool.get_mut("uv"))
1020                .map(|tool_uv| tool_uv.as_table_mut().ok_or(Error::MalformedDependencies))
1021                .transpose()?
1022                .and_then(|tool_uv| tool_uv.get_mut("dev-dependencies"))
1023                .map(|dependencies| {
1024                    dependencies
1025                        .as_array_mut()
1026                        .ok_or(Error::MalformedDependencies)
1027                })
1028                .transpose()?,
1029            DependencyType::Optional(extra) => self
1030                .project_mut()?
1031                .and_then(|project| project.get_mut("optional-dependencies"))
1032                .map(|extras| {
1033                    extras
1034                        .as_table_like_mut()
1035                        .ok_or(Error::MalformedDependencies)
1036                })
1037                .transpose()?
1038                .and_then(|extras| {
1039                    extras.iter_mut().find_map(|(key, value)| {
1040                        if ExtraName::from_str(key.get()).is_ok_and(|name| name == *extra) {
1041                            Some(value)
1042                        } else {
1043                            None
1044                        }
1045                    })
1046                })
1047                .map(|dependencies| {
1048                    dependencies
1049                        .as_array_mut()
1050                        .ok_or(Error::MalformedDependencies)
1051                })
1052                .transpose()?,
1053            DependencyType::Group(group) => self
1054                .doc
1055                .get_mut("dependency-groups")
1056                .map(|groups| {
1057                    groups
1058                        .as_table_like_mut()
1059                        .ok_or(Error::MalformedDependencies)
1060                })
1061                .transpose()?
1062                .and_then(|groups| {
1063                    groups.iter_mut().find_map(|(key, value)| {
1064                        if GroupName::from_str(key.get()).is_ok_and(|name| name == *group) {
1065                            Some(value)
1066                        } else {
1067                            None
1068                        }
1069                    })
1070                })
1071                .map(|dependencies| {
1072                    dependencies
1073                        .as_array_mut()
1074                        .ok_or(Error::MalformedDependencies)
1075                })
1076                .transpose()?,
1077        };
1078
1079        Ok(dependencies)
1080    }
1081
1082    /// Adds a source to `tool.uv.sources`.
1083    fn add_source(&mut self, name: &PackageName, source: &Source) -> Result<(), Error> {
1084        // Get or create `tool.uv.sources`.
1085        let sources = self
1086            .doc
1087            .entry("tool")
1088            .or_insert(implicit())
1089            .as_table_mut()
1090            .ok_or(Error::MalformedSources)?
1091            .entry("uv")
1092            .or_insert(implicit())
1093            .as_table_mut()
1094            .ok_or(Error::MalformedSources)?
1095            .entry("sources")
1096            .or_insert(Item::Table(Table::new()))
1097            .as_table_mut()
1098            .ok_or(Error::MalformedSources)?;
1099
1100        if let Some(key) = find_source(name, sources) {
1101            sources.remove(&key);
1102        }
1103        add_source(name, source, sources)?;
1104
1105        Ok(())
1106    }
1107
1108    /// Removes all occurrences of dependencies with the given name.
1109    pub fn remove_dependency(&mut self, name: &PackageName) -> Result<Vec<Requirement>, Error> {
1110        // Try to get `project.dependencies`.
1111        let Some(dependencies) = self
1112            .project_mut()?
1113            .and_then(|project| project.get_mut("dependencies"))
1114            .map(|dependencies| {
1115                dependencies
1116                    .as_array_mut()
1117                    .ok_or(Error::MalformedDependencies)
1118            })
1119            .transpose()?
1120        else {
1121            return Ok(Vec::new());
1122        };
1123
1124        let requirements = remove_dependency(name, dependencies);
1125        self.remove_source(name)?;
1126
1127        Ok(requirements)
1128    }
1129
1130    /// Removes all occurrences of development dependencies with the given name.
1131    pub fn remove_dev_dependency(&mut self, name: &PackageName) -> Result<Vec<Requirement>, Error> {
1132        // Try to get `tool.uv.dev-dependencies`.
1133        let Some(dev_dependencies) = self
1134            .doc
1135            .get_mut("tool")
1136            .map(|tool| tool.as_table_mut().ok_or(Error::MalformedDependencies))
1137            .transpose()?
1138            .and_then(|tool| tool.get_mut("uv"))
1139            .map(|tool_uv| tool_uv.as_table_mut().ok_or(Error::MalformedDependencies))
1140            .transpose()?
1141            .and_then(|tool_uv| tool_uv.get_mut("dev-dependencies"))
1142            .map(|dependencies| {
1143                dependencies
1144                    .as_array_mut()
1145                    .ok_or(Error::MalformedDependencies)
1146            })
1147            .transpose()?
1148        else {
1149            return Ok(Vec::new());
1150        };
1151
1152        let requirements = remove_dependency(name, dev_dependencies);
1153        self.remove_source(name)?;
1154
1155        Ok(requirements)
1156    }
1157
1158    /// Removes all occurrences of optional dependencies in the group with the given name.
1159    pub fn remove_optional_dependency(
1160        &mut self,
1161        name: &PackageName,
1162        group: &ExtraName,
1163    ) -> Result<Vec<Requirement>, Error> {
1164        // Try to get `project.optional-dependencies.<group>`.
1165        let Some(optional_dependencies) = self
1166            .project_mut()?
1167            .and_then(|project| project.get_mut("optional-dependencies"))
1168            .map(|extras| {
1169                extras
1170                    .as_table_like_mut()
1171                    .ok_or(Error::MalformedDependencies)
1172            })
1173            .transpose()?
1174            .and_then(|extras| {
1175                extras.iter_mut().find_map(|(key, value)| {
1176                    if ExtraName::from_str(key.get()).is_ok_and(|g| g == *group) {
1177                        Some(value)
1178                    } else {
1179                        None
1180                    }
1181                })
1182            })
1183            .map(|dependencies| {
1184                dependencies
1185                    .as_array_mut()
1186                    .ok_or(Error::MalformedDependencies)
1187            })
1188            .transpose()?
1189        else {
1190            return Ok(Vec::new());
1191        };
1192
1193        let requirements = remove_dependency(name, optional_dependencies);
1194        self.remove_source(name)?;
1195
1196        Ok(requirements)
1197    }
1198
1199    /// Removes all occurrences of the dependency in the group with the given name.
1200    pub fn remove_dependency_group_requirement(
1201        &mut self,
1202        name: &PackageName,
1203        group: &GroupName,
1204    ) -> Result<Vec<Requirement>, Error> {
1205        // Try to get `project.optional-dependencies.<group>`.
1206        let Some(group_dependencies) = self
1207            .doc
1208            .get_mut("dependency-groups")
1209            .map(|groups| {
1210                groups
1211                    .as_table_like_mut()
1212                    .ok_or(Error::MalformedDependencies)
1213            })
1214            .transpose()?
1215            .and_then(|groups| {
1216                groups.iter_mut().find_map(|(key, value)| {
1217                    if GroupName::from_str(key.get()).is_ok_and(|g| g == *group) {
1218                        Some(value)
1219                    } else {
1220                        None
1221                    }
1222                })
1223            })
1224            .map(|dependencies| {
1225                dependencies
1226                    .as_array_mut()
1227                    .ok_or(Error::MalformedDependencies)
1228            })
1229            .transpose()?
1230        else {
1231            return Ok(Vec::new());
1232        };
1233
1234        let requirements = remove_dependency(name, group_dependencies);
1235        self.remove_source(name)?;
1236
1237        Ok(requirements)
1238    }
1239
1240    /// Remove a matching source from `tool.uv.sources`, if it exists.
1241    fn remove_source(&mut self, name: &PackageName) -> Result<(), Error> {
1242        // If the dependency is still in use, don't remove the source.
1243        if !self.find_dependency(name, None).is_empty() {
1244            return Ok(());
1245        }
1246
1247        if let Some(sources) = self
1248            .doc
1249            .get_mut("tool")
1250            .map(|tool| tool.as_table_mut().ok_or(Error::MalformedSources))
1251            .transpose()?
1252            .and_then(|tool| tool.get_mut("uv"))
1253            .map(|tool_uv| tool_uv.as_table_mut().ok_or(Error::MalformedSources))
1254            .transpose()?
1255            .and_then(|tool_uv| tool_uv.get_mut("sources"))
1256            .map(|sources| sources.as_table_mut().ok_or(Error::MalformedSources))
1257            .transpose()?
1258        {
1259            if let Some(key) = find_source(name, sources) {
1260                sources.remove(&key);
1261
1262                // Remove the `tool.uv.sources` table if it is empty.
1263                if sources.is_empty() {
1264                    self.doc
1265                        .entry("tool")
1266                        .or_insert(implicit())
1267                        .as_table_mut()
1268                        .ok_or(Error::MalformedSources)?
1269                        .entry("uv")
1270                        .or_insert(implicit())
1271                        .as_table_mut()
1272                        .ok_or(Error::MalformedSources)?
1273                        .remove("sources");
1274                }
1275            }
1276        }
1277
1278        Ok(())
1279    }
1280
1281    /// Returns `true` if the `tool.uv.dev-dependencies` table is present.
1282    pub fn has_dev_dependencies(&self) -> bool {
1283        self.doc
1284            .get("tool")
1285            .and_then(Item::as_table)
1286            .and_then(|tool| tool.get("uv"))
1287            .and_then(Item::as_table)
1288            .and_then(|uv| uv.get("dev-dependencies"))
1289            .is_some()
1290    }
1291
1292    /// Returns `true` if the `dependency-groups` table is present and contains the given group.
1293    pub fn has_dependency_group(&self, group: &GroupName) -> bool {
1294        self.doc
1295            .get("dependency-groups")
1296            .and_then(Item::as_table)
1297            .and_then(|groups| groups.get(group.as_ref()))
1298            .is_some()
1299    }
1300
1301    /// Returns all the places in this `pyproject.toml` that contain a dependency with the given
1302    /// name.
1303    ///
1304    /// This method searches `project.dependencies`, `tool.uv.dev-dependencies`, and
1305    /// `tool.uv.optional-dependencies`.
1306    pub fn find_dependency(
1307        &self,
1308        name: &PackageName,
1309        marker: Option<&MarkerTree>,
1310    ) -> Vec<DependencyType> {
1311        let mut types = Vec::new();
1312
1313        if let Some(project) = self.doc.get("project").and_then(Item::as_table) {
1314            // Check `project.dependencies`.
1315            if let Some(dependencies) = project.get("dependencies").and_then(Item::as_array)
1316                && !find_dependencies(name, marker, dependencies).is_empty()
1317            {
1318                types.push(DependencyType::Production);
1319            }
1320
1321            // Check `project.optional-dependencies`.
1322            if let Some(extras) = project
1323                .get("optional-dependencies")
1324                .and_then(Item::as_table)
1325            {
1326                for (extra, dependencies) in extras {
1327                    let Some(dependencies) = dependencies.as_array() else {
1328                        continue;
1329                    };
1330                    let Ok(extra) = ExtraName::from_str(extra) else {
1331                        continue;
1332                    };
1333
1334                    if !find_dependencies(name, marker, dependencies).is_empty() {
1335                        types.push(DependencyType::Optional(extra));
1336                    }
1337                }
1338            }
1339        }
1340
1341        // Check `dependency-groups`.
1342        if let Some(groups) = self.doc.get("dependency-groups").and_then(Item::as_table) {
1343            for (group, dependencies) in groups {
1344                let Some(dependencies) = dependencies.as_array() else {
1345                    continue;
1346                };
1347                let Ok(group) = GroupName::from_str(group) else {
1348                    continue;
1349                };
1350
1351                if !find_dependencies(name, marker, dependencies).is_empty() {
1352                    types.push(DependencyType::Group(group));
1353                }
1354            }
1355        }
1356
1357        // Check `tool.uv.dev-dependencies`.
1358        if let Some(dev_dependencies) = self
1359            .doc
1360            .get("tool")
1361            .and_then(Item::as_table)
1362            .and_then(|tool| tool.get("uv"))
1363            .and_then(Item::as_table)
1364            .and_then(|uv| uv.get("dev-dependencies"))
1365            .and_then(Item::as_array)
1366            && !find_dependencies(name, marker, dev_dependencies).is_empty()
1367        {
1368            types.push(DependencyType::Dev);
1369        }
1370
1371        types
1372    }
1373
1374    pub fn version(&mut self) -> Result<Version, Error> {
1375        let version = self
1376            .doc
1377            .get("project")
1378            .and_then(Item::as_table)
1379            .and_then(|project| project.get("version"))
1380            .and_then(Item::as_str)
1381            .ok_or(Error::MalformedWorkspace)?;
1382
1383        Ok(Version::from_str(version)?)
1384    }
1385
1386    pub fn has_dynamic_version(&mut self) -> bool {
1387        let Some(dynamic) = self
1388            .doc
1389            .get("project")
1390            .and_then(Item::as_table)
1391            .and_then(|project| project.get("dynamic"))
1392            .and_then(Item::as_array)
1393        else {
1394            return false;
1395        };
1396
1397        dynamic.iter().any(|val| val.as_str() == Some("version"))
1398    }
1399
1400    pub fn set_version(&mut self, version: &Version) -> Result<(), Error> {
1401        let project = self
1402            .doc
1403            .get_mut("project")
1404            .and_then(Item::as_table_mut)
1405            .ok_or(Error::MalformedWorkspace)?;
1406
1407        if let Some(existing) = project.get_mut("version") {
1408            if let Some(value) = existing.as_value_mut() {
1409                let mut formatted = Value::from(version.to_string());
1410                *formatted.decor_mut() = value.decor().clone();
1411                *value = formatted;
1412            } else {
1413                *existing = Item::Value(Value::from(version.to_string()));
1414            }
1415        } else {
1416            project.insert("version", Item::Value(Value::from(version.to_string())));
1417        }
1418
1419        Ok(())
1420    }
1421}
1422
1423/// Returns an implicit table.
1424fn implicit() -> Item {
1425    let mut table = Table::new();
1426    table.set_implicit(true);
1427    Item::Table(table)
1428}
1429
1430/// Adds a dependency to the given `deps` array.
1431///
1432/// Returns `true` if the dependency was added, `false` if it was updated.
1433fn add_dependency(
1434    req: &Requirement,
1435    deps: &mut Array,
1436    has_source: bool,
1437    raw: bool,
1438) -> Result<ArrayEdit, Error> {
1439    let mut to_replace = find_dependencies(&req.name, Some(&req.marker), deps);
1440
1441    match to_replace.as_slice() {
1442        [] => {
1443            #[derive(Debug, Copy, Clone)]
1444            enum Sort {
1445                /// The list is sorted in a case-insensitive manner.
1446                CaseInsensitive,
1447                /// The list is sorted naively in a case-insensitive manner.
1448                CaseInsensitiveNaive,
1449                /// The list is sorted in a case-sensitive manner.
1450                CaseSensitive,
1451                /// The list is sorted naively in a case-sensitive manner.
1452                CaseSensitiveNaive,
1453                /// The list is unsorted.
1454                Unsorted,
1455            }
1456
1457            fn is_sorted<T, I>(items: I) -> bool
1458            where
1459                I: IntoIterator<Item = T>,
1460                T: PartialOrd + Copy,
1461            {
1462                items.into_iter().tuple_windows().all(|(a, b)| a <= b)
1463            }
1464
1465            // `deps` are either requirements (strings) or include groups (inline tables).
1466            // Here we pull out just the requirements for determining the sort.
1467            let reqs: Vec<_> = deps.iter().filter_map(Value::as_str).collect();
1468            let reqs_lowercase: Vec<_> = reqs.iter().copied().map(str::to_lowercase).collect();
1469
1470            // Determine if the dependency list is sorted prior to
1471            // adding the new dependency; the new dependency list
1472            // will be sorted only when the original list is sorted
1473            // so that user's custom dependency ordering is preserved.
1474            //
1475            // Any items which aren't strings are ignored, e.g.
1476            // `{ include-group = "..." }` in dependency-groups.
1477            //
1478            // We account for both case-sensitive and case-insensitive sorting.
1479            let sort = if is_sorted(
1480                reqs_lowercase
1481                    .iter()
1482                    .map(String::as_str)
1483                    .map(split_specifiers),
1484            ) {
1485                Sort::CaseInsensitive
1486            } else if is_sorted(reqs.iter().copied().map(split_specifiers)) {
1487                Sort::CaseSensitive
1488            } else if is_sorted(reqs_lowercase.iter().map(String::as_str)) {
1489                Sort::CaseInsensitiveNaive
1490            } else if is_sorted(reqs) {
1491                Sort::CaseSensitiveNaive
1492            } else {
1493                Sort::Unsorted
1494            };
1495
1496            let req_string = if raw {
1497                req.displayable_with_credentials().to_string()
1498            } else {
1499                req.to_string()
1500            };
1501            let index = match sort {
1502                Sort::CaseInsensitive => deps.iter().position(|dep| {
1503                    dep.as_str().is_some_and(|dep| {
1504                        split_specifiers(&dep.to_lowercase())
1505                            > split_specifiers(&req_string.to_lowercase())
1506                    })
1507                }),
1508                Sort::CaseInsensitiveNaive => deps.iter().position(|dep| {
1509                    dep.as_str()
1510                        .is_some_and(|dep| dep.to_lowercase() > req_string.to_lowercase())
1511                }),
1512                Sort::CaseSensitive => deps.iter().position(|dep| {
1513                    dep.as_str()
1514                        .is_some_and(|dep| split_specifiers(dep) > split_specifiers(&req_string))
1515                }),
1516                Sort::CaseSensitiveNaive => deps
1517                    .iter()
1518                    .position(|dep| dep.as_str().is_some_and(|dep| *dep > *req_string)),
1519                Sort::Unsorted => None,
1520            };
1521            let index = index.unwrap_or_else(|| {
1522                // The dependency should be added to the end, ignoring any
1523                // `include-group` items. This preserves the order for users who
1524                // keep their `include-groups` at the bottom.
1525                deps.iter()
1526                    .enumerate()
1527                    .filter_map(|(i, dep)| if dep.is_str() { Some(i + 1) } else { None })
1528                    .last()
1529                    .unwrap_or(deps.len())
1530            });
1531
1532            let mut value = Value::from(req_string.as_str());
1533
1534            let decor = value.decor_mut();
1535
1536            // Ensure comments remain on the correct line, post-insertion
1537            match index {
1538                val if val == deps.len() => {
1539                    // If we're adding to the end of the list, treat trailing comments as leading comments
1540                    // on the added dependency.
1541                    //
1542                    // For example, given:
1543                    // ```toml
1544                    // dependencies = [
1545                    //     "anyio", # trailing comment
1546                    // ]
1547                    // ```
1548                    //
1549                    // If we add `flask` to the end, we want to retain the comment on `anyio`:
1550                    // ```toml
1551                    // dependencies = [
1552                    //     "anyio", # trailing comment
1553                    //     "flask",
1554                    // ]
1555                    // ```
1556                    decor.set_prefix(deps.trailing().clone());
1557                    deps.set_trailing("");
1558                }
1559                0 => {
1560                    // If the dependency is prepended to a non-empty list, do nothing
1561                }
1562                val => {
1563                    // Retain position of end-of-line comments when a dependency is inserted right below it.
1564                    //
1565                    // For example, given:
1566                    // ```toml
1567                    // dependencies = [
1568                    //     "anyio", # end-of-line comment
1569                    //     "flask",
1570                    // ]
1571                    // ```
1572                    //
1573                    // If we add `pydantic` (between `anyio` and `flask`), we want to retain the comment on `anyio`:
1574                    // ```toml
1575                    // dependencies = [
1576                    //     "anyio", # end-of-line comment
1577                    //     "pydantic",
1578                    //     "flask",
1579                    // ]
1580                    // ```
1581                    let targeted_decor = deps.get_mut(val).unwrap().decor_mut();
1582                    decor.set_prefix(targeted_decor.prefix().unwrap().clone());
1583                    targeted_decor.set_prefix(""); // Re-formatted later by `reformat_array_multiline`
1584                }
1585            }
1586
1587            deps.insert_formatted(index, value);
1588
1589            // `reformat_array_multiline` uses the indentation of the first dependency entry.
1590            // Therefore, we retrieve the indentation of the first dependency entry and apply it to
1591            // the new entry. Note that it is only necessary if the newly added dependency is going
1592            // to be the first in the list _and_ the dependency list was not empty prior to adding
1593            // the new dependency.
1594            if deps.len() > 1 && index == 0 {
1595                let prefix = deps
1596                    .clone()
1597                    .get(index + 1)
1598                    .unwrap()
1599                    .decor()
1600                    .prefix()
1601                    .unwrap()
1602                    .clone();
1603
1604                // However, if the prefix includes a comment, we don't want to duplicate it.
1605                // Depending on the location of the comment, we either want to leave it as-is, or
1606                // attach it to the entry that's being moved to the next line.
1607                //
1608                // For example, given:
1609                // ```toml
1610                // dependencies = [ # comment
1611                //     "flask",
1612                // ]
1613                // ```
1614                //
1615                // If we add `anyio` to the beginning, we want to retain the comment on the open
1616                // bracket:
1617                // ```toml
1618                // dependencies = [ # comment
1619                //     "anyio",
1620                //     "flask",
1621                // ]
1622                // ```
1623                //
1624                // However, given:
1625                // ```toml
1626                // dependencies = [
1627                //     # comment
1628                //     "flask",
1629                // ]
1630                // ```
1631                //
1632                // If we add `anyio` to the beginning, we want the comment to move down with the
1633                // existing entry:
1634                // entry:
1635                // ```toml
1636                // dependencies = [
1637                //     "anyio",
1638                //     # comment
1639                //     "flask",
1640                // ]
1641                if let Some(prefix) = prefix.as_str() {
1642                    // Treat anything before the first own-line comment as a prefix on the new
1643                    // entry; anything after the first own-line comment is a prefix on the existing
1644                    // entry.
1645                    //
1646                    // This is equivalent to using the first and last line content as the prefix for
1647                    // the new entry, and the rest as the prefix for the existing entry.
1648                    if let Some((first_line, rest)) = prefix.split_once(['\r', '\n']) {
1649                        // Determine the appropriate newline character.
1650                        let newline = {
1651                            let mut chars = prefix[first_line.len()..].chars();
1652                            match (chars.next(), chars.next()) {
1653                                (Some('\r'), Some('\n')) => "\r\n",
1654                                (Some('\r'), _) => "\r",
1655                                (Some('\n'), _) => "\n",
1656                                _ => "\n",
1657                            }
1658                        };
1659                        let last_line = rest.lines().last().unwrap_or_default();
1660                        let prefix = format!("{first_line}{newline}{last_line}");
1661                        deps.get_mut(index).unwrap().decor_mut().set_prefix(prefix);
1662
1663                        let prefix = format!("{newline}{rest}");
1664                        deps.get_mut(index + 1)
1665                            .unwrap()
1666                            .decor_mut()
1667                            .set_prefix(prefix);
1668                    } else {
1669                        deps.get_mut(index).unwrap().decor_mut().set_prefix(prefix);
1670                    }
1671                } else {
1672                    deps.get_mut(index).unwrap().decor_mut().set_prefix(prefix);
1673                }
1674            }
1675
1676            reformat_array_multiline(deps);
1677
1678            Ok(ArrayEdit::Add(index))
1679        }
1680        [_] => {
1681            let (i, mut old_req) = to_replace.remove(0);
1682            update_requirement(&mut old_req, req, has_source);
1683            deps.replace(i, old_req.to_string());
1684            reformat_array_multiline(deps);
1685            Ok(ArrayEdit::Update(i))
1686        }
1687        // Cannot perform ambiguous updates.
1688        _ => Err(Error::Ambiguous {
1689            package_name: req.name.clone(),
1690            requirements: to_replace
1691                .into_iter()
1692                .map(|(_, requirement)| requirement)
1693                .collect(),
1694        }),
1695    }
1696}
1697
1698/// Update an existing requirement.
1699fn update_requirement(old: &mut Requirement, new: &Requirement, has_source: bool) {
1700    // Add any new extras.
1701    let mut extras = old.extras.to_vec();
1702    extras.extend(new.extras.iter().cloned());
1703    extras.sort_unstable();
1704    extras.dedup();
1705    old.extras = extras.into_boxed_slice();
1706
1707    // Clear the requirement source if we are going to add to `tool.uv.sources`.
1708    if has_source {
1709        old.clear_url();
1710    }
1711
1712    // Update the source if a new one was specified.
1713    match &new.version_or_url {
1714        None => {}
1715        Some(VersionOrUrl::VersionSpecifier(specifier)) if specifier.is_empty() => {}
1716        Some(version_or_url) => old.version_or_url = Some(version_or_url.clone()),
1717    }
1718
1719    // Update the marker expression.
1720    if new.marker.contents().is_some() {
1721        old.marker = new.marker;
1722    }
1723}
1724
1725/// Removes all occurrences of dependencies with the given name from the given `deps` array.
1726fn remove_dependency(name: &PackageName, deps: &mut Array) -> Vec<Requirement> {
1727    // Remove in reverse to preserve indices. Before each removal, transfer the item's
1728    // prefix (which may contain end-of-line comments belonging to the previous line) to
1729    // the next item or array trailing so comments are not lost.
1730    //
1731    // For example, in:
1732    // ```toml
1733    // dependencies = [
1734    //     "numpy>=2.4.3", # essential comment
1735    //     "requests>=2.32.5",
1736    // ]
1737    // ```
1738    //
1739    // The comment `# essential comment` is stored by `toml_edit` in the prefix of
1740    // `requests`. When `requests` is removed, we transfer it so it remains on the
1741    // `numpy` line.
1742    let removed = find_dependencies(name, None, deps)
1743        .into_iter()
1744        .rev()
1745        .filter_map(|(i, _)| remove_dependency_at(i, deps))
1746        .collect::<Vec<_>>();
1747
1748    if !removed.is_empty() {
1749        reformat_array_multiline(deps);
1750    }
1751
1752    removed
1753}
1754
1755fn remove_dependency_at(index: usize, deps: &mut Array) -> Option<Requirement> {
1756    if let Some(prefix) = deps
1757        .get(index)
1758        .and_then(|item| item.decor().prefix().and_then(|s| s.as_str()))
1759        .filter(|s| !s.is_empty())
1760    {
1761        let prefix = prefix.to_string();
1762        if let Some(next) = deps.get(index + 1)
1763            && let Some(existing) = next.decor().prefix().and_then(|s| s.as_str())
1764        {
1765            // Transfer removed item's prefix to the next item's prefix.
1766            let existing = existing.to_string();
1767            if let Some(next) = deps.get_mut(index + 1) {
1768                next.decor_mut().set_prefix(format!("{prefix}{existing}"));
1769            }
1770        } else if let Some(next) = deps.get_mut(index + 1) {
1771            // Next item exists but has no prefix; use ours directly.
1772            next.decor_mut().set_prefix(&prefix);
1773        } else if let Some(existing) = deps.trailing().as_str() {
1774            // No next item; move comments to the array trailing.
1775            deps.set_trailing(format!("{prefix}{existing}"));
1776        } else {
1777            deps.set_trailing(&prefix);
1778        }
1779    }
1780
1781    deps.remove(index)
1782        .as_str()
1783        .and_then(|req| Requirement::from_str(req).ok())
1784}
1785
1786/// Returns a `Vec` containing the all dependencies with the given name, along with their positions
1787/// in the array.
1788fn find_dependencies(
1789    name: &PackageName,
1790    marker: Option<&MarkerTree>,
1791    deps: &Array,
1792) -> Vec<(usize, Requirement)> {
1793    let mut to_replace = Vec::new();
1794    for (i, dep) in deps.iter().enumerate() {
1795        if let Some(req) = dep.as_str().and_then(try_parse_requirement)
1796            && marker.is_none_or(|m| *m == req.marker)
1797            && *name == req.name
1798        {
1799            to_replace.push((i, req));
1800        }
1801    }
1802    to_replace
1803}
1804
1805/// Return whether two requirements have the same serialized fields, ignoring their parsed origin.
1806fn same_requirement_declaration(left: &Requirement, right: &Requirement) -> bool {
1807    left.name == right.name
1808        && left.extras == right.extras
1809        && left.version_or_url == right.version_or_url
1810        && left.marker == right.marker
1811}
1812
1813/// Returns the key in `tool.uv.sources` that matches the given package name.
1814fn find_source(name: &PackageName, sources: &Table) -> Option<String> {
1815    for (key, _) in sources {
1816        if PackageName::from_str(key).is_ok_and(|ref key| key == name) {
1817            return Some(key.to_string());
1818        }
1819    }
1820    None
1821}
1822
1823// Add a source to `tool.uv.sources`.
1824fn add_source(req: &PackageName, source: &Source, sources: &mut Table) -> Result<(), Error> {
1825    // Serialize as an inline table.
1826    let mut doc = toml::to_string(&source)
1827        .map_err(Box::new)?
1828        .parse::<DocumentMut>()
1829        .unwrap();
1830    let table = mem::take(doc.as_table_mut()).into_inline_table();
1831
1832    sources.insert(req.as_ref(), Item::Value(Value::InlineTable(table)));
1833
1834    Ok(())
1835}
1836
1837impl fmt::Display for PyProjectTomlMut {
1838    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1839        self.doc.fmt(f)
1840    }
1841}
1842
1843fn try_parse_requirement(req: &str) -> Option<Requirement> {
1844    Requirement::from_str(req).ok()
1845}
1846
1847/// Reformats a TOML array to multi line while trying to preserve all comments
1848/// and move them around. This also formats the array to have a trailing comma.
1849fn reformat_array_multiline(deps: &mut Array) {
1850    fn find_comments(s: Option<&RawString>) -> Box<dyn Iterator<Item = Comment> + '_> {
1851        let iter = s
1852            .and_then(|x| x.as_str())
1853            .unwrap_or("")
1854            .lines()
1855            .scan(
1856                (false, false),
1857                |(prev_line_was_empty, prev_line_was_comment), line| {
1858                    let trimmed_line = line.trim();
1859
1860                    if let Some((before, comment)) = line.split_once('#') {
1861                        let comment_text = format!("#{}", comment.trim_end());
1862
1863                        let comment_kind = if (*prev_line_was_empty) || (*prev_line_was_comment) {
1864                            CommentType::OwnLine
1865                        } else {
1866                            CommentType::EndOfLine {
1867                                leading_whitespace: before
1868                                    .chars()
1869                                    .rev()
1870                                    .take_while(|c| c.is_whitespace())
1871                                    .collect::<String>()
1872                                    .chars()
1873                                    .rev()
1874                                    .collect(),
1875                            }
1876                        };
1877
1878                        *prev_line_was_empty = trimmed_line.is_empty();
1879                        *prev_line_was_comment = true;
1880
1881                        Some(Some(Comment {
1882                            text: comment_text,
1883                            kind: comment_kind,
1884                        }))
1885                    } else {
1886                        *prev_line_was_empty = trimmed_line.is_empty();
1887                        *prev_line_was_comment = false;
1888                        Some(None)
1889                    }
1890                },
1891            )
1892            .flatten();
1893
1894        Box::new(iter)
1895    }
1896
1897    // Without a trailing comma, `toml_edit` stores comments after the final item in its
1898    // suffix. Once we add a trailing comma, those comments must follow the comma instead.
1899    if !deps.trailing_comma()
1900        && let Some(last) = deps.iter_mut().last()
1901        && let Some(suffix) = last.decor().suffix().and_then(RawString::as_str)
1902        && suffix.contains('#')
1903    {
1904        let suffix = suffix.to_string();
1905        last.decor_mut().set_suffix("");
1906        let trailing = deps.trailing().as_str().unwrap_or_default();
1907        deps.set_trailing(format!("{suffix}{trailing}"));
1908    }
1909
1910    let mut indentation_prefix = None;
1911
1912    // Calculate the indentation prefix based on the indentation of the first dependency entry.
1913    if let Some(first_item) = deps.iter().next() {
1914        let decor_prefix = first_item
1915            .decor()
1916            .prefix()
1917            .and_then(|s| s.as_str())
1918            .and_then(|s| s.lines().last())
1919            .unwrap_or_default();
1920
1921        let decor_prefix = decor_prefix
1922            .split_once('#')
1923            .map(|(s, _)| s)
1924            .unwrap_or(decor_prefix);
1925
1926        indentation_prefix = (!decor_prefix.is_empty()).then_some(decor_prefix.to_string());
1927    }
1928
1929    let indentation_prefix_str = format!("\n{}", indentation_prefix.as_deref().unwrap_or("    "));
1930
1931    for item in deps.iter_mut() {
1932        let decor = item.decor_mut();
1933        let mut prefix = String::new();
1934
1935        for comment in find_comments(decor.prefix()).chain(find_comments(decor.suffix())) {
1936            match &comment.kind {
1937                CommentType::OwnLine => {
1938                    prefix.push_str(&indentation_prefix_str);
1939                }
1940                CommentType::EndOfLine { leading_whitespace } => {
1941                    prefix.push_str(leading_whitespace);
1942                }
1943            }
1944            prefix.push_str(&comment.text);
1945        }
1946        prefix.push_str(&indentation_prefix_str);
1947        decor.set_prefix(prefix);
1948        decor.set_suffix("");
1949    }
1950
1951    deps.set_trailing(&{
1952        let mut comments = find_comments(Some(deps.trailing())).peekable();
1953        let mut rv = String::new();
1954        if comments.peek().is_some() {
1955            for comment in comments {
1956                match &comment.kind {
1957                    CommentType::OwnLine => {
1958                        let indentation_prefix_str =
1959                            format!("\n{}", indentation_prefix.as_deref().unwrap_or("    "));
1960                        rv.push_str(&indentation_prefix_str);
1961                    }
1962                    CommentType::EndOfLine { leading_whitespace } => {
1963                        rv.push_str(leading_whitespace);
1964                    }
1965                }
1966                rv.push_str(&comment.text);
1967            }
1968        }
1969        if !rv.is_empty() || !deps.is_empty() {
1970            rv.push('\n');
1971        }
1972        rv
1973    });
1974    deps.set_trailing_comma(true);
1975}
1976
1977/// Split a requirement into the package name and its dependency specifiers.
1978///
1979/// E.g., given `flask>=1.0`, this function returns `("flask", ">=1.0")`. But given
1980/// `Flask>=1.0`, this function returns `("Flask", ">=1.0")`.
1981///
1982/// Extras are retained, such that `flask[dotenv]>=1.0` returns `("flask[dotenv]", ">=1.0")`.
1983fn split_specifiers(req: &str) -> (&str, &str) {
1984    let (name, specifiers) = req
1985        .find(['>', '<', '=', '~', '!', '@'])
1986        .map_or((req, ""), |pos| {
1987            let (name, specifiers) = req.split_at(pos);
1988            (name, specifiers)
1989        });
1990    (name.trim(), specifiers.trim())
1991}
1992
1993#[cfg(test)]
1994mod test {
1995    use crate::pyproject::DependencyType;
1996
1997    use super::{
1998        AddBoundsKind, ArrayEdit, DependencyTarget, PyProjectTomlMut, reformat_array_multiline,
1999        remove_dependency, split_specifiers,
2000    };
2001    use anyhow::Result;
2002    use insta::assert_snapshot;
2003    use std::path::Path;
2004    use std::str::FromStr;
2005    use toml_edit::DocumentMut;
2006    use uv_distribution_types::Index;
2007    use uv_normalize::{ExtraName, GroupName, PackageName};
2008    use uv_pep440::Version;
2009    use uv_pep508::{Requirement, RequirementOrigin};
2010
2011    #[test]
2012    fn split() {
2013        assert_eq!(split_specifiers("flask>=1.0"), ("flask", ">=1.0"));
2014        assert_eq!(split_specifiers("Flask>=1.0"), ("Flask", ">=1.0"));
2015        assert_eq!(
2016            split_specifiers("flask[dotenv]>=1.0"),
2017            ("flask[dotenv]", ">=1.0")
2018        );
2019        assert_eq!(split_specifiers("flask[dotenv]"), ("flask[dotenv]", ""));
2020        assert_eq!(
2021            split_specifiers(
2022                "flask @ https://files.pythonhosted.org/packages/af/47/93213ee66ef8fae3b93b3e29206f6b251e65c97bd91d8e1c5596ef15af0a/flask-3.1.0-py3-none-any.whl"
2023            ),
2024            (
2025                "flask",
2026                "@ https://files.pythonhosted.org/packages/af/47/93213ee66ef8fae3b93b3e29206f6b251e65c97bd91d8e1c5596ef15af0a/flask-3.1.0-py3-none-any.whl"
2027            )
2028        );
2029    }
2030
2031    #[test]
2032    fn reformat_preserves_inline_comment_spacing() {
2033        let mut doc: DocumentMut = r#"
2034[project]
2035dependencies = [
2036    "attrs>=25.4.0",     # comment
2037]
2038"#
2039        .parse()
2040        .unwrap();
2041
2042        reformat_array_multiline(
2043            doc["project"]["dependencies"]
2044                .as_array_mut()
2045                .expect("dependencies array"),
2046        );
2047
2048        let serialized = doc.to_string();
2049
2050        assert!(
2051            serialized.contains("\"attrs>=25.4.0\",     # comment"),
2052            "inline comment spacing should be preserved:\n{serialized}"
2053        );
2054    }
2055
2056    #[test]
2057    fn reformat_preserves_inline_comment_without_padding() {
2058        let mut doc: DocumentMut = r#"
2059[project]
2060dependencies = [
2061    "attrs>=25.4.0",#comment
2062]
2063"#
2064        .parse()
2065        .unwrap();
2066
2067        reformat_array_multiline(
2068            doc["project"]["dependencies"]
2069                .as_array_mut()
2070                .expect("dependencies array"),
2071        );
2072
2073        let serialized = doc.to_string();
2074
2075        assert!(
2076            serialized.contains("\"attrs>=25.4.0\",#comment"),
2077            "inline comment spacing without padding should be preserved:\n{serialized}"
2078        );
2079    }
2080
2081    #[test]
2082    fn bound_kind_to_specifiers_exact() {
2083        let tests = [
2084            ("0", "==0"),
2085            ("0.0", "==0.0"),
2086            ("0.0.0", "==0.0.0"),
2087            ("0.1", "==0.1"),
2088            ("0.0.1", "==0.0.1"),
2089            ("0.0.0.1", "==0.0.0.1"),
2090            ("1.0.0", "==1.0.0"),
2091            ("1.2", "==1.2"),
2092            ("1.2.3", "==1.2.3"),
2093            ("1.2.3.4", "==1.2.3.4"),
2094            ("1.2.3.4a1.post1", "==1.2.3.4a1.post1"),
2095        ];
2096
2097        for (version, expected) in tests {
2098            let actual = AddBoundsKind::Exact
2099                .specifiers(Version::from_str(version).unwrap())
2100                .to_string();
2101            assert_eq!(actual, expected, "{version}");
2102        }
2103    }
2104
2105    #[test]
2106    fn bound_kind_to_specifiers_lower() {
2107        let tests = [
2108            ("0", ">=0"),
2109            ("0.0", ">=0.0"),
2110            ("0.0.0", ">=0.0.0"),
2111            ("0.1", ">=0.1"),
2112            ("0.0.1", ">=0.0.1"),
2113            ("0.0.0.1", ">=0.0.0.1"),
2114            ("1", ">=1"),
2115            ("1.0.0", ">=1.0.0"),
2116            ("1.2", ">=1.2"),
2117            ("1.2.3", ">=1.2.3"),
2118            ("1.2.3.4", ">=1.2.3.4"),
2119            ("1.2.3.4a1.post1", ">=1.2.3.4a1.post1"),
2120        ];
2121
2122        for (version, expected) in tests {
2123            let actual = AddBoundsKind::Lower
2124                .specifiers(Version::from_str(version).unwrap())
2125                .to_string();
2126            assert_eq!(actual, expected, "{version}");
2127        }
2128    }
2129
2130    #[test]
2131    fn bound_kind_to_specifiers_major() {
2132        let tests = [
2133            ("0", ">=0, <0.1"),
2134            ("0.0", ">=0.0, <0.1"),
2135            ("0.0.0", ">=0.0.0, <0.1.0"),
2136            ("0.0.0.0", ">=0.0.0.0, <0.1.0.0"),
2137            ("0.1", ">=0.1, <0.2"),
2138            ("0.0.1", ">=0.0.1, <0.0.2"),
2139            ("0.0.1.1", ">=0.0.1.1, <0.0.2.0"),
2140            ("0.0.0.1", ">=0.0.0.1, <0.0.0.2"),
2141            ("1", ">=1, <2"),
2142            ("1.0.0", ">=1.0.0, <2.0.0"),
2143            ("1.2", ">=1.2, <2.0"),
2144            ("1.2.3", ">=1.2.3, <2.0.0"),
2145            ("1.2.3.4", ">=1.2.3.4, <2.0.0.0"),
2146            ("1.2.3.4a1.post1", ">=1.2.3.4a1.post1, <2.0.0.0"),
2147        ];
2148
2149        for (version, expected) in tests {
2150            let actual = AddBoundsKind::Major
2151                .specifiers(Version::from_str(version).unwrap())
2152                .to_string();
2153            assert_eq!(actual, expected, "{version}");
2154        }
2155    }
2156
2157    #[test]
2158    fn bound_kind_to_specifiers_minor() {
2159        let tests = [
2160            ("0", ">=0, <0.0.1"),
2161            ("0.0", ">=0.0, <0.0.1"),
2162            ("0.0.0", ">=0.0.0, <0.0.1"),
2163            ("0.0.0.0", ">=0.0.0.0, <0.0.1.0"),
2164            ("0.1", ">=0.1, <0.1.1"),
2165            ("0.0.1", ">=0.0.1, <0.0.2"),
2166            ("0.0.1.1", ">=0.0.1.1, <0.0.2.0"),
2167            ("0.0.0.1", ">=0.0.0.1, <0.0.0.2"),
2168            ("1", ">=1, <1.1"),
2169            ("1.0.0", ">=1.0.0, <1.1.0"),
2170            ("1.2", ">=1.2, <1.3"),
2171            ("1.2.3", ">=1.2.3, <1.3.0"),
2172            ("1.2.3.4", ">=1.2.3.4, <1.3.0.0"),
2173            ("1.2.3.4a1.post1", ">=1.2.3.4a1.post1, <1.3.0.0"),
2174        ];
2175
2176        for (version, expected) in tests {
2177            let actual = AddBoundsKind::Minor
2178                .specifiers(Version::from_str(version).unwrap())
2179                .to_string();
2180            assert_eq!(actual, expected, "{version}");
2181        }
2182    }
2183
2184    #[test]
2185    fn replace_dependency_updates_every_exact_match() -> Result<()> {
2186        let mut pyproject = PyProjectTomlMut::from_toml(
2187            r#"[project]
2188dependencies = ["anyio<=2", "anyio>=1", "anyio<=2"]
2189
2190[tool.uv.sources]
2191anyio = { index = "internal" }
2192            "#,
2193            DependencyTarget::PyProjectToml,
2194        )?;
2195        let existing = Requirement::from_str("anyio<=2")?.with_origin(RequirementOrigin::Workspace);
2196        let replacement = Requirement::from_str("anyio<3")?;
2197
2198        let replaced = pyproject.replace_dependency_declaration(
2199            &DependencyType::Production,
2200            &existing,
2201            &replacement,
2202        )?;
2203        assert_eq!(replaced, vec![ArrayEdit::Update(0), ArrayEdit::Update(2)]);
2204
2205        assert_snapshot!(
2206            pyproject.to_string(),
2207            @r#"
2208[project]
2209dependencies = ["anyio<3", "anyio>=1", "anyio<3"]
2210
2211[tool.uv.sources]
2212anyio = { index = "internal" }
2213"#
2214        );
2215        Ok(())
2216    }
2217
2218    #[test]
2219    fn replace_dependency_declaration_updates_selected_type() -> Result<()> {
2220        let mut pyproject = PyProjectTomlMut::from_toml(
2221            r#"[project]
2222dependencies = ["anyio<=2"]
2223
2224[project.optional-dependencies]
2225test = ["anyio<=2"]
2226
2227[dependency-groups]
2228dev = ["anyio<=2"]
2229            "#,
2230            DependencyTarget::PyProjectToml,
2231        )?;
2232        let existing = Requirement::from_str("anyio<=2")?;
2233        let optional_replacement = Requirement::from_str("anyio<3")?;
2234        let group_replacement = Requirement::from_str("anyio<4")?;
2235
2236        let replaced = pyproject.replace_dependency_declaration(
2237            &DependencyType::Optional(ExtraName::from_str("test")?),
2238            &existing,
2239            &optional_replacement,
2240        )?;
2241        assert_eq!(replaced, vec![ArrayEdit::Update(0)]);
2242
2243        let replaced = pyproject.replace_dependency_declaration(
2244            &DependencyType::Group(GroupName::from_str("dev")?),
2245            &existing,
2246            &group_replacement,
2247        )?;
2248        assert_eq!(replaced, vec![ArrayEdit::Update(0)]);
2249
2250        assert_snapshot!(
2251            pyproject.to_string(),
2252            @r#"
2253[project]
2254dependencies = ["anyio<=2"]
2255
2256[project.optional-dependencies]
2257test = ["anyio<3"]
2258
2259[dependency-groups]
2260dev = ["anyio<4"]
2261"#
2262        );
2263        Ok(())
2264    }
2265
2266    #[test]
2267    fn remove_preserves_end_of_line_comment_on_previous_item() {
2268        let toml = r#"
2269[project]
2270dependencies = [
2271    "numpy>=2.4.3", # this comment is clearly essential
2272    "requests>=2.32.5",
2273]
2274"#;
2275        let mut doc: DocumentMut = toml.parse().unwrap();
2276        let deps = doc["project"]["dependencies"]
2277            .as_array_mut()
2278            .expect("dependencies array");
2279
2280        let name = PackageName::from_str("requests").unwrap();
2281        remove_dependency(&name, deps);
2282
2283        assert_snapshot!(
2284            doc.to_string(),
2285            @r#"
2286[project]
2287dependencies = [
2288    "numpy>=2.4.3", # this comment is clearly essential
2289]
2290"#
2291        );
2292    }
2293
2294    #[test]
2295    fn remove_preserves_end_of_line_comment_on_previous_item_middle() {
2296        let toml = r#"
2297[project]
2298dependencies = [
2299    "numpy>=2.4.3", # numpy comment
2300    "requests>=2.32.5",
2301    "flask>=3.0.0",
2302]
2303"#;
2304        let mut doc: DocumentMut = toml.parse().unwrap();
2305        let deps = doc["project"]["dependencies"]
2306            .as_array_mut()
2307            .expect("dependencies array");
2308
2309        let name = PackageName::from_str("requests").unwrap();
2310        remove_dependency(&name, deps);
2311
2312        assert_snapshot!(
2313            doc.to_string(),
2314            @r#"
2315[project]
2316dependencies = [
2317    "numpy>=2.4.3", # numpy comment
2318    "flask>=3.0.0",
2319]
2320"#
2321        );
2322    }
2323
2324    #[test]
2325    fn remove_preserves_own_line_comment_above_removed_item() {
2326        let toml = r#"
2327[project]
2328dependencies = [
2329    "numpy>=2.4.3",
2330    # This is a comment about requests
2331    "requests>=2.32.5",
2332]
2333"#;
2334        let mut doc: DocumentMut = toml.parse().unwrap();
2335        let deps = doc["project"]["dependencies"]
2336            .as_array_mut()
2337            .expect("dependencies array");
2338
2339        let name = PackageName::from_str("requests").unwrap();
2340        remove_dependency(&name, deps);
2341
2342        assert_snapshot!(
2343            doc.to_string(),
2344            @r#"
2345[project]
2346dependencies = [
2347    "numpy>=2.4.3",
2348    # This is a comment about requests
2349]
2350"#
2351        );
2352    }
2353
2354    #[test]
2355    fn remove_item_with_trailing_comment_last() {
2356        // When the removed item itself has an end-of-line comment and is the last item,
2357        // toml_edit stores the comment in the array trailing. The comment is preserved
2358        // (as an own-line comment in the trailing section) but moves position since it
2359        // can no longer be on the removed item's line.
2360        let toml = r#"
2361[project]
2362dependencies = [
2363    "requests>=2.32.5",
2364    "numpy>=2.4.3", # comment on numpy
2365]
2366"#;
2367        let mut doc: DocumentMut = toml.parse().unwrap();
2368        let deps = doc["project"]["dependencies"]
2369            .as_array_mut()
2370            .expect("dependencies array");
2371
2372        let name = PackageName::from_str("numpy").unwrap();
2373        remove_dependency(&name, deps);
2374
2375        assert_snapshot!(
2376            doc.to_string(),
2377            @r#"
2378[project]
2379dependencies = [
2380    "requests>=2.32.5",
2381    # comment on numpy
2382]
2383"#
2384        );
2385    }
2386
2387    #[test]
2388    fn remove_last_item_with_trailing_comment_preserves_previous_comment() {
2389        let toml = r#"
2390[project]
2391dependencies = [
2392    "boto3", # this is boto3
2393    "requests", # this is requests
2394]
2395"#;
2396        let mut doc: DocumentMut = toml.parse().unwrap();
2397        let deps = doc["project"]["dependencies"]
2398            .as_array_mut()
2399            .expect("dependencies array");
2400
2401        let name = PackageName::from_str("requests").unwrap();
2402        remove_dependency(&name, deps);
2403
2404        assert_snapshot!(
2405            doc.to_string(),
2406            @r#"
2407[project]
2408dependencies = [
2409    "boto3", # this is boto3
2410    # this is requests
2411]
2412"#
2413        );
2414    }
2415
2416    #[test]
2417    fn remove_item_with_trailing_comment_middle() {
2418        // When the removed item has an end-of-line comment and is in the middle,
2419        // toml_edit stores the comment in the next item's prefix. After removal,
2420        // reformat_array_multiline repositions it as an own-line comment.
2421        let toml = r#"
2422[project]
2423dependencies = [
2424    "requests>=2.32.5",
2425    "numpy>=2.4.3", # comment on numpy
2426    "flask>=3.0.0",
2427]
2428"#;
2429        let mut doc: DocumentMut = toml.parse().unwrap();
2430        let deps = doc["project"]["dependencies"]
2431            .as_array_mut()
2432            .expect("dependencies array");
2433
2434        let name = PackageName::from_str("numpy").unwrap();
2435        remove_dependency(&name, deps);
2436
2437        assert_snapshot!(
2438            doc.to_string(),
2439            @r#"
2440[project]
2441dependencies = [
2442    "requests>=2.32.5",
2443    # comment on numpy
2444    "flask>=3.0.0",
2445]
2446"#
2447        );
2448    }
2449
2450    #[test]
2451    fn remove_first_item_with_trailing_comment_preserves_leading_comments() {
2452        let toml = r#"
2453[project]
2454dependencies = [
2455    # should be in alphabetical order
2456    "basedmypy[faster-cache]>=2.8.1", # this is a comment
2457    "basedpyright>=1.18.2,<2.0.0",
2458]
2459"#;
2460        let mut doc: DocumentMut = toml.parse().unwrap();
2461        let deps = doc["project"]["dependencies"]
2462            .as_array_mut()
2463            .expect("dependencies array");
2464
2465        let name = PackageName::from_str("basedmypy").unwrap();
2466        remove_dependency(&name, deps);
2467
2468        assert_snapshot!(
2469            doc.to_string(),
2470            @r#"
2471[project]
2472dependencies = [
2473    # should be in alphabetical order
2474    # this is a comment
2475    "basedpyright>=1.18.2,<2.0.0",
2476]
2477"#
2478        );
2479    }
2480
2481    #[test]
2482    fn remove_multiple_adjacent_matches_preserves_comment_order() {
2483        let toml = r#"
2484[project]
2485dependencies = [
2486    "iniconfig>=2.0.0", # comment on iniconfig
2487    "typing-extensions>=4.0.0 ; python_version < '3.11'", # comment on first typing-extensions
2488    "typing-extensions>=4.0.0 ; python_version >= '3.11'",
2489    "sniffio>=1.3.0",
2490]
2491"#;
2492        let mut doc: DocumentMut = toml.parse().unwrap();
2493        let deps = doc["project"]["dependencies"]
2494            .as_array_mut()
2495            .expect("dependencies array");
2496
2497        let name = PackageName::from_str("typing-extensions").unwrap();
2498        remove_dependency(&name, deps);
2499
2500        assert_snapshot!(
2501            doc.to_string(),
2502            @r#"
2503[project]
2504dependencies = [
2505    "iniconfig>=2.0.0", # comment on iniconfig
2506    # comment on first typing-extensions
2507    "sniffio>=1.3.0",
2508]
2509"#
2510        );
2511    }
2512
2513    #[test]
2514    fn add_index_syncs_format_on_url_update() {
2515        let toml = r#"
2516[[tool.uv.index]]
2517name = "index"
2518url = "https://example.com/flat/"
2519format = "flat"
2520"#;
2521
2522        let mut doc = PyProjectTomlMut::from_toml(toml, DependencyTarget::PyProjectToml).unwrap();
2523
2524        // The URL spelling changes, but the canonical URL does not, so format should be preserved.
2525        let equivalent_index = Index::from_str("index=https://example.com/flat").unwrap();
2526        doc.add_index(&equivalent_index, Path::new(".")).unwrap();
2527
2528        assert_snapshot!(doc.to_string(), @r#"
2529
2530[[tool.uv.index]]
2531name = "index"
2532url = "https://example.com/flat"
2533format = "flat"
2534"#);
2535
2536        let new_index = Index::from_str("index=https://pypi.org/simple").unwrap();
2537        doc.add_index(&new_index, Path::new(".")).unwrap();
2538
2539        assert_snapshot!(doc.to_string(), @r#"
2540
2541[[tool.uv.index]]
2542name = "index"
2543url = "https://pypi.org/simple"
2544"#);
2545    }
2546
2547    #[cfg(windows)]
2548    #[test]
2549    fn add_index_preserves_format_when_windows_path_unchanged() -> Result<()> {
2550        let toml = r#"
2551[[tool.uv.index]]
2552name = "index"
2553url = 'C:\links'
2554format = "flat"
2555"#;
2556
2557        let mut doc = PyProjectTomlMut::from_toml(toml, DependencyTarget::PyProjectToml)?;
2558
2559        let new_index = Index::from_str(r"index=C:\links")?;
2560        doc.add_index(&new_index, &std::env::current_dir()?)?;
2561
2562        let index = doc.doc["tool"]["uv"]["index"]
2563            .as_array_of_tables()
2564            .and_then(|indexes| indexes.get(0))
2565            .expect("index table");
2566        assert_eq!(
2567            index.get("url").and_then(|item| item.as_str()),
2568            Some("C:/links")
2569        );
2570        assert_eq!(
2571            index.get("format").and_then(|item| item.as_str()),
2572            Some("flat")
2573        );
2574
2575        Ok(())
2576    }
2577}