feat(client): implement graceful shutdown for Client

Cargo.toml

@@ -28,7 +28,8 @@ abi3-py313 = ["pyo3/abi3-py313"]
 abi3-py314 = ["pyo3/abi3-py314"]
 
 [dependencies]
-tokio = { version = "1.49.0", features = ["sync"]}
+tokio = { version = "1.50.0", features = ["sync"]}
+tokio-util = "0.7.18"
 pyo3 = { version = "0.28.2", features = [
     "indexmap",
     "multiple-pymethods",

python/wreq/blocking.py

@@ -256,6 +256,28 @@ def __init__(
         """
         ...
 
+    def close(self) -> None:
+        r"""
+        Closes the client and any associated resources.
+
+        After calling this method, the client should not be used to make further requests.
+
+        Examples:
+
+        ```python
+        import asyncio
+        import wreq
+
+        client = wreq.blocking.Client()
+
+        response = client.get('https://httpbin.io/get')
+        print(response.text())
+
+        client.close()
+        ```
+        """
+        ...
+
     def request(
         self,
         method: Method,

python/wreq/cookie.py

@@ -9,7 +9,6 @@
 import datetime
 from enum import Enum, auto
 from typing import Sequence, final
-from warnings import deprecated
 
 __all__ = ["SameSite", "Cookie", "Jar"]
 

python/wreq/wreq.py

@@ -1136,6 +1136,31 @@ async def main():
         """
         ...
 
+    def close(self) -> None:
+        r"""
+        Closes the client and any associated resources.
+
+        After calling this method, the client should not be used to make further requests.
+
+        Examples:
+
+        ```python
+        import asyncio
+        import wreq
+
+        async def main():
+            client = wreq.Client()
+
+            response = await client.get('https://httpbin.io/get')
+            print(await response.text())
+
+            client.close()
+
+        asyncio.run(main())
+        ```
+        """
+        ...
+
     async def request(
         self,
         method: Method,

src/client.rs

@@ -15,6 +15,7 @@ use std::{
 
 use pyo3::{IntoPyObjectExt, coroutine::CancelHandle, prelude::*, pybacked::PyBackedStr};
 use req::{Request, WebSocketRequest};
+use tokio_util::sync::CancellationToken;
 use wreq::{Proxy, tls::CertStore};
 use wreq_util::EmulationOption;
 
@@ -234,6 +235,7 @@ impl FromPyObject<'_, '_> for Builder {
 #[pyclass(subclass, frozen, skip_from_py_object)]
 pub struct Client {
     inner: wreq::Client,
+    cancel: CancellationToken,
 
     /// Get the cookie jar of the client.
     #[pyo3(get)]
@@ -449,12 +451,22 @@ impl Client {
 
             builder
                 .build()
-                .map(|inner| Client { inner, cookie_jar })
+                .map(|inner| Client {
+                    inner,
+                    cancel: CancellationToken::new(),
+                    cookie_jar,
+                })
                 .map_err(Error::Library)
                 .map_err(Into::into)
         })
     }
 
+    /// Close the client, preventing any new requests.
+    #[inline]
+    pub fn close(&self) {
+        self.cancel.cancel();
+    }
+
     /// Make a GET request to the given URL.
     #[inline(always)]
     #[pyo3(signature = (url, **kwds))]
@@ -561,9 +573,10 @@ impl Client {
         url: PyBackedStr,
         kwds: Option<Request>,
     ) -> PyResult<Response> {
-        NoGIL::new(
+        NoGIL::new_with_token(
             execute_request(self.inner.clone(), method, url, kwds),
             cancel,
+            self.cancel.clone(),
         )
         .await
     }
@@ -577,9 +590,10 @@ impl Client {
         url: PyBackedStr,
         kwds: Option<WebSocketRequest>,
     ) -> PyResult<WebSocket> {
-        NoGIL::new(
+        NoGIL::new_with_token(
             execute_websocket_request(self.inner.clone(), url, kwds),
             cancel,
+            self.cancel.clone(),
         )
         .await
     }
@@ -594,7 +608,7 @@ impl Client {
 
     #[inline]
     async fn __aexit__(&self, _exc_type: Py<PyAny>, _exc_val: Py<PyAny>, _traceback: Py<PyAny>) {
-        // TODO: Implement connection closing logic if necessary.
+        self.cancel.cancel();
     }
 }
 
@@ -617,6 +631,12 @@ impl BlockingClient {
         self.0.cookie_jar.clone()
     }
 
+    /// Close the client, preventing any new requests.
+    #[inline]
+    pub fn close(&self) {
+        self.0.close();
+    }
+
     /// Make a GET request to the specified URL.
     #[inline(always)]
     #[pyo3(signature = (url, **kwds))]

src/client/nogil.rs

@@ -11,6 +11,7 @@ use pyo3::{
     prelude::*,
 };
 use tokio::task::JoinHandle;
+use tokio_util::sync::CancellationToken;
 
 pin_project! {
     /// A future that allows Python threads to run while it is being polled or executed.
@@ -38,6 +39,25 @@ where
             }
         }) }
     }
+
+    /// Create [`NoGIL`] from a future and a cancellation token
+    #[inline]
+    pub fn new_with_token<Fut>(
+        fut: Fut,
+        mut cancel: CancelHandle,
+        cancel_token: CancellationToken,
+    ) -> Self
+    where
+        Fut: Future<Output = PyResult<T>> + Send + 'static,
+    {
+        Self { handle:  pyo3_async_runtimes::tokio::get_runtime().spawn(async move {
+            tokio::select! {
+                result = fut => result,
+                _ = cancel.cancelled() => Err(CancelledError::new_err("Operation was cancelled")),
+                _ = cancel_token.cancelled() => Err(CancelledError::new_err("Operation was cancelled: client has been closed")),
+            }
+        }) }
+    }
 }
 
 impl<T> Future for NoGIL<T>

src/client/resp/http.rs

@@ -44,7 +44,7 @@ enum Body {
 }
 
 /// A blocking response from a request.
-#[pyclass(name = "Response", subclass, frozen, str)]
+#[pyclass(name = "Response", subclass, frozen, str, skip_from_py_object)]
 pub struct BlockingResponse(Response);
 
 // ===== impl Response =====
@@ -115,6 +115,7 @@ impl Response {
 
     /// Forcefully destroys the response [`Body`], preventing any further reads.
     fn destroy(&self) {
+        #[allow(clippy::option_map_unit_fn)]
         self.body
             .swap(None)
             .and_then(Arc::into_inner)