Use .with* for MiddlewareStack

Author: konstin

crates/uv-client/src/base_client.rs

@@ -22,39 +22,38 @@ use crate::middleware::OfflineMiddleware;
 use crate::Connectivity;
 
 /// Newtype to implement [`Debug`] on [`Middleware`].
-#[derive(Clone)]
+#[derive(Clone, Default)]
 pub struct MiddlewareStack(Vec<Arc<dyn Middleware>>);
 
 impl MiddlewareStack {
-    /// Use a custom middleware stack, skipping the default retry and auth layers.
+    /// Add an arbitrary middleware layer.
     // This function exists for rustlib users.
-    pub fn custom(middleware_stack: Vec<Arc<dyn Middleware>>) -> Self {
-        Self(middleware_stack)
+    pub fn with(mut self, middleware: impl Into<Arc<dyn Middleware>>) -> Self {
+        self.0.push(middleware.into());
+        self
     }
 
-    /// Initialize the middleware stack, using an [`ExponentialBackoff`] with the given number of
-    /// retries and an [`AuthMiddleware`] layer with the given keyring provider.
-    pub fn new(retries: u32, keyring: KeyringProviderType) -> Self {
-        let retry_policy = ExponentialBackoff::builder().build_with_max_retries(retries);
-        let retry_strategy = RetryTransientMiddleware::new_with_policy(retry_policy);
-
-        let auth_middleware = AuthMiddleware::new().with_keyring(keyring.to_provider());
-
-        Self(vec![Arc::new(retry_strategy), Arc::new(auth_middleware)])
+    /// Add an [`ExponentialBackoff`] layer with the given number of retries.
+    pub fn with_retries(mut self, retries: u32) -> Self {
+        if retries > 0 {
+            let retry_policy = ExponentialBackoff::builder().build_with_max_retries(retries);
+            let retry_strategy = RetryTransientMiddleware::new_with_policy(retry_policy);
+            self.0.push(Arc::new(retry_strategy))
+        }
+        self
     }
-}
 
-impl Default for MiddlewareStack {
-    /// Initialize the default middleware stack, consisting of an [`ExponentialBackoff`] retry
-    /// strategy and an [`AuthMiddleware`] layer not using the keyring.
-    fn default() -> Self {
-        Self::new(3, KeyringProviderType::default())
+    /// Add an [`AuthMiddleware`] layer with the given keyring provider.
+    pub fn with_auth(mut self, keyring: KeyringProviderType) -> Self {
+        let auth_middleware = AuthMiddleware::new().with_keyring(keyring.to_provider());
+        self.0.push(Arc::new(auth_middleware));
+        self
     }
 }
 
 impl Debug for MiddlewareStack {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        f.debug_struct("CustomMiddleware").finish()
+        f.debug_struct("MiddlewareStack").finish()
     }
 }
 

crates/uv/src/commands/pip/compile.rs

@@ -116,7 +116,11 @@ pub(crate) async fn pip_compile(
     let client_builder = BaseClientBuilder::new()
         .connectivity(connectivity)
         .native_tls(native_tls)
-        .middleware_stack(MiddlewareStack::new(3, keyring_provider));
+        .middleware_stack(
+            MiddlewareStack::default()
+                .with_retries(3)
+                .with_auth(keyring_provider),
+        );
 
     // Read all requirements from the provided sources.
     let RequirementsSpecification {
@@ -287,7 +291,11 @@ pub(crate) async fn pip_compile(
         .connectivity(connectivity)
         .index_urls(index_locations.index_urls())
         .index_strategy(index_strategy)
-        .middleware_stack(MiddlewareStack::new(3, keyring_provider))
+        .middleware_stack(
+            MiddlewareStack::default()
+                .with_retries(3)
+                .with_auth(keyring_provider),
+        )
         .markers(&markers)
         .platform(interpreter.platform())
         .build();

crates/uv/src/commands/pip/install.rs

@@ -85,7 +85,11 @@ pub(crate) async fn pip_install(
     let client_builder = BaseClientBuilder::new()
         .connectivity(connectivity)
         .native_tls(native_tls)
-        .middleware_stack(MiddlewareStack::new(3, keyring_provider));
+        .middleware_stack(
+            MiddlewareStack::default()
+                .with_retries(3)
+                .with_auth(keyring_provider),
+        );
 
     // Read all requirements from the provided sources.
     let RequirementsSpecification {
@@ -287,7 +291,11 @@ pub(crate) async fn pip_install(
         .connectivity(connectivity)
         .index_urls(index_locations.index_urls())
         .index_strategy(index_strategy)
-        .middleware_stack(MiddlewareStack::new(3, keyring_provider))
+        .middleware_stack(
+            MiddlewareStack::default()
+                .with_retries(3)
+                .with_auth(keyring_provider),
+        )
         .markers(&markers)
         .platform(interpreter.platform())
         .build();

crates/uv/src/commands/pip/sync.rs

@@ -72,7 +72,11 @@ pub(crate) async fn pip_sync(
     let client_builder = BaseClientBuilder::new()
         .connectivity(connectivity)
         .native_tls(native_tls)
-        .middleware_stack(MiddlewareStack::new(3, keyring_provider));
+        .middleware_stack(
+            MiddlewareStack::default()
+                .with_retries(3)
+                .with_auth(keyring_provider),
+        );
 
     // Initialize a few defaults.
     let overrides = &[];
@@ -225,7 +229,11 @@ pub(crate) async fn pip_sync(
         .connectivity(connectivity)
         .index_urls(index_locations.index_urls())
         .index_strategy(index_strategy)
-        .middleware_stack(MiddlewareStack::new(3, keyring_provider))
+        .middleware_stack(
+            MiddlewareStack::default()
+                .with_retries(3)
+                .with_auth(keyring_provider),
+        )
         .markers(&markers)
         .platform(interpreter.platform())
         .build();

crates/uv/src/commands/pip/uninstall.rs

@@ -38,7 +38,11 @@ pub(crate) async fn pip_uninstall(
     let client_builder = BaseClientBuilder::new()
         .connectivity(connectivity)
         .native_tls(native_tls)
-        .middleware_stack(MiddlewareStack::new(3, keyring_provider));
+        .middleware_stack(
+            MiddlewareStack::default()
+                .with_retries(3)
+                .with_auth(keyring_provider),
+        );
 
     // Read all requirements from the provided sources.
     let spec = RequirementsSpecification::from_simple_sources(sources, &client_builder).await?;

crates/uv/src/commands/venv.rs

@@ -174,7 +174,11 @@ async fn venv_impl(
             .native_tls(native_tls)
             .index_urls(index_locations.index_urls())
             .index_strategy(index_strategy)
-            .middleware_stack(MiddlewareStack::new(3, keyring_provider))
+            .middleware_stack(
+                MiddlewareStack::default()
+                    .with_retries(3)
+                    .with_auth(keyring_provider),
+            )
             .connectivity(connectivity)
             .markers(interpreter.markers())
             .platform(interpreter.platform())