From 368276d7d132694c3d9a81023e53fa7e029e14de Mon Sep 17 00:00:00 2001 From: Charlie Marsh Date: Mon, 1 Jul 2024 20:27:01 -0400 Subject: [PATCH] Shared Git and in-memory index across operations (#4715) ## Summary I ended up needing this for https://github.com/astral-sh/uv/issues/4664 but I think it's a good change more broadly. We should be able to share this cached information across operations within a given invocation. --- crates/uv/src/commands/project/add.rs | 6 ++++++ crates/uv/src/commands/project/lock.rs | 23 ++++++++++++----------- crates/uv/src/commands/project/mod.rs | 22 +++++++++++++++------- crates/uv/src/commands/project/remove.rs | 6 ++++++ crates/uv/src/commands/project/run.rs | 8 ++++++++ crates/uv/src/commands/project/sync.rs | 13 ++++++------- crates/uv/src/commands/tool/install.rs | 3 ++- crates/uv/src/commands/tool/run.rs | 3 ++- 8 files changed, 57 insertions(+), 27 deletions(-) diff --git a/crates/uv/src/commands/project/add.rs b/crates/uv/src/commands/project/add.rs index 09da51574..be2cb3ed8 100644 --- a/crates/uv/src/commands/project/add.rs +++ b/crates/uv/src/commands/project/add.rs @@ -18,6 +18,7 @@ use uv_warnings::warn_user_once; use crate::commands::pip::operations::Modifications; use crate::commands::pip::resolution_environment; +use crate::commands::project::SharedState; use crate::commands::reporters::ResolverReporter; use crate::commands::{project, ExitStatus}; use crate::printer::Printer; @@ -205,11 +206,15 @@ pub(crate) async fn add( pyproject.to_string(), )?; + // Initialize any shared state. + let state = SharedState::default(); + // Lock and sync the environment. let lock = project::lock::do_lock( project.workspace(), venv.interpreter(), settings.as_ref().into(), + &state, preview, connectivity, concurrency, @@ -232,6 +237,7 @@ pub(crate) async fn add( dev, Modifications::Sufficient, settings.as_ref().into(), + &state, preview, connectivity, concurrency, diff --git a/crates/uv/src/commands/project/lock.rs b/crates/uv/src/commands/project/lock.rs index 2d58ba071..1ebbd20a9 100644 --- a/crates/uv/src/commands/project/lock.rs +++ b/crates/uv/src/commands/project/lock.rs @@ -8,16 +8,14 @@ use uv_client::{Connectivity, FlatIndexClient, RegistryClientBuilder}; use uv_configuration::{Concurrency, ExtrasSpecification, PreviewMode, Reinstall, SetupPyStrategy}; use uv_dispatch::BuildDispatch; use uv_distribution::{Workspace, DEV_DEPENDENCIES}; -use uv_git::GitResolver; +use uv_git::ResolvedRepositoryReference; use uv_requirements::upgrade::{read_lockfile, LockedRequirements}; -use uv_resolver::{ - FlatIndex, InMemoryIndex, Lock, OptionsBuilder, PythonRequirement, RequiresPython, -}; +use uv_resolver::{FlatIndex, Lock, OptionsBuilder, PythonRequirement, RequiresPython}; use uv_toolchain::{Interpreter, ToolchainPreference, ToolchainRequest}; use uv_types::{BuildIsolation, EmptyInstalledPackages, HashStrategy, InFlight}; use uv_warnings::{warn_user, warn_user_once}; -use crate::commands::project::{find_requires_python, FoundInterpreter, ProjectError}; +use crate::commands::project::{find_requires_python, FoundInterpreter, ProjectError, SharedState}; use crate::commands::{pip, ExitStatus}; use crate::printer::Printer; use crate::settings::{ResolverSettings, ResolverSettingsRef}; @@ -59,6 +57,7 @@ pub(crate) async fn lock( &workspace, &interpreter, settings.as_ref(), + &SharedState::default(), preview, connectivity, concurrency, @@ -86,6 +85,7 @@ pub(super) async fn do_lock( workspace: &Workspace, interpreter: &Interpreter, settings: ResolverSettingsRef<'_>, + state: &SharedState, preview: PreviewMode, connectivity: Connectivity, concurrency: Concurrency, @@ -163,7 +163,6 @@ pub(super) async fn do_lock( // Initialize any shared state. let in_flight = InFlight::default(); - let index = InMemoryIndex::default(); // TODO(charlie): These are all default values. We should consider whether we want to make them // optional on the downstream APIs. @@ -181,8 +180,10 @@ pub(super) async fn do_lock( // If an existing lockfile exists, build up a set of preferences. let LockedRequirements { preferences, git } = read_lockfile(workspace, upgrade).await?; - // Create the Git resolver. - let git = GitResolver::from_refs(git); + // Populate the Git resolver. + for ResolvedRepositoryReference { reference, sha } in git { + state.git.insert(reference, sha); + } // Create a build dispatch. let build_dispatch = BuildDispatch::new( @@ -191,8 +192,8 @@ pub(super) async fn do_lock( interpreter, index_locations, &flat_index, - &index, - &git, + &state.index, + &state.git, &in_flight, index_strategy, setup_py, @@ -224,7 +225,7 @@ pub(super) async fn do_lock( python_requirement, &client, &flat_index, - &index, + &state.index, &build_dispatch, concurrency, options, diff --git a/crates/uv/src/commands/project/mod.rs b/crates/uv/src/commands/project/mod.rs index 8bc7da6c5..466a4936d 100644 --- a/crates/uv/src/commands/project/mod.rs +++ b/crates/uv/src/commands/project/mod.rs @@ -273,11 +273,21 @@ pub(crate) async fn get_or_init_environment( } } +/// Shared state used during resolution and installation. +#[derive(Default)] +pub(crate) struct SharedState { + /// The resolved Git references. + git: GitResolver, + /// The fetched package versions and metadata. + index: InMemoryIndex, +} + /// Update a [`PythonEnvironment`] to satisfy a set of [`RequirementsSource`]s. pub(crate) async fn update_environment( venv: PythonEnvironment, spec: RequirementsSpecification, settings: &ResolverInstallerSettings, + state: &SharedState, preview: PreviewMode, connectivity: Connectivity, concurrency: Concurrency, @@ -350,9 +360,7 @@ pub(crate) async fn update_environment( .build(); // Initialize any shared state. - let git = GitResolver::default(); let in_flight = InFlight::default(); - let index = InMemoryIndex::default(); // TODO(charlie): These are all default values. We should consider whether we want to make them // optional on the downstream APIs. @@ -378,8 +386,8 @@ pub(crate) async fn update_environment( interpreter, index_locations, &flat_index, - &index, - &git, + &state.index, + &state.git, &in_flight, *index_strategy, setup_py, @@ -411,7 +419,7 @@ pub(crate) async fn update_environment( python_requirement, &client, &flat_index, - &index, + &state.index, &resolve_dispatch, concurrency, options, @@ -438,8 +446,8 @@ pub(crate) async fn update_environment( interpreter, index_locations, &flat_index, - &index, - &git, + &state.index, + &state.git, &in_flight, *index_strategy, setup_py, diff --git a/crates/uv/src/commands/project/remove.rs b/crates/uv/src/commands/project/remove.rs index c19d10ed2..49eb390b7 100644 --- a/crates/uv/src/commands/project/remove.rs +++ b/crates/uv/src/commands/project/remove.rs @@ -11,6 +11,7 @@ use uv_toolchain::{ToolchainPreference, ToolchainRequest}; use uv_warnings::{warn_user, warn_user_once}; use crate::commands::pip::operations::Modifications; +use crate::commands::project::SharedState; use crate::commands::{project, ExitStatus}; use crate::printer::Printer; use crate::settings::{InstallerSettings, ResolverSettings}; @@ -95,11 +96,15 @@ pub(crate) async fn remove( // Use the default settings. let settings = ResolverSettings::default(); + // Initialize any shared state. + let state = SharedState::default(); + // Lock and sync the environment. let lock = project::lock::do_lock( project.workspace(), venv.interpreter(), settings.as_ref(), + &state, preview, connectivity, concurrency, @@ -123,6 +128,7 @@ pub(crate) async fn remove( dev, Modifications::Exact, settings.as_ref(), + &state, preview, connectivity, concurrency, diff --git a/crates/uv/src/commands/project/run.rs b/crates/uv/src/commands/project/run.rs index 52ddfa338..d19da2656 100644 --- a/crates/uv/src/commands/project/run.rs +++ b/crates/uv/src/commands/project/run.rs @@ -21,6 +21,7 @@ use uv_toolchain::{ use uv_warnings::warn_user_once; use crate::commands::pip::operations::Modifications; +use crate::commands::project::SharedState; use crate::commands::{project, ExitStatus}; use crate::printer::Printer; use crate::settings::ResolverInstallerSettings; @@ -50,6 +51,9 @@ pub(crate) async fn run( // Parse the input command. let command = RunCommand::from(command); + // Initialize any shared state. + let state = SharedState::default(); + // Determine whether the command to execute is a PEP 723 script. let temp_dir; let script_interpreter = if let RunCommand::Python(target, _) = &command { @@ -106,6 +110,7 @@ pub(crate) async fn run( venv, spec, &settings, + &state, preview, connectivity, concurrency, @@ -177,6 +182,7 @@ pub(crate) async fn run( project.workspace(), venv.interpreter(), settings.as_ref().into(), + &state, preview, connectivity, concurrency, @@ -193,6 +199,7 @@ pub(crate) async fn run( dev, Modifications::Sufficient, settings.as_ref().into(), + &state, preview, connectivity, concurrency, @@ -289,6 +296,7 @@ pub(crate) async fn run( venv, spec, &settings, + &state, preview, connectivity, concurrency, diff --git a/crates/uv/src/commands/project/sync.rs b/crates/uv/src/commands/project/sync.rs index 61eec69d1..fc327f344 100644 --- a/crates/uv/src/commands/project/sync.rs +++ b/crates/uv/src/commands/project/sync.rs @@ -5,15 +5,14 @@ use uv_client::{Connectivity, FlatIndexClient, RegistryClientBuilder}; use uv_configuration::{Concurrency, ExtrasSpecification, PreviewMode, SetupPyStrategy}; use uv_dispatch::BuildDispatch; use uv_distribution::{VirtualProject, DEV_DEPENDENCIES}; -use uv_git::GitResolver; use uv_installer::SitePackages; -use uv_resolver::{FlatIndex, InMemoryIndex, Lock}; +use uv_resolver::{FlatIndex, Lock}; use uv_toolchain::{PythonEnvironment, ToolchainPreference, ToolchainRequest}; use uv_types::{BuildIsolation, HashStrategy, InFlight}; use uv_warnings::warn_user_once; use crate::commands::pip::operations::Modifications; -use crate::commands::project::ProjectError; +use crate::commands::project::{ProjectError, SharedState}; use crate::commands::{pip, project, ExitStatus}; use crate::printer::Printer; use crate::settings::{InstallerSettings, InstallerSettingsRef}; @@ -68,6 +67,7 @@ pub(crate) async fn sync( dev, modifications, settings.as_ref(), + &SharedState::default(), preview, connectivity, concurrency, @@ -89,6 +89,7 @@ pub(super) async fn do_sync( dev: bool, modifications: Modifications, settings: InstallerSettingsRef<'_>, + state: &SharedState, preview: PreviewMode, connectivity: Connectivity, concurrency: Concurrency, @@ -143,9 +144,7 @@ pub(super) async fn do_sync( .build(); // Initialize any shared state. - let git = GitResolver::default(); let in_flight = InFlight::default(); - let index = InMemoryIndex::default(); // TODO(charlie): These are all default values. We should consider whether we want to make them // optional on the downstream APIs. @@ -169,8 +168,8 @@ pub(super) async fn do_sync( venv.interpreter(), index_locations, &flat_index, - &index, - &git, + &state.index, + &state.git, &in_flight, index_strategy, setup_py, diff --git a/crates/uv/src/commands/tool/install.rs b/crates/uv/src/commands/tool/install.rs index 26f3d981b..52f197e42 100644 --- a/crates/uv/src/commands/tool/install.rs +++ b/crates/uv/src/commands/tool/install.rs @@ -21,7 +21,7 @@ use uv_tool::{entrypoint_paths, find_executable_directory, InstalledTools, Tool, use uv_toolchain::{EnvironmentPreference, Toolchain, ToolchainPreference, ToolchainRequest}; use uv_warnings::warn_user_once; -use crate::commands::project::update_environment; +use crate::commands::project::{update_environment, SharedState}; use crate::commands::ExitStatus; use crate::printer::Printer; use crate::settings::ResolverInstallerSettings; @@ -146,6 +146,7 @@ pub(crate) async fn install( environment, spec, &settings, + &SharedState::default(), preview, connectivity, concurrency, diff --git a/crates/uv/src/commands/tool/run.rs b/crates/uv/src/commands/tool/run.rs index 380c8b8e5..cc58b162b 100644 --- a/crates/uv/src/commands/tool/run.rs +++ b/crates/uv/src/commands/tool/run.rs @@ -19,7 +19,7 @@ use uv_toolchain::{ }; use uv_warnings::warn_user_once; -use crate::commands::project::update_environment; +use crate::commands::project::{update_environment, SharedState}; use crate::commands::ExitStatus; use crate::printer::Printer; use crate::settings::ResolverInstallerSettings; @@ -102,6 +102,7 @@ pub(crate) async fn run( venv, spec, &settings, + &SharedState::default(), preview, connectivity, concurrency,