feat: add a function to extend version with 0s (#689)

This function will extend a given version to a certain length. 

If the version is already longer than the specified length, it will just
return it.

For example:

```
Extending to length of 3:
1 -> 1.0.0
1.2 -> 1.2.0

1!1.2 -> 1!1.2.0
1!1.2+3.4 -> 1!1.2.0+3.4
```

This is in preparation of making the version bumping extend the version.

crates/rattler_conda_types/src/lib.rs

@@ -47,7 +47,7 @@ pub use repo_data_record::RepoDataRecord;
 pub use run_export::RunExportKind;
 pub use version::{
     Component, ParseVersionError, ParseVersionErrorKind, StrictVersion, Version, VersionBumpError,
-    VersionBumpType, VersionWithSource,
+    VersionBumpType, VersionExtendError, VersionWithSource,
 };
 pub use version_spec::VersionSpec;
 

crates/rattler_conda_types/src/version/mod.rs

@@ -28,6 +28,7 @@ pub use bump::{VersionBumpError, VersionBumpType};
 use flags::Flags;
 use segment::Segment;
 
+use thiserror::Error;
 pub use with_source::VersionWithSource;
 
 /// This class implements an order relation between version strings. Version strings can contain the
@@ -168,6 +169,15 @@ pub struct Version {
 type ComponentVec = SmallVec<[Component; 3]>;
 type SegmentVec = SmallVec<[Segment; 4]>;
 
+/// Error that can occur when extending a version to a certain length.
+#[derive(Error, Debug, PartialEq)]
+
+pub enum VersionExtendError {
+    /// The version is too long (there is a maximum number of segments allowed)
+    #[error("the version is too long")]
+    VersionTooLong,
+}
+
 impl Version {
     /// Constructs a version with just a major component and no other components, e.g. "1".
     pub fn major(major: u64) -> Version {
@@ -567,6 +577,51 @@ impl Version {
             Cow::Borrowed(self)
         }
     }
+
+    /// Extend the version to the specified length by adding default components (0s).
+    /// If the version is already longer than the specified length it is returned as is.
+    pub fn extend_to_length(&self, length: usize) -> Result<Cow<'_, Version>, VersionExtendError> {
+        if self.segment_count() >= length {
+            return Ok(Cow::Borrowed(self));
+        }
+
+        // copy everything up to local version
+        let mut segments = self.segments[..self.segment_count()].to_vec();
+        let components_end = segments.iter().map(|s| s.len() as usize).sum::<usize>()
+            + usize::from(self.has_epoch());
+        let mut components = self.components.clone()[..components_end].to_vec();
+
+        // unwrap is OK here because these should be fine to construct
+        let segment = Segment::new(1).unwrap().with_separator(Some('.')).unwrap();
+
+        for _ in 0..(length - self.segment_count()) {
+            components.push(Component::Numeral(0));
+            segments.push(segment);
+        }
+
+        // add local version if it exists
+        let flags = if self.has_local() {
+            let flags = self
+                .flags
+                .with_local_segment_index(segments.len() as u8)
+                .ok_or(VersionExtendError::VersionTooLong)?;
+            for segment_iter in self.local_segments() {
+                for component in segment_iter.components().cloned() {
+                    components.push(component);
+                }
+                segments.push(segment_iter.segment);
+            }
+            flags
+        } else {
+            self.flags
+        };
+
+        Ok(Cow::Owned(Version {
+            components: components.into(),
+            segments: segments.into(),
+            flags,
+        }))
+    }
 }
 
 /// Returns true if the specified segments are considered to start with the other segments.
@@ -1076,6 +1131,7 @@ mod test {
     use std::hash::{Hash, Hasher};
 
     use rand::seq::SliceRandom;
+    use rstest::rstest;
 
     use crate::version::{StrictVersion, VersionBumpError, VersionBumpType};
 
@@ -1599,4 +1655,22 @@ mod test {
 
         assert_eq!(err, VersionBumpError::InvalidSegment { index: -3 });
     }
+
+    #[rstest]
+    #[case("1", 3, "1.0.0")]
+    #[case("1.2", 3, "1.2.0")]
+    #[case("1.2+3.4", 3, "1.2.0+3.4")]
+    #[case("4!1.2+3.4", 3, "4!1.2.0+3.4")]
+    #[case("4!1.2+3.4", 5, "4!1.2.0.0.0+3.4")]
+    #[test]
+    fn extend_to_length(#[case] version: &str, #[case] elements: usize, #[case] expected: &str) {
+        assert_eq!(
+            Version::from_str(version)
+                .unwrap()
+                .extend_to_length(elements)
+                .unwrap()
+                .to_string(),
+            expected
+        );
+    }
 }

py-rattler/rattler/install/installer.py

@@ -35,14 +35,18 @@ async def install(
     ```python
     >>> import asyncio
     >>> from rattler import solve, install
+    >>> from tempfile import TemporaryDirectory
+    >>> temp_dir = TemporaryDirectory()
+    >>>
     >>> async def main():
     ...     # Solve an environment with python 3.9 for the current platform
     ...     records = await solve(channels=["conda-forge"], specs=["python=3.9"])
     ...
-    ...     # Link the environment in the directory `my-env`.
-    ...     await install(records, target_prefix="my-env")
+    ...     # Link the environment in a temporary directory (you can pass any kind of path here).
+    ...     await install(records, target_prefix=temp_dir.name)
     ...
     ...     # That's it! The environment is now created.
+    ...     # You will find Python under `f"{temp_dir.name}/bin/python"` or `f"{temp_dir.name}/python.exe"` on Windows.
     >>> asyncio.run(main())
 
     ```

py-rattler/rattler/version/version.py

@@ -146,6 +146,24 @@ def bump_segment(self, index: int) -> Version:
         """
         return Version._from_py_version(self._version.bump_segment(index))
 
+    def extend_to_length(self, length: int) -> Version:
+        """
+        Returns a new version that is extended with `0s` to the specified length.
+
+        Examples
+        --------
+        ```python
+        >>> v = Version('1')
+        >>> v.extend_to_length(3)
+        Version("1.0.0")
+        >>> v = Version('4!1.2+3.4')
+        >>> v.extend_to_length(4)
+        Version("4!1.2.0.0+3.4")
+        >>>
+        ```
+        """
+        return Version._from_py_version(self._version.extend_to_length(length))
+
     @property
     def has_local(self) -> bool:
         """

py-rattler/src/error.rs

@@ -5,6 +5,7 @@ use rattler::install::TransactionError;
 use rattler_conda_types::{
     ConvertSubdirError, InvalidPackageNameError, ParseArchError, ParseChannelError,
     ParseMatchSpecError, ParsePlatformError, ParseVersionError, VersionBumpError,
+    VersionExtendError,
 };
 use rattler_lock::{ConversionError, ParseCondaLockError};
 use rattler_package_streaming::ExtractError;
@@ -52,6 +53,8 @@ pub enum PyRattlerError {
     #[error(transparent)]
     VersionBumpError(#[from] VersionBumpError),
     #[error(transparent)]
+    VersionExtendError(#[from] VersionExtendError),
+    #[error(transparent)]
     ParseCondaLockError(#[from] ParseCondaLockError),
     #[error(transparent)]
     ConversionError(#[from] ConversionError),
@@ -126,6 +129,9 @@ impl From<PyRattlerError> for PyErr {
             PyRattlerError::VersionBumpError(err) => {
                 VersionBumpException::new_err(pretty_print_error(&err))
             }
+            PyRattlerError::VersionExtendError(err) => {
+                VersionExtendException::new_err(pretty_print_error(&err))
+            }
             PyRattlerError::ParseCondaLockError(err) => {
                 ParseCondaLockException::new_err(pretty_print_error(&err))
             }
@@ -169,6 +175,7 @@ create_exception!(exceptions, TransactionException, PyException);
 create_exception!(exceptions, LinkException, PyException);
 create_exception!(exceptions, ConvertSubdirException, PyException);
 create_exception!(exceptions, VersionBumpException, PyException);
+create_exception!(exceptions, VersionExtendException, PyException);
 create_exception!(exceptions, ParseCondaLockException, PyException);
 create_exception!(exceptions, ConversionException, PyException);
 create_exception!(exceptions, RequirementException, PyException);

py-rattler/src/version/mod.rs

@@ -83,6 +83,17 @@ impl PyVersion {
         })
     }
 
+    /// Extend a version to a specified length by adding `0s` if necessary
+    pub fn extend_to_length(&self, length: usize) -> PyResult<Self> {
+        Ok(Self {
+            inner: self
+                .inner
+                .extend_to_length(length)
+                .map_err(PyRattlerError::from)?
+                .into_owned(),
+        })
+    }
+
     /// Returns a list of segments of the version. It does not contain
     /// the local segment of the version. See `local_segments` for
     /// local segments in version.