From e6b42adee232f9cad64a9dac644c5c23a3504a47 Mon Sep 17 00:00:00 2001
From: gngpp <gngppz@gmail.com>
Date: Thu, 5 Mar 2026 17:12:42 +0800
Subject: [PATCH] feat(client): implement context manager protocol

---
 python/rnet/__init__.pyi |  5 +++++
 python/rnet/blocking.py  |  7 +++++--
 src/client.rs            | 34 ++++++++++++++++++++++++++++++++++
 3 files changed, 44 insertions(+), 2 deletions(-)

diff --git a/python/rnet/__init__.pyi b/python/rnet/__init__.pyi
index 2b753044..9ed6e3c3 100644
--- a/python/rnet/__init__.pyi
+++ b/python/rnet/__init__.pyi
@@ -1354,6 +1354,11 @@ class Client:
         asyncio.run(main())
         ```
         """
+        
+    async def __aenter__(self) -> Any: ...
+    async def __aexit__(
+        self, _exc_type: Any, _exc_value: Any, _traceback: Any
+    ) -> Any: ...
 
 async def delete(
     url: str,
diff --git a/python/rnet/blocking.py b/python/rnet/blocking.py
index 837f06fd..d793e794 100644
--- a/python/rnet/blocking.py
+++ b/python/rnet/blocking.py
@@ -127,7 +127,7 @@ def close(self) -> None:
         in future versions.
         """
 
-    def __enter__(self) -> "Response": ...
+    def __enter__(self) -> Any: ...
     def __exit__(self, _exc_type: Any, _exc_value: Any, _traceback: Any) -> None: ...
     def __str__(self) -> str: ...
 
@@ -204,7 +204,7 @@ def close(
         * `reason` - An optional reason for closing.
         """
 
-    def __enter__(self) -> "WebSocket": ...
+    def __enter__(self) -> Any: ...
     def __exit__(self, _exc_type: Any, _exc_value: Any, _traceback: Any) -> None: ...
     def __str__(self) -> str: ...
 
@@ -463,6 +463,9 @@ def get(
         ```
         """
         ...
+        
+    def __enter__(self) -> Any: ...
+    def __exit__(self, _exc_type: Any, _exc_value: Any, _traceback: Any) -> None: ...
 
 
 def delete(
diff --git a/src/client.rs b/src/client.rs
index 0f1081ce..3a8e00f1 100644
--- a/src/client.rs
+++ b/src/client.rs
@@ -585,6 +585,21 @@ impl Client {
     }
 }
 
+#[pymethods]
+impl Client {
+    #[inline]
+    async fn __aenter__(slf: Py<Self>) -> PyResult<Py<Self>> {
+        Ok(slf)
+    }
+
+    #[inline]
+    async fn __aexit__(&self, _exc_type: Py<PyAny>, _exc_val: Py<PyAny>, _traceback: Py<PyAny>) {
+        // TODO: Implement connection closing logic if necessary.
+    }
+}
+
+// ===== impl BlockingClient =====
+
 #[pymethods]
 impl BlockingClient {
     /// Creates a new blocking Client instance.
@@ -729,3 +744,22 @@ impl BlockingClient {
         })
     }
 }
+
+#[pymethods]
+impl BlockingClient {
+    #[inline]
+    fn __enter__(slf: PyRef<Self>) -> PyRef<Self> {
+        slf
+    }
+
+    #[inline]
+    fn __exit__<'py>(
+        &self,
+        _py: Python<'py>,
+        _exc_type: &Bound<'py, PyAny>,
+        _exc_value: &Bound<'py, PyAny>,
+        _traceback: &Bound<'py, PyAny>,
+    ) {
+        // TODO: Implement connection closing logic if necessary.
+    }
+}