Move update_environment from run to the project namespace (#3659)

Prompted by #3657 (comment) There's still some level of discomfort here, as the `tool` module needs needs to import the `project` module to manage an environment. We should probably move most of the basic operations in the `project` module root into some sort of shared module for behind the scenes operations? Regardless, this change should simplify that future move.
astral-sh · May 20, 2024 · d8971c1 · d8971c1
1 parent d7dc184
commit d8971c1
Show file tree

Hide file tree

Showing 2 changed files with 171 additions and 173 deletions.
diff --git a/crates/uv/src/commands/project/mod.rs b/crates/uv/src/commands/project/mod.rs
@@ -9,23 +9,27 @@ use install_wheel_rs::linker::LinkMode;
 use pep508_rs::MarkerEnvironment;
 use platform_tags::Tags;
 use pypi_types::Yanked;
+use tracing::debug;
 use uv_cache::Cache;
-use uv_client::RegistryClient;
-use uv_configuration::{Concurrency, Constraints, NoBinary, Overrides, Reinstall};
+use uv_client::{BaseClientBuilder, RegistryClient, RegistryClientBuilder};
+use uv_configuration::{
+    Concurrency, ConfigSettings, Constraints, NoBinary, NoBuild, Overrides, PreviewMode, Reinstall,
+    SetupPyStrategy,
+};
 use uv_dispatch::BuildDispatch;
 use uv_distribution::DistributionDatabase;
 use uv_fs::Simplified;
-use uv_installer::{Downloader, Plan, Planner, SitePackages};
+use uv_installer::{Downloader, Plan, Planner, SatisfiesResult, SitePackages};
 use uv_interpreter::{find_default_python, Interpreter, PythonEnvironment};
 use uv_requirements::{
-    ExtrasSpecification, LookaheadResolver, NamedRequirementsResolver, RequirementsSpecification,
-    SourceTreeResolver,
+    ExtrasSpecification, LookaheadResolver, NamedRequirementsResolver, RequirementsSource,
+    RequirementsSpecification, SourceTreeResolver,
 };
 use uv_resolver::{
-    Exclusions, FlatIndex, InMemoryIndex, Manifest, Options, PythonRequirement, ResolutionGraph,
-    Resolver,
+    Exclusions, FlatIndex, InMemoryIndex, Manifest, Options, OptionsBuilder, PythonRequirement,
+    ResolutionGraph, Resolver,
 };
-use uv_types::{HashStrategy, InFlight, InstalledPackagesProvider};
+use uv_types::{BuildIsolation, HashStrategy, InFlight, InstalledPackagesProvider};
 
 use crate::commands::project::discovery::Project;
 use crate::commands::reporters::{DownloadReporter, InstallReporter, ResolverReporter};
@@ -434,3 +438,155 @@ pub(crate) async fn install(
 
     Ok(())
 }
+
+/// Update a [`PythonEnvironment`] to satisfy a set of [`RequirementsSource`]s.
+async fn update_environment(
+    venv: PythonEnvironment,
+    requirements: &[RequirementsSource],
+    preview: PreviewMode,
+    cache: &Cache,
+    printer: Printer,
+) -> Result<PythonEnvironment> {
+    // TODO(zanieb): Support client configuration
+    let client_builder = BaseClientBuilder::default();
+
+    // Read all requirements from the provided sources.
+    // TODO(zanieb): Consider allowing constraints and extras
+    // TODO(zanieb): Allow specifying extras somehow
+    let spec = RequirementsSpecification::from_sources(
+        requirements,
+        &[],
+        &[],
+        &ExtrasSpecification::None,
+        &client_builder,
+        preview,
+    )
+    .await?;
+
+    // Check if the current environment satisfies the requirements
+    let site_packages = SitePackages::from_executable(&venv)?;
+
+    // If the requirements are already satisfied, we're done.
+    if spec.source_trees.is_empty() {
+        match site_packages.satisfies(&spec.requirements, &spec.editables, &spec.constraints)? {
+            SatisfiesResult::Fresh {
+                recursive_requirements,
+            } => {
+                debug!(
+                    "All requirements satisfied: {}",
+                    recursive_requirements
+                        .iter()
+                        .map(|entry| entry.requirement.to_string())
+                        .sorted()
+                        .join(" | ")
+                );
+                debug!(
+                    "All editables satisfied: {}",
+                    spec.editables.iter().map(ToString::to_string).join(", ")
+                );
+                return Ok(venv);
+            }
+            SatisfiesResult::Unsatisfied(requirement) => {
+                debug!("At least one requirement is not satisfied: {requirement}");
+            }
+        }
+    }
+
+    // Determine the tags, markers, and interpreter to use for resolution.
+    let interpreter = venv.interpreter().clone();
+    let tags = venv.interpreter().tags()?;
+    let markers = venv.interpreter().markers();
+
+    // Initialize the registry client.
+    // TODO(zanieb): Support client options e.g. offline, tls, etc.
+    let client = RegistryClientBuilder::new(cache.clone())
+        .markers(markers)
+        .platform(venv.interpreter().platform())
+        .build();
+
+    // TODO(charlie): Respect project configuration.
+    let build_isolation = BuildIsolation::default();
+    let config_settings = ConfigSettings::default();
+    let flat_index = FlatIndex::default();
+    let hasher = HashStrategy::default();
+    let in_flight = InFlight::default();
+    let index = InMemoryIndex::default();
+    let index_locations = IndexLocations::default();
+    let link_mode = LinkMode::default();
+    let no_binary = NoBinary::default();
+    let no_build = NoBuild::default();
+    let setup_py = SetupPyStrategy::default();
+    let concurrency = Concurrency::default();
+
+    // Create a build dispatch.
+    let build_dispatch = BuildDispatch::new(
+        &client,
+        cache,
+        &interpreter,
+        &index_locations,
+        &flat_index,
+        &index,
+        &in_flight,
+        setup_py,
+        &config_settings,
+        build_isolation,
+        link_mode,
+        &no_build,
+        &no_binary,
+        concurrency,
+    );
+
+    let options = OptionsBuilder::new()
+        // TODO(zanieb): Support resolver options
+        // .resolution_mode(resolution_mode)
+        // .prerelease_mode(prerelease_mode)
+        // .dependency_mode(dependency_mode)
+        // .exclude_newer(exclude_newer)
+        .build();
+
+    // Resolve the requirements.
+    let resolution = match resolve(
+        spec,
+        site_packages.clone(),
+        &hasher,
+        &interpreter,
+        tags,
+        markers,
+        &client,
+        &flat_index,
+        &index,
+        &build_dispatch,
+        options,
+        printer,
+        concurrency,
+    )
+    .await
+    {
+        Ok(resolution) => Resolution::from(resolution),
+        Err(err) => return Err(err.into()),
+    };
+
+    // Re-initialize the in-flight map.
+    let in_flight = InFlight::default();
+
+    // Sync the environment.
+    install(
+        &resolution,
+        site_packages,
+        &no_binary,
+        link_mode,
+        &index_locations,
+        &hasher,
+        tags,
+        &client,
+        &in_flight,
+        &build_dispatch,
+        cache,
+        &venv,
+        printer,
+        concurrency,
+    )
+    .await?;
+
+    Ok(venv)
+}
diff --git a/crates/uv/src/commands/project/run.rs b/crates/uv/src/commands/project/run.rs
@@ -7,19 +7,10 @@ use tempfile::tempdir_in;
 use tokio::process::Command;
 use tracing::debug;
 
-use distribution_types::{IndexLocations, Resolution};
-use install_wheel_rs::linker::LinkMode;
 use uv_cache::Cache;
-use uv_client::{BaseClientBuilder, RegistryClientBuilder};
-use uv_configuration::{
-    Concurrency, ConfigSettings, NoBinary, NoBuild, PreviewMode, SetupPyStrategy,
-};
-use uv_dispatch::BuildDispatch;
-use uv_installer::{SatisfiesResult, SitePackages};
+use uv_configuration::PreviewMode;
 use uv_interpreter::PythonEnvironment;
-use uv_requirements::{ExtrasSpecification, RequirementsSource, RequirementsSpecification};
-use uv_resolver::{FlatIndex, InMemoryIndex, OptionsBuilder};
-use uv_types::{BuildIsolation, HashStrategy, InFlight};
+use uv_requirements::RequirementsSource;
 use uv_warnings::warn_user;
 
 use crate::commands::project::discovery::Project;
@@ -73,7 +64,10 @@ pub(crate) async fn run(
         let venv = project::init(&project, cache, printer)?;
 
         // Install the project requirements.
-        Some(update_environment(venv, &project.requirements(), preview, cache, printer).await?)
+        Some(
+            project::update_environment(venv, &project.requirements(), preview, cache, printer)
+                .await?,
+        )
     };
 
     // If necessary, create an environment for the ephemeral requirements.
@@ -111,7 +105,7 @@ pub(crate) async fn run(
         )?;
 
         // Install the ephemeral requirements.
-        Some(update_environment(venv, &requirements, preview, cache, printer).await?)
+        Some(project::update_environment(venv, &requirements, preview, cache, printer).await?)
     };
 
     // Construct the command
@@ -183,155 +177,3 @@ pub(crate) async fn run(
         Ok(ExitStatus::Failure)
     }
 }
-
-/// Update a [`PythonEnvironment`] to satisfy a set of [`RequirementsSource`]s.
-async fn update_environment(
-    venv: PythonEnvironment,
-    requirements: &[RequirementsSource],
-    preview: PreviewMode,
-    cache: &Cache,
-    printer: Printer,
-) -> Result<PythonEnvironment> {
-    // TODO(zanieb): Support client configuration
-    let client_builder = BaseClientBuilder::default();
-
-    // Read all requirements from the provided sources.
-    // TODO(zanieb): Consider allowing constraints and extras
-    // TODO(zanieb): Allow specifying extras somehow
-    let spec = RequirementsSpecification::from_sources(
-        requirements,
-        &[],
-        &[],
-        &ExtrasSpecification::None,
-        &client_builder,
-        preview,
-    )
-    .await?;
-
-    // Check if the current environment satisfies the requirements
-    let site_packages = SitePackages::from_executable(&venv)?;
-
-    // If the requirements are already satisfied, we're done.
-    if spec.source_trees.is_empty() {
-        match site_packages.satisfies(&spec.requirements, &spec.editables, &spec.constraints)? {
-            SatisfiesResult::Fresh {
-                recursive_requirements,
-            } => {
-                debug!(
-                    "All requirements satisfied: {}",
-                    recursive_requirements
-                        .iter()
-                        .map(|entry| entry.requirement.to_string())
-                        .sorted()
-                        .join(" | ")
-                );
-                debug!(
-                    "All editables satisfied: {}",
-                    spec.editables.iter().map(ToString::to_string).join(", ")
-                );
-                return Ok(venv);
-            }
-            SatisfiesResult::Unsatisfied(requirement) => {
-                debug!("At least one requirement is not satisfied: {requirement}");
-            }
-        }
-    }
-
-    // Determine the tags, markers, and interpreter to use for resolution.
-    let interpreter = venv.interpreter().clone();
-    let tags = venv.interpreter().tags()?;
-    let markers = venv.interpreter().markers();
-
-    // Initialize the registry client.
-    // TODO(zanieb): Support client options e.g. offline, tls, etc.
-    let client = RegistryClientBuilder::new(cache.clone())
-        .markers(markers)
-        .platform(venv.interpreter().platform())
-        .build();
-
-    // TODO(charlie): Respect project configuration.
-    let build_isolation = BuildIsolation::default();
-    let config_settings = ConfigSettings::default();
-    let flat_index = FlatIndex::default();
-    let hasher = HashStrategy::default();
-    let in_flight = InFlight::default();
-    let index = InMemoryIndex::default();
-    let index_locations = IndexLocations::default();
-    let link_mode = LinkMode::default();
-    let no_binary = NoBinary::default();
-    let no_build = NoBuild::default();
-    let setup_py = SetupPyStrategy::default();
-    let concurrency = Concurrency::default();
-
-    // Create a build dispatch.
-    let build_dispatch = BuildDispatch::new(
-        &client,
-        cache,
-        &interpreter,
-        &index_locations,
-        &flat_index,
-        &index,
-        &in_flight,
-        setup_py,
-        &config_settings,
-        build_isolation,
-        link_mode,
-        &no_build,
-        &no_binary,
-        concurrency,
-    );
-
-    let options = OptionsBuilder::new()
-        // TODO(zanieb): Support resolver options
-        // .resolution_mode(resolution_mode)
-        // .prerelease_mode(prerelease_mode)
-        // .dependency_mode(dependency_mode)
-        // .exclude_newer(exclude_newer)
-        .build();
-
-    // Resolve the requirements.
-    let resolution = match project::resolve(
-        spec,
-        site_packages.clone(),
-        &hasher,
-        &interpreter,
-        tags,
-        markers,
-        &client,
-        &flat_index,
-        &index,
-        &build_dispatch,
-        options,
-        printer,
-        concurrency,
-    )
-    .await
-    {
-        Ok(resolution) => Resolution::from(resolution),
-        Err(err) => return Err(err.into()),
-    };
-
-    // Re-initialize the in-flight map.
-    let in_flight = InFlight::default();
-
-    // Sync the environment.
-    project::install(
-        &resolution,
-        site_packages,
-        &no_binary,
-        link_mode,
-        &index_locations,
-        &hasher,
-        tags,
-        &client,
-        &in_flight,
-        &build_dispatch,
-        cache,
-        &venv,
-        printer,
-        concurrency,
-    )
-    .await?;
-
-    Ok(venv)
-}