feat: allow passing pypi paths (#572)

Allow passing paths as pypi package locations instead of just urls.

Author: baszalmstra

crates/rattler_lock/src/lib.rs

@@ -81,6 +81,7 @@ mod conda;
 mod hash;
 mod parse;
 mod pypi;
+mod url_or_path;
 mod utils;
 
 pub use builder::LockFileBuilder;
@@ -89,6 +90,7 @@ pub use conda::{CondaPackageData, ConversionError};
 pub use hash::PackageHashes;
 pub use parse::ParseCondaLockError;
 pub use pypi::{PypiPackageData, PypiPackageEnvironmentData};
+pub use url_or_path::UrlOrPath;
 
 /// The name of the default environment in a [`LockFile`]. This is the environment name that is used
 /// when no explicit environment name is specified.
@@ -444,11 +446,11 @@ impl Package {
         }
     }
 
-    /// Returns the URL of the package
-    pub fn url(&self) -> &Url {
+    /// Returns the URL or relative path to the package
+    pub fn url_or_path(&self) -> Cow<'_, UrlOrPath> {
         match self {
-            Package::Conda(value) => value.url(),
-            Package::Pypi(value) => value.url(),
+            Self::Conda(value) => Cow::Owned(UrlOrPath::Url(value.url().clone())),
+            Self::Pypi(value) => Cow::Borrowed(value.url()),
         }
     }
 }
@@ -545,8 +547,8 @@ impl PypiPackage {
     }
 
     /// Returns the URL of the package
-    pub fn url(&self) -> &Url {
-        &self.package_data().url
+    pub fn url(&self) -> &UrlOrPath {
+        &self.package_data().url_or_path
     }
 
     /// Returns the extras enabled for this package
@@ -586,6 +588,7 @@ mod test {
     #[case("v4/python-lock.yml")]
     #[case("v4/pypi-matplotlib-lock.yml")]
     #[case("v4/turtlesim-lock.yml")]
+    #[case("v4/path-based-lock.yml")]
     fn test_parse(#[case] file_name: &str) {
         let path = Path::new(env!("CARGO_MANIFEST_DIR"))
             .join("../../test-data/conda-lock")
@@ -608,15 +611,15 @@ mod test {
             .unwrap()
             .packages(Platform::Linux64)
             .unwrap()
-            .map(|p| p.url().clone())
+            .map(|p| p.url_or_path().into_owned())
             .collect::<Vec<_>>());
 
         insta::assert_yaml_snapshot!(conda_lock
             .environment(DEFAULT_ENVIRONMENT_NAME)
             .unwrap()
             .packages(Platform::Osx64)
             .unwrap()
-            .map(|p| p.url().clone())
+            .map(|p| p.url_or_path().into_owned())
             .collect::<Vec<_>>());
     }
 }

crates/rattler_lock/src/parse/deserialize.rs

@@ -1,7 +1,7 @@
 use crate::utils::serde::RawCondaPackageData;
 use crate::{
     Channel, CondaPackageData, EnvironmentData, EnvironmentPackageData, LockFile, LockFileInner,
-    ParseCondaLockError, PypiPackageData, PypiPackageEnvironmentData,
+    ParseCondaLockError, PypiPackageData, PypiPackageEnvironmentData, UrlOrPath,
 };
 use fxhash::FxHashMap;
 use indexmap::IndexSet;
@@ -40,7 +40,7 @@ enum DeserializablePackageSelector {
         conda: Url,
     },
     Pypi {
-        pypi: Url,
+        pypi: UrlOrPath,
         #[serde(flatten)]
         runtime: DeserializablePypiPackageEnvironmentData,
     },
@@ -81,7 +81,7 @@ pub fn parse_from_document(document: Value) -> Result<LockFile, ParseCondaLockEr
     let pypi_url_lookup = pypi_packages
         .iter()
         .enumerate()
-        .map(|(idx, p)| (&p.url, idx))
+        .map(|(idx, p)| (&p.url_or_path, idx))
         .collect::<FxHashMap<_, _>>();
     let mut pypi_runtime_lookup = IndexSet::new();
 
@@ -107,7 +107,7 @@ pub fn parse_from_document(document: Value) -> Result<LockFile, ParseCondaLockEr
                                                     ParseCondaLockError::MissingPackage(
                                                         name.clone(),
                                                         platform,
-                                                        conda,
+                                                        UrlOrPath::Url(conda),
                                                     )
                                                 })?,
                                             )

crates/rattler_lock/src/parse/mod.rs

@@ -2,12 +2,11 @@ mod deserialize;
 mod serialize;
 mod v3;
 
-use super::LockFile;
+use super::{LockFile, UrlOrPath};
 use rattler_conda_types::Platform;
 use serde::de::Error;
 use serde_yaml::Value;
 use std::str::FromStr;
-use url::Url;
 use v3::parse_v3_or_lower;
 
 // Version 2: dependencies are now arrays instead of maps
@@ -31,7 +30,7 @@ pub enum ParseCondaLockError {
     },
 
     #[error("environment {0} and platform {1} refers to a package that does not exist: {2}")]
-    MissingPackage(String, Platform, Url),
+    MissingPackage(String, Platform, UrlOrPath),
 
     #[error(transparent)]
     InvalidPypiPackageName(#[from] pep508_rs::InvalidNameError),

crates/rattler_lock/src/parse/serialize.rs

@@ -1,10 +1,11 @@
 use super::FILE_VERSION;
 use crate::utils::serde::RawCondaPackageData;
-use crate::{Channel, EnvironmentPackageData, LockFile, PypiPackageData};
+use crate::{Channel, EnvironmentPackageData, LockFile, PypiPackageData, UrlOrPath};
 use itertools::Itertools;
 use pep508_rs::ExtraName;
 use rattler_conda_types::Platform;
 use serde::{Serialize, Serializer};
+use std::borrow::Cow;
 use std::collections::{BTreeSet, HashSet};
 use std::{cmp::Ordering, collections::BTreeMap};
 use url::Url;
@@ -37,17 +38,19 @@ enum SerializablePackageSelector<'a> {
         conda: &'a Url,
     },
     Pypi {
-        pypi: &'a Url,
+        pypi: &'a UrlOrPath,
         #[serde(skip_serializing_if = "BTreeSet::is_empty")]
         extras: &'a BTreeSet<ExtraName>,
     },
 }
 
 impl<'a> SerializablePackageSelector<'a> {
-    fn url(&self) -> &Url {
+    fn url(&self) -> Cow<'_, UrlOrPath> {
         match self {
-            SerializablePackageSelector::Conda { conda } => conda,
-            SerializablePackageSelector::Pypi { pypi, .. } => pypi,
+            SerializablePackageSelector::Conda { conda } => {
+                Cow::Owned(UrlOrPath::Url((*conda).clone()))
+            }
+            SerializablePackageSelector::Pypi { pypi, .. } => Cow::Borrowed(pypi),
         }
     }
 }
@@ -78,34 +81,41 @@ impl<'a> Ord for SerializablePackageSelector<'a> {
             (
                 SerializablePackageSelector::Conda { conda: a },
                 SerializablePackageSelector::Conda { conda: b },
-            )
-            | (
+            ) => compare_url_by_filename(a, b),
+            (
                 SerializablePackageSelector::Pypi { pypi: a, .. },
                 SerializablePackageSelector::Pypi { pypi: b, .. },
-            ) => {
-                // First sort packages just by their filename. Since most of the time the urls end
-                // in the packages filename this causes the urls to be sorted by package name.
-                if let (Some(a), Some(b)) = (
-                    a.path_segments()
-                        .and_then(Iterator::last)
-                        .map(str::to_lowercase),
-                    b.path_segments()
-                        .and_then(Iterator::last)
-                        .map(str::to_lowercase),
-                ) {
-                    match a.cmp(&b) {
-                        Ordering::Equal => {}
-                        ordering => return ordering,
-                    }
-                }
-
-                // Otherwise just sort by their full URL
-                a.cmp(b)
-            }
+            ) => match (a, b) {
+                (UrlOrPath::Url(a), UrlOrPath::Url(b)) => compare_url_by_filename(a, b),
+                (UrlOrPath::Url(_), UrlOrPath::Path(_)) => Ordering::Less,
+                (UrlOrPath::Path(_), UrlOrPath::Url(_)) => Ordering::Greater,
+                (UrlOrPath::Path(a), UrlOrPath::Path(b)) => a.cmp(b),
+            },
         }
     }
 }
 
+/// First sort packages just by their filename. Since most of the time the urls end
+/// in the packages filename this causes the urls to be sorted by package name.
+fn compare_url_by_filename(a: &Url, b: &Url) -> Ordering {
+    if let (Some(a), Some(b)) = (
+        a.path_segments()
+            .and_then(Iterator::last)
+            .map(str::to_lowercase),
+        b.path_segments()
+            .and_then(Iterator::last)
+            .map(str::to_lowercase),
+    ) {
+        match a.cmp(&b) {
+            Ordering::Equal => {}
+            ordering => return ordering,
+        }
+    }
+
+    // Otherwise just sort by their full URL
+    a.cmp(b)
+}
+
 impl<'a> SerializablePackageData<'a> {
     fn source_name(&self) -> &str {
         match self {
@@ -114,10 +124,12 @@ impl<'a> SerializablePackageData<'a> {
         }
     }
 
-    fn url(&self) -> &Url {
+    fn url(&self) -> Cow<'_, UrlOrPath> {
         match self {
-            SerializablePackageData::Conda(p) => &p.url,
-            SerializablePackageData::Pypi(p) => &p.url,
+            SerializablePackageData::Conda(p) => {
+                Cow::Owned(UrlOrPath::Url(p.url.clone().into_owned()))
+            }
+            SerializablePackageData::Pypi(p) => Cow::Borrowed(&p.url_or_path),
         }
     }
 }
@@ -197,7 +209,7 @@ impl Serialize for LockFile {
                                                     .pypi_environment_package_datas
                                                     [pypi_runtime_index];
                                                 SerializablePackageSelector::Pypi {
-                                                    pypi: &pypi_package.url,
+                                                    pypi: &pypi_package.url_or_path,
                                                     extras: &pypi_runtime.extras,
                                                 }
                                             }

crates/rattler_lock/src/parse/v3.rs

@@ -3,7 +3,8 @@
 use super::ParseCondaLockError;
 use crate::{
     Channel, CondaPackageData, EnvironmentData, EnvironmentPackageData, LockFile, LockFileInner,
-    PackageHashes, PypiPackageData, PypiPackageEnvironmentData, DEFAULT_ENVIRONMENT_NAME,
+    PackageHashes, PypiPackageData, PypiPackageEnvironmentData, UrlOrPath,
+    DEFAULT_ENVIRONMENT_NAME,
 };
 use fxhash::FxHashMap;
 use indexmap::IndexSet;
@@ -182,7 +183,7 @@ pub fn parse_v3_or_lower(document: serde_yaml::Value) -> Result<LockFile, ParseC
                         version: pkg.version,
                         requires_dist: pkg.requires_dist,
                         requires_python: pkg.requires_python,
-                        url: pkg.url,
+                        url_or_path: UrlOrPath::Url(pkg.url),
                         hash: pkg.hash,
                     })
                     .0;

crates/rattler_lock/src/pypi.rs

@@ -1,11 +1,10 @@
-use crate::PackageHashes;
+use crate::{PackageHashes, UrlOrPath};
 use pep440_rs::VersionSpecifiers;
 use pep508_rs::{ExtraName, PackageName, Requirement};
 use serde::{Deserialize, Serialize};
 use serde_with::{serde_as, skip_serializing_none};
 use std::cmp::Ordering;
 use std::collections::BTreeSet;
-use url::Url;
 
 /// A pinned Pypi package
 #[serde_as]
@@ -19,7 +18,8 @@ pub struct PypiPackageData {
     pub version: pep440_rs::Version,
 
     /// The URL that points to where the artifact can be downloaded from.
-    pub url: Url,
+    #[serde(with = "crate::utils::serde::url_or_path", flatten)]
+    pub url_or_path: UrlOrPath,
 
     /// Hashes of the file pointed to by `url`.
     #[serde(flatten)]