feat: add a with_alpha function that adds 0a0 to the version (#696)

I also added a `remove_local` function that removes the local segments
of a version.

crates/rattler_conda_types/src/version/bump.rs

@@ -1,3 +1,5 @@
+use std::borrow::Cow;
+
 use thiserror::Error;
 
 use crate::{Component, Version};
@@ -35,6 +37,67 @@ pub enum VersionBumpError {
 }
 
 impl Version {
+    /// Add alpha specifier to the end of the version when the last element does not contain an `iden` component.
+    ///
+    /// For example, `1.0.0` will become `1.0.0.0a0`.
+    /// If the last version element contains a character, it's not modified (e.g. `1.0.0a` will remain `1.0.0a`).
+    pub fn with_alpha(&self) -> Cow<'_, Self> {
+        let last_segment = self.segments().last().expect("at least one segment");
+        // check if there is an iden component in the last segment
+        let has_iden = last_segment.components().any(|c| c.as_iden().is_some());
+        if has_iden {
+            return Cow::Borrowed(self);
+        }
+        let local_segment_index = self.local_segment_index().unwrap_or(self.segments.len());
+        let mut segments = self.segments[0..local_segment_index].to_vec();
+        let components_offset = segments.iter().map(|s| s.len() as usize).sum::<usize>()
+            + usize::from(self.has_epoch());
+
+        segments.push(Segment::new(3).unwrap().with_separator(Some('.')).unwrap());
+        segments.extend(self.segments[local_segment_index..].iter());
+
+        let mut components = self.components.clone();
+        components.insert(components_offset, Component::Numeral(0));
+        components.insert(components_offset + 1, Component::Iden("a".into()));
+        components.insert(components_offset + 2, Component::Numeral(0));
+
+        let flags = if let Some(local_segment_index) = self.local_segment_index() {
+            self.flags
+                .with_local_segment_index((local_segment_index + 1) as u8)
+                .unwrap()
+        } else {
+            self.flags
+        };
+
+        Cow::Owned(Version {
+            components,
+            segments: segments.into(),
+            flags,
+        })
+    }
+
+    /// Remove the local segment from the version if it exists.
+    /// Returns a new version without the local segment.
+    ///
+    /// For example, `1.0.0+3.4` will become `1.0.0`.
+    pub fn remove_local(&self) -> Cow<'_, Self> {
+        if let Some(local_segment_index) = self.local_segment_index() {
+            let segments = self.segments[0..local_segment_index].to_vec();
+            let components_offset = segments.iter().map(|s| s.len() as usize).sum::<usize>()
+                + usize::from(self.has_epoch());
+            let mut components = self.components.clone();
+            components.drain(components_offset..);
+
+            Cow::Owned(Version {
+                components,
+                segments: segments.into(),
+                flags: self.flags.with_local_segment_index(0).unwrap(),
+            })
+        } else {
+            return Cow::Borrowed(self);
+        }
+    }
+
     /// Returns a new version after bumping it according to the specified bump type.
     /// Note: if a version ends with a character, the next bigger version will use `a` as the character.
     /// For example: `1.1l` -> `1.2a`, but also `1.1.0alpha` -> `1.1.1a`.
@@ -236,4 +299,39 @@ mod test {
             Version::from_str(expected).unwrap()
         );
     }
+
+    #[rstest]
+    #[case(0, "1.1.9", "2.1.9.0a0")]
+    #[case(2, "1.0.0", "1.0.1.0a0")]
+    #[case(2, "1.0.0a", "1.0.1a")]
+    #[case(2, "1.0.0f", "1.0.1a")]
+    #[case(2, "5!1.0.0", "5!1.0.1.0a0")]
+    #[case(2, "5!1.0.0+3.4", "5!1.0.1.0a0+3.4")]
+    fn with_alpha(#[case] idx: i32, #[case] input: &str, #[case] expected: &str) {
+        assert_eq!(
+            Version::from_str(input)
+                .unwrap()
+                .bump(VersionBumpType::Segment(idx))
+                .unwrap()
+                .with_alpha()
+                .into_owned(),
+            Version::from_str(expected).unwrap()
+        );
+    }
+
+    #[rstest]
+    #[case("1.1.9", "1.1.9")]
+    #[case("1.0.0+3", "1.0.0")]
+    #[case("1.0.0+3.4", "1.0.0")]
+    #[case("1.0.0+3.4alpha.2.4", "1.0.0")]
+    #[case("5!1.0.0+3.4alpha.2.4", "5!1.0.0")]
+    fn remove_local(#[case] input: &str, #[case] expected: &str) {
+        assert_eq!(
+            Version::from_str(input)
+                .unwrap()
+                .remove_local()
+                .into_owned(),
+            Version::from_str(expected).unwrap()
+        );
+    }
 }

py-rattler/rattler/version/version.py

@@ -149,6 +149,45 @@ def bump_segment(self, index: int) -> Version:
         """
         return Version._from_py_version(self._version.bump_segment(index))
 
+    def with_alpha(self) -> Version:
+        """
+        Returns a new version where the last segment of this version has
+        been bumped with an alpha character. If the last segment contains a
+        character, nothing is added.
+
+        Examples
+        --------
+        ```python
+        >>> v = Version('1.0')
+        >>> v.with_alpha()
+        Version("1.0.0a0")
+        >>> v = Version('1.0.f')
+        >>> v.with_alpha()
+        Version("1.0.f")
+        >>>
+        ```
+        """
+        return Version._from_py_version(self._version.with_alpha())
+
+    def remove_local(self) -> Version:
+        """
+        Returns a new version where the local segment of the version has been removed.
+        Leaves the version unchanged if it does not have a local segment.
+
+        Examples
+        --------
+        ```python
+        >>> v = Version('1.0+3.4')
+        >>> v.remove_local()
+        Version("1.0")
+        >>> v = Version('1.0')
+        >>> v.remove_local()
+        Version("1.0")
+        >>>
+        ```
+        """
+        return Version._from_py_version(self._version.remove_local())
+
     def extend_to_length(self, length: int) -> Version:
         """
         Returns a new version that is extended with `0s` to the specified length.

py-rattler/src/version/mod.rs

@@ -184,6 +184,20 @@ impl PyVersion {
             .map_err(PyRattlerError::from)?)
     }
 
+    /// Returns a new version where the last segment is an "alpha" segment (ie. `.0a0`)
+    pub fn with_alpha(&self) -> Self {
+        Self {
+            inner: self.inner.with_alpha().into_owned(),
+        }
+    }
+
+    /// Returns a new version where the local segment is removed (e.g. `1.0+local` -> `1.0`)
+    pub fn remove_local(&self) -> Self {
+        Self {
+            inner: self.inner.remove_local().into_owned(),
+        }
+    }
+
     /// Compute the hash of the version.
     fn __hash__(&self) -> u64 {
         let mut hasher = DefaultHasher::new();