fix: add channels to constructor of Environment (#1051)

conda · Feb 6, 2025 · 46afb12 · 46afb12
1 parent dcc0d99
commit 46afb12
Show file tree

Hide file tree

Showing 4 changed files with 38 additions and 29 deletions.
diff --git a/py-rattler/Cargo.lock b/py-rattler/Cargo.lock
diff --git a/py-rattler/rattler/lock/environment.py b/py-rattler/rattler/lock/environment.py
@@ -1,5 +1,6 @@
 from __future__ import annotations
 from typing import Dict, List, Optional
+
 from rattler.lock.channel import LockChannel
 from rattler.lock.package import LockedPackage, PypiLockedPackage
 from rattler.platform.platform import Platform
@@ -15,14 +16,19 @@ class Environment:
 
     _env: PyEnvironment
 
-    def __init__(self, name: str, requirements: Dict[Platform, List[RepoDataRecord]]) -> None:
+    def __init__(
+        self, name: str, requirements: Dict[Platform, List[RepoDataRecord]], channels: List[LockChannel]
+    ) -> None:
         """
         Create a new environment.
         """
         self._env = PyEnvironment(
-            name,
+            name=name,
             # TODO: move this logic to rust
-            {platform._inner: [record._record for record in records] for (platform, records) in requirements.items()},
+            records={
+                platform._inner: [record._record for record in records] for (platform, records) in requirements.items()
+            },
+            channels=[channel._channel for channel in channels],
         )
 
     def platforms(self) -> List[Platform]:

diff --git a/py-rattler/src/lock/mod.rs b/py-rattler/src/lock/mod.rs
@@ -1,3 +1,4 @@
+use crate::channel::PyChannel;
 use crate::match_spec::PyMatchSpec;
 use crate::version::PyVersion;
 use crate::{error::PyRattlerError, platform::PyPlatform, record::PyRecord};
@@ -9,7 +10,7 @@ use rattler_lock::{
     PackageHashes, PypiPackageData, PypiPackageEnvironmentData, DEFAULT_ENVIRONMENT_NAME,
 };
 use std::{
-    collections::{BTreeSet, HashMap, HashSet},
+    collections::{BTreeSet, HashMap},
     path::PathBuf,
     str::FromStr,
 };
@@ -127,21 +128,23 @@ impl PyEnvironment {
 #[pymethods]
 impl PyEnvironment {
     #[new]
-    pub fn new(name: String, req: HashMap<PyPlatform, Vec<PyRecord>>) -> PyResult<Self> {
+    pub fn new(
+        name: String,
+        records: HashMap<PyPlatform, Vec<PyRecord>>,
+        channels: Vec<PyChannel>,
+    ) -> PyResult<Self> {
         let mut lock = LockFile::builder();
-        let channels = req
-            .values()
-            .flat_map(|records| {
-                records
-                    .iter()
-                    .filter_map(|r| r.channel().transpose())
-                    .collect::<Vec<PyResult<_>>>()
-            })
-            .collect::<PyResult<HashSet<_>>>()?;
 
-        lock.set_channels(&name, channels);
+        lock.set_channels(
+            &name,
+            channels.into_iter().map(|c| {
+                rattler_lock::Channel::from(
+                    c.inner.base_url.as_str().trim_end_matches('/').to_string(),
+                )
+            }),
+        );
 
-        for (platform, records) in req {
+        for (platform, records) in records {
             for record in records {
                 lock.add_conda_package(
                     &name,

diff --git a/py-rattler/tests/unit/test_prefix_record.py b/py-rattler/tests/unit/test_prefix_record.py
@@ -117,7 +117,7 @@ def test_prefix_paths() -> None:
         size_in_bytes=1024,
     )
 
-    assert str(prefix_paths_entry.relative_path) == "foo/bar/baz"
+    assert str(prefix_paths_entry.relative_path) == str(Path("foo/bar/baz"))
     assert prefix_paths_entry.path_type.hardlink
     assert prefix_paths_entry.prefix_placeholder == "placeholder_foo_bar"
     assert prefix_paths_entry.file_mode.binary