From 4ac1858337f0039d950dcda5408caea378abe917 Mon Sep 17 00:00:00 2001
From: Hadrien Mary <hadrien.mary@gmail.com>
Date: Wed, 20 Dec 2023 11:44:13 -0500
Subject: [PATCH] python binding for `bump_segment`

---
 py-rattler/rattler/version/version.py | 16 ++++++++++++++++
 py-rattler/src/version/mod.rs         |  8 ++++++++
 py-rattler/tests/unit/test_version.py |  8 ++++++++
 3 files changed, 32 insertions(+)

diff --git a/py-rattler/rattler/version/version.py b/py-rattler/rattler/version/version.py
index 0ce9c2ae8..887ca8218 100644
--- a/py-rattler/rattler/version/version.py
+++ b/py-rattler/rattler/version/version.py
@@ -113,6 +113,22 @@ def bump_last(self) -> Version:
         """
         return Version._from_py_version(self._version.bump_last())
 
+    def bump_segment(self, index: int) -> Version:
+        """
+        Returns a new version where the last segment of this version has
+        been bumped.
+
+        Examples
+        --------
+        ```python
+        >>> v = Version('1.0')
+        >>> v.bump_segment(index=1)
+        Version("1.1")
+        >>>
+        ```
+        """
+        return Version._from_py_version(self._version.bump_segment(index))
+
     @property
     def has_local(self) -> bool:
         """
diff --git a/py-rattler/src/version/mod.rs b/py-rattler/src/version/mod.rs
index 0fb479d80..441547105 100644
--- a/py-rattler/src/version/mod.rs
+++ b/py-rattler/src/version/mod.rs
@@ -164,6 +164,14 @@ impl PyVersion {
         }
     }
 
+    /// Returns a new version where the given segment of this version has been bumped.
+    pub fn bump_segment(&self, index: i32) -> PyResult<Self> {
+        match self.inner.bump(VersionBumpType::Segment(index)) {
+            Ok(v) => Ok(Self { inner: v }),
+            Err(e) => Err(PyRattlerError::from(e).into()),
+        }
+    }
+
     /// Compute the hash of the version.
     fn __hash__(&self) -> u64 {
         let mut hasher = DefaultHasher::new();
diff --git a/py-rattler/tests/unit/test_version.py b/py-rattler/tests/unit/test_version.py
index 7b46ef0b4..f2b24744a 100644
--- a/py-rattler/tests/unit/test_version.py
+++ b/py-rattler/tests/unit/test_version.py
@@ -25,6 +25,8 @@ def test_bump() -> None:
     assert Version("0.5.5").bump_minor() == Version("0.6.5")
     assert Version("0.5.5").bump_patch() == Version("0.5.6")
     assert Version("0.5.5").bump_last() == Version("0.5.6")
+    assert Version("0.5.5").bump_segment(1) == Version("0.6.5")
+    assert Version("0.5.5").bump_segment(-1) == Version("0.5.6")
 
 
 def test_bump_fail() -> None:
@@ -33,3 +35,9 @@ def test_bump_fail() -> None:
 
     with pytest.raises(VersionBumpError):
         Version("1.5").bump_patch()
+
+    with pytest.raises(VersionBumpError):
+        Version("1.5").bump_segment(-5)
+
+    with pytest.raises(VersionBumpError):
+        Version("1.5").bump_segment(5)