From 35afcfd053ac3d4a96b90d512143fb2c44b0c959 Mon Sep 17 00:00:00 2001 From: Danny <33044223+danielenricocahall@users.noreply.github.com> Date: Thu, 4 Jul 2024 11:53:05 -0400 Subject: [PATCH] Enable Registry Client Builder to be created from Base Client Builder (#4729) ## Summary Addresses https://github.com/astral-sh/uv/issues/4330, to reduce duplication in the client creation logic. ## Test Plan https://github.com/astral-sh/uv/pull/4729#issuecomment-2204681655 --- crates/uv-client/src/registry_client.rs | 62 +++++++++---------------- crates/uv/src/commands/pip/compile.rs | 7 ++- crates/uv/src/commands/pip/install.rs | 6 +-- crates/uv/src/commands/pip/sync.rs | 6 +-- crates/uv/src/commands/project/add.rs | 5 +- crates/uv/src/commands/venv.rs | 7 +-- 6 files changed, 34 insertions(+), 59 deletions(-) diff --git a/crates/uv-client/src/registry_client.rs b/crates/uv-client/src/registry_client.rs index ef20831ef..29cd70eb0 100644 --- a/crates/uv-client/src/registry_client.rs +++ b/crates/uv-client/src/registry_client.rs @@ -37,14 +37,8 @@ use crate::{CachedClient, CachedClientError, Error, ErrorKind}; pub struct RegistryClientBuilder<'a> { index_urls: IndexUrls, index_strategy: IndexStrategy, - keyring: KeyringProviderType, - native_tls: bool, - retries: u32, - connectivity: Connectivity, cache: Cache, - client: Option, - markers: Option<&'a MarkerEnvironment>, - platform: Option<&'a Platform>, + base_client_builder: BaseClientBuilder<'a>, } impl RegistryClientBuilder<'_> { @@ -52,14 +46,8 @@ impl RegistryClientBuilder<'_> { Self { index_urls: IndexUrls::default(), index_strategy: IndexStrategy::default(), - keyring: KeyringProviderType::default(), - native_tls: false, cache, - connectivity: Connectivity::Online, - retries: 3, - client: None, - markers: None, - platform: None, + base_client_builder: BaseClientBuilder::new(), } } } @@ -79,25 +67,25 @@ impl<'a> RegistryClientBuilder<'a> { #[must_use] pub fn keyring(mut self, keyring_type: KeyringProviderType) -> Self { - self.keyring = keyring_type; + self.base_client_builder = self.base_client_builder.keyring(keyring_type); self } #[must_use] pub fn connectivity(mut self, connectivity: Connectivity) -> Self { - self.connectivity = connectivity; + self.base_client_builder = self.base_client_builder.connectivity(connectivity); self } #[must_use] pub fn retries(mut self, retries: u32) -> Self { - self.retries = retries; + self.base_client_builder = self.base_client_builder.retries(retries); self } #[must_use] pub fn native_tls(mut self, native_tls: bool) -> Self { - self.native_tls = native_tls; + self.base_client_builder = self.base_client_builder.native_tls(native_tls); self } @@ -109,44 +97,27 @@ impl<'a> RegistryClientBuilder<'a> { #[must_use] pub fn client(mut self, client: Client) -> Self { - self.client = Some(client); + self.base_client_builder = self.base_client_builder.client(client); self } #[must_use] pub fn markers(mut self, markers: &'a MarkerEnvironment) -> Self { - self.markers = Some(markers); + self.base_client_builder = self.base_client_builder.markers(markers); self } #[must_use] pub fn platform(mut self, platform: &'a Platform) -> Self { - self.platform = Some(platform); + self.base_client_builder = self.base_client_builder.platform(platform); self } pub fn build(self) -> RegistryClient { // Build a base client - let mut builder = BaseClientBuilder::new(); + let builder = self.base_client_builder; - if let Some(client) = self.client { - builder = builder.client(client); - } - - if let Some(markers) = self.markers { - builder = builder.markers(markers); - } - - if let Some(platform) = self.platform { - builder = builder.platform(platform); - } - - let client = builder - .retries(self.retries) - .connectivity(self.connectivity) - .native_tls(self.native_tls) - .keyring(self.keyring) - .build(); + let client = builder.build(); let timeout = client.timeout(); let connectivity = client.connectivity(); @@ -165,6 +136,17 @@ impl<'a> RegistryClientBuilder<'a> { } } +impl<'a> From> for RegistryClientBuilder<'a> { + fn from(value: BaseClientBuilder<'a>) -> Self { + Self { + index_urls: IndexUrls::default(), + index_strategy: IndexStrategy::default(), + cache: Cache::temp().unwrap(), + base_client_builder: value, + } + } +} + /// A client for fetching packages from a `PyPI`-compatible index. #[derive(Debug, Clone)] pub struct RegistryClient { diff --git a/crates/uv/src/commands/pip/compile.rs b/crates/uv/src/commands/pip/compile.rs index 7439e69c1..abaa60459 100644 --- a/crates/uv/src/commands/pip/compile.rs +++ b/crates/uv/src/commands/pip/compile.rs @@ -257,12 +257,11 @@ pub(crate) async fn pip_compile( } // Initialize the registry client. - let client = RegistryClientBuilder::new(cache.clone()) - .native_tls(native_tls) - .connectivity(connectivity) + + let client = RegistryClientBuilder::from(client_builder) + .cache(cache.clone()) .index_urls(index_locations.index_urls()) .index_strategy(index_strategy) - .keyring(keyring_provider) .markers(interpreter.markers()) .platform(interpreter.platform()) .build(); diff --git a/crates/uv/src/commands/pip/install.rs b/crates/uv/src/commands/pip/install.rs index e1d2deaaa..4c875eda2 100644 --- a/crates/uv/src/commands/pip/install.rs +++ b/crates/uv/src/commands/pip/install.rs @@ -256,12 +256,10 @@ pub(crate) async fn pip_install( } // Initialize the registry client. - let client = RegistryClientBuilder::new(cache.clone()) - .native_tls(native_tls) - .connectivity(connectivity) + let client = RegistryClientBuilder::from(client_builder) + .cache(cache.clone()) .index_urls(index_locations.index_urls()) .index_strategy(index_strategy) - .keyring(keyring_provider) .markers(&markers) .platform(interpreter.platform()) .build(); diff --git a/crates/uv/src/commands/pip/sync.rs b/crates/uv/src/commands/pip/sync.rs index 0c04490e5..a1970f4f5 100644 --- a/crates/uv/src/commands/pip/sync.rs +++ b/crates/uv/src/commands/pip/sync.rs @@ -203,12 +203,10 @@ pub(crate) async fn pip_sync( } // Initialize the registry client. - let client = RegistryClientBuilder::new(cache.clone()) - .native_tls(native_tls) - .connectivity(connectivity) + let client = RegistryClientBuilder::from(client_builder) + .cache(cache.clone()) .index_urls(index_locations.index_urls()) .index_strategy(index_strategy) - .keyring(keyring_provider) .markers(&markers) .platform(interpreter.platform()) .build(); diff --git a/crates/uv/src/commands/project/add.rs b/crates/uv/src/commands/project/add.rs index 5700c97ce..cdaeed4e5 100644 --- a/crates/uv/src/commands/project/add.rs +++ b/crates/uv/src/commands/project/add.rs @@ -94,12 +94,9 @@ pub(crate) async fn add( resolution_environment(python_version, python_platform, venv.interpreter())?; // Initialize the registry client. - let client = RegistryClientBuilder::new(cache.clone()) - .native_tls(native_tls) - .connectivity(connectivity) + let client = RegistryClientBuilder::from(client_builder) .index_urls(settings.index_locations.index_urls()) .index_strategy(settings.index_strategy) - .keyring(settings.keyring_provider) .markers(&markers) .platform(venv.interpreter().platform()) .build(); diff --git a/crates/uv/src/commands/venv.rs b/crates/uv/src/commands/venv.rs index 055d4134a..08369288f 100644 --- a/crates/uv/src/commands/venv.rs +++ b/crates/uv/src/commands/venv.rs @@ -129,6 +129,8 @@ async fn venv_impl( .connectivity(connectivity) .native_tls(native_tls); + let client_builder_clone = client_builder.clone(); + let mut interpreter_request = python_request.map(PythonRequest::parse); if preview.is_enabled() && interpreter_request.is_none() { interpreter_request = request_from_version_file().await.into_diagnostic()?; @@ -184,12 +186,11 @@ async fn venv_impl( let interpreter = venv.interpreter(); // Instantiate a client. - let client = RegistryClientBuilder::new(cache.clone()) - .native_tls(native_tls) + let client = RegistryClientBuilder::from(client_builder_clone) + .cache(cache.clone()) .index_urls(index_locations.index_urls()) .index_strategy(index_strategy) .keyring(keyring_provider) - .connectivity(connectivity) .markers(interpreter.markers()) .platform(interpreter.platform()) .build();