feat: use modified branch with environment variable passing

examples/pypi/pixi.toml

@@ -32,7 +32,7 @@ black = { version = "~=23.10", extras = ["jupyter"] }
 pyliblzfse = "*"
 pycosat = "*"
 # TODO(tim): fix env-passing in uv
-#env_test_package = "==0.0.3"
+env_test_package = "==0.0.3"
 plot-antenna = "==1.8"
 
 [system-requirements]

src/environment.rs

@@ -180,14 +180,14 @@ pub async fn update_prefix_pypi(
     status: &PythonStatus,
     system_requirements: &SystemRequirements,
     uv_context: UvResolutionContext,
+    environment_variables: &HashMap<String, String>,
 ) -> miette::Result<()> {
     // Remove python packages from a previous python distribution if the python version changed.
-    // install_pypi::remove_old_python_distributions(prefix, platform, status)?;
 
     // Install and/or remove python packages
     progress::await_in_progress(
         format!(
-            "updating pypi package in '{}'",
+            "updating pypi packages in '{}'",
             environment_name.fancy_display()
         ),
         |_| {
@@ -198,6 +198,7 @@ pub async fn update_prefix_pypi(
                 status,
                 system_requirements,
                 uv_context,
+                environment_variables,
             )
         },
     )

src/install_pypi.rs

@@ -20,6 +20,7 @@ use install_wheel_rs::linker::LinkMode;
 use rattler_conda_types::{Platform, RepoDataRecord};
 use rattler_lock::{PypiPackageData, PypiPackageEnvironmentData};
 
+use std::collections::HashMap;
 use std::time::Duration;
 
 use uv_client::{FlatIndex, FlatIndexClient};
@@ -29,7 +30,7 @@ use uv_installer::{Downloader, SitePackages};
 use uv_interpreter::{Interpreter, Virtualenv};
 use uv_normalize::PackageName;
 
-use uv_traits::{NoBinary, NoBuild, SetupPyStrategy};
+use uv_traits::{ConfigSettings, NoBinary, NoBuild, SetupPyStrategy};
 
 type CombinedPypiPackageData = (PypiPackageData, PypiPackageEnvironmentData);
 
@@ -233,6 +234,7 @@ pub async fn update_python_distributions(
     status: &PythonStatus,
     system_requirements: &SystemRequirements,
     uv_context: UvResolutionContext,
+    environment_variables: &HashMap<String, String>,
 ) -> miette::Result<()> {
     let start = std::time::Instant::now();
     let Some(python_info) = status.current_info() else {
@@ -252,7 +254,7 @@ pub async fn update_python_distributions(
 
     let platform = platform_host::Platform::current().expect("unsupported platform");
     let interpreter =
-        Interpreter::query(&python_location, &platform, &uv_context.cache).into_diagnostic()?;
+        Interpreter::query(&python_location, platform, &uv_context.cache).into_diagnostic()?;
 
     tracing::debug!("[Install] Using Python Interpreter: {:?}", interpreter);
 
@@ -281,6 +283,7 @@ pub async fn update_python_distributions(
     let no_binary = NoBinary::None;
 
     let in_memory_index = InMemoryIndex::default();
+    let config_settings = ConfigSettings::default();
 
     // Prep the build context.
     let build_dispatch = BuildDispatch::new(
@@ -291,11 +294,12 @@ pub async fn update_python_distributions(
         &flat_index,
         &in_memory_index,
         &uv_context.in_flight,
-        venv.python_executable(),
         SetupPyStrategy::default(),
+        &config_settings,
         &no_build,
         &no_binary,
-    );
+    )
+    .with_sdist_build_env_vars(environment_variables.iter());
 
     let _lock = venv.lock().into_diagnostic()?;
     // TODO: need to resolve editables?

src/lock_file/resolve.rs

@@ -50,7 +50,7 @@ use uv_resolver::{
     DefaultResolverProvider, DistFinder, InMemoryIndex, Manifest, Options, PythonRequirement,
     Resolver, ResolverProvider, VersionMap, VersionsResponse,
 };
-use uv_traits::{BuildContext, InFlight, NoBinary, NoBuild, SetupPyStrategy};
+use uv_traits::{BuildContext, ConfigSettings, InFlight, NoBinary, NoBuild, SetupPyStrategy};
 
 /// Objects that are needed for resolutions which can be shared between different resolutions.
 #[derive(Clone)]
@@ -214,7 +214,7 @@ pub async fn resolve_pypi(
     platform: rattler_conda_types::Platform,
     pb: &ProgressBar,
     python_location: &Path,
-    _venv_root: &Path,
+    env_variables: &HashMap<String, String>,
 ) -> miette::Result<LockedPypiPackages> {
     // Solve python packages
     pb.set_message("resolving pypi dependencies");
@@ -275,7 +275,7 @@ pub async fn resolve_pypi(
     // TODO: Should we look into using the actual interpreter here?
     let platform = Platform::current().expect("unsupported platform");
     let interpreter =
-        Interpreter::query(python_location, &platform, &context.cache).into_diagnostic()?;
+        Interpreter::query(python_location, platform, &context.cache).into_diagnostic()?;
 
     tracing::debug!("[Resolve] Using Python Interpreter: {:?}", interpreter);
 
@@ -290,6 +290,7 @@ pub async fn resolve_pypi(
     };
 
     let in_memory_index = InMemoryIndex::default();
+    let config_settings = ConfigSettings::default();
 
     // Create a shared in-memory index.
     let options = Options::default();
@@ -301,12 +302,13 @@ pub async fn resolve_pypi(
         &flat_index,
         &in_memory_index,
         &context.in_flight,
-        interpreter.sys_executable().to_path_buf(),
         SetupPyStrategy::default(),
+        &config_settings,
         &NoBuild::None,
         &NoBinary::None,
     )
-    .with_options(options);
+    .with_options(options)
+    .with_sdist_build_env_vars(env_variables.iter());
 
     let constraints = conda_python_packages
         .values()

src/lock_file/update.rs

@@ -117,7 +117,7 @@ impl<'p> LockFileDerivedData<'p> {
         };
 
         // TODO(tim): get support for this somehow with uv
-        let _env_variables = environment.project().get_env_variables(environment).await?;
+        let env_variables = environment.project().get_env_variables(environment).await?;
 
         // Update the prefix with Pypi records
         environment::update_prefix_pypi(
@@ -129,6 +129,7 @@ impl<'p> LockFileDerivedData<'p> {
             &python_status,
             &environment.system_requirements(),
             uv_context,
+            env_variables,
         )
         .await?;
 
@@ -782,7 +783,7 @@ pub async fn ensure_up_to_date_lock_file(
                     Some(context) => context.clone(),
                 };
                 // Get environment variables from the activation
-                let _env_variables = project.get_env_variables(&environment).await?;
+                let env_variables = project.get_env_variables(&environment).await?;
 
                 // Spawn a task to solve the pypi environment
                 let pypi_solve_future = spawn_solve_pypi_task(
@@ -791,6 +792,7 @@ pub async fn ensure_up_to_date_lock_file(
                     platform,
                     repodata_future,
                     prefix_future,
+                    env_variables,
                 );
 
                 pending_futures.push(pypi_solve_future.boxed_local());
@@ -1264,6 +1266,7 @@ async fn spawn_solve_pypi_task(
     platform: Platform,
     repodata_records: impl Future<Output = Arc<RepoDataRecordsByName>>,
     prefix: impl Future<Output = (Prefix, PythonStatus)>,
+    env_variables: &HashMap<String, String>,
 ) -> miette::Result<TaskResult> {
     // Get the Pypi dependencies for this environment
     let dependencies = environment.pypi_dependencies(Some(platform));
@@ -1308,7 +1311,7 @@ async fn spawn_solve_pypi_task(
             platform,
             &pb.pb,
             &python_path,
-            prefix.root(),
+            env_variables,
         )
         .await
         .with_context(|| {