Add context manager support to clients (#5693)

* minimum context

Signed-off-by: redoubts <Redoubts@users.noreply.github.com>

* tests

Signed-off-by: redoubts <Redoubts@users.noreply.github.com>

* test_api_consistency

Signed-off-by: redoubts <Redoubts@users.noreply.github.com>

* xShinnRyuu rc

Signed-off-by: redoubts <Redoubts@users.noreply.github.com>

---------

Signed-off-by: redoubts <Redoubts@users.noreply.github.com>
Co-authored-by: Thomas Zhou <54688146+xShinnRyuu@users.noreply.github.com>

Author: Redoubts

python/glide-async/python/glide/glide_client.py

@@ -2,6 +2,7 @@
 
 import sys
 import threading
+from types import TracebackType
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -288,6 +289,12 @@ async def _create_uds_connection(self) -> None:
             raise ClosingError("Failed to create UDS connection") from e
 
     async def close(self, err_message: Optional[str] = None) -> None:
+        """
+        Forwards to `aclose`, the more common method for async resources.
+        """
+        await self.aclose(err_message)
+
+    async def aclose(self, err_message: Optional[str] = None) -> None:
         """
         Terminate the client by closing all associated resources, including the socket and any active futures.
         All open futures will be closed with an exception.
@@ -313,6 +320,17 @@ async def close(self, err_message: Optional[str] = None) -> None:
 
             await self._stream.aclose()
 
+    async def __aenter__(self) -> Self:
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: Optional[type[BaseException]],
+        exc: Optional[BaseException],
+        tb: Optional[TracebackType],
+    ) -> None:
+        await self.aclose()
+
     def _get_future(self, callback_idx: int) -> "TFuture":
         response_future: "TFuture" = _get_new_future_instance()
         self._available_futures.update({callback_idx: response_future})

python/glide-sync/glide_sync/glide_client.py

@@ -3,6 +3,7 @@
 import os
 import sys
 import threading
+from types import TracebackType
 from typing import Any, List, Optional, Tuple, Union
 
 from glide_shared._fast_response import parse_response as _fast_parse_response
@@ -911,7 +912,7 @@ def _parse_pubsub_state(self, result, is_cluster):
             actual_subscriptions=actual_subscriptions,
         )
 
-    def close(self):
+    def close(self) -> None:
         if not self._is_closed:
             self._is_closed = True
             with self._pubsub_condition:
@@ -920,6 +921,17 @@ def close(self):
             self._core_client = self._ffi.NULL
             self._pubsub_callback_ref = None
 
+    def __enter__(self) -> Self:
+        return self
+
+    def __exit__(
+        self,
+        exc_type: Optional[type[BaseException]],
+        exc: Optional[BaseException],
+        tb: Optional[TracebackType],
+    ) -> None:
+        self.close()
+
 
 class GlideClusterClient(BaseClient, ClusterCommands):
     """

python/tests/async_tests/test_async_client.py

@@ -135,6 +135,16 @@ async def test_register_client_name_and_version(self, glide_client: TGlideClient
         assert "lib-name=GlidePy" in info_str
         assert "lib-ver=unknown" in info_str
 
+    @pytest.mark.parametrize("cluster_mode", [True, False])
+    @pytest.mark.parametrize("protocol", [ProtocolVersion.RESP2, ProtocolVersion.RESP3])
+    async def test_context_manager(self, request, cluster_mode, protocol):
+        async with await create_client(
+            request, cluster_mode=cluster_mode, protocol=protocol, request_timeout=5000
+        ) as client:
+            assert not client._is_closed
+
+        assert client._is_closed
+
     @pytest.mark.parametrize("cluster_mode", [True, False])
     @pytest.mark.parametrize("protocol", [ProtocolVersion.RESP2, ProtocolVersion.RESP3])
     async def test_send_and_receive_large_values(self, request, cluster_mode, protocol):

python/tests/sync_tests/test_sync_client.py

@@ -139,6 +139,16 @@ def test_sync_register_client_name_and_version(
         assert "lib-name=GlidePySync" in info_str
         assert "lib-ver=unknown" in info_str
 
+    @pytest.mark.parametrize("cluster_mode", [True, False])
+    @pytest.mark.parametrize("protocol", [ProtocolVersion.RESP2, ProtocolVersion.RESP3])
+    def test_context_manager(self, request, cluster_mode, protocol):
+        with create_sync_client(
+            request, cluster_mode=cluster_mode, protocol=protocol, request_timeout=5000
+        ) as client:
+            assert not client._is_closed
+
+        assert client._is_closed
+
     @pytest.mark.parametrize("cluster_mode", [True, False])
     @pytest.mark.parametrize("protocol", [ProtocolVersion.RESP2, ProtocolVersion.RESP3])
     def test_sync_send_and_receive_large_values(self, request, cluster_mode, protocol):

python/tests/test_api_consistency.py

@@ -37,6 +37,7 @@
         "create_leaked_value",
         "start_socket_listener_external",
         "value_from_pointer",
+        "aclose",
     ],
     "sync_only": [],
 }