feat(client): support with and async with on the clients

qdrant_client/async_qdrant_client.py

@@ -10,6 +10,7 @@
 # ******  WARNING: THIS FILE IS AUTOGENERATED  ******
 
 import warnings
+from types import TracebackType
 from typing import Any, Awaitable, Callable, Iterable, Mapping, Sequence
 import numpy as np
 from qdrant_client import grpc as grpc
@@ -155,6 +156,17 @@ def __init__(
             is_local_mode=isinstance(self._client, AsyncQdrantLocal),
         )
 
+    async def __aenter__(self) -> "AsyncQdrantClient":
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_value: BaseException | None,
+        traceback: TracebackType | None,
+    ) -> None:
+        await self.close()
+
     async def close(self, grpc_grace: float | None = None, **kwargs: Any) -> None:
         """Closes the connection to Qdrant
 

qdrant_client/qdrant_client.py

@@ -1,4 +1,5 @@
 import warnings
+from types import TracebackType
 from typing import (
     Any,
     Awaitable,
@@ -168,6 +169,17 @@ def __init__(
     def __del__(self) -> None:
         self.close()
 
+    def __enter__(self) -> "QdrantClient":
+        return self
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_value: BaseException | None,
+        traceback: TracebackType | None,
+    ) -> None:
+        self.close()
+
     def close(self, grpc_grace: float | None = None, **kwargs: Any) -> None:
         """Closes the connection to Qdrant
 

tests/test_context_manager.py

@@ -0,0 +1,67 @@
+import pytest
+
+from qdrant_client import AsyncQdrantClient, QdrantClient, models
+
+
+def test_sync_context_manager_closes_on_exit():
+    with QdrantClient(":memory:") as client:
+        assert isinstance(client, QdrantClient)
+        client.create_collection(
+            collection_name="ctx",
+            vectors_config=models.VectorParams(size=2, distance=models.Distance.DOT),
+        )
+        assert client._client.closed is False
+
+    assert client._client.closed is True
+
+
+def test_sync_context_manager_closes_when_body_raises():
+    client = QdrantClient(":memory:")
+
+    with pytest.raises(RuntimeError, match="boom"):
+        with client:
+            raise RuntimeError("boom")
+
+    assert client._client.closed is True
+
+
+def test_sync_context_manager_exit_is_idempotent():
+    with QdrantClient(":memory:") as client:
+        pass
+
+    # A second close() must not raise: QdrantLocal/QdrantRemote.close() are idempotent.
+    client.close()
+    assert client._client.closed is True
+
+
+@pytest.mark.asyncio
+async def test_async_context_manager_closes_on_exit():
+    async with AsyncQdrantClient(":memory:") as client:
+        assert isinstance(client, AsyncQdrantClient)
+        await client.create_collection(
+            collection_name="ctx",
+            vectors_config=models.VectorParams(size=2, distance=models.Distance.DOT),
+        )
+        assert client._client.closed is False
+
+    assert client._client.closed is True
+
+
+@pytest.mark.asyncio
+async def test_async_context_manager_closes_when_body_raises():
+    client = AsyncQdrantClient(":memory:")
+
+    with pytest.raises(RuntimeError, match="boom"):
+        async with client:
+            raise RuntimeError("boom")
+
+    assert client._client.closed is True
+
+
+@pytest.mark.asyncio
+async def test_async_context_manager_exit_is_idempotent():
+    async with AsyncQdrantClient(":memory:") as client:
+        pass
+
+    await client.close()
+    assert client._client.closed is True

tools/async_client_generator/client_generator.py

@@ -21,6 +21,7 @@ def __init__(
         class_replace_map: dict[str, str] | None = None,
         import_replace_map: dict[str, str] | None = None,
         exclude_methods: list[str] | None = None,
+        rename_methods: dict[str, str] | None = None,
     ):
         super().__init__()
         self._async_methods: list[str] | None = None
@@ -33,6 +34,7 @@ def __init__(
                 class_replace_map=class_replace_map,
                 exclude_methods=exclude_methods,
                 async_methods=self.async_methods,
+                rename_methods=rename_methods,
             )
         )
         self.transformers.append(ClassDefTransformer(class_replace_map=class_replace_map))
@@ -94,6 +96,10 @@ def get_async_methods(class_obj: type) -> list[str]:
             "__del__",
             "migrate",
         ],
+        rename_methods={
+            "__enter__": "__aenter__",
+            "__exit__": "__aexit__",
+        },
     )
 
     modified_code = generator.generate(code)

tools/async_client_generator/transformers/client/function_def_transformer.py

@@ -10,17 +10,34 @@ def __init__(
         class_replace_map: dict[str, str] | None = None,
         exclude_methods: list[str] | None = None,
         async_methods: list[str] | None = None,
+        rename_methods: dict[str, str] | None = None,
     ):
         super().__init__(keep_sync)
         self.class_replace_map = class_replace_map if class_replace_map is not None else {}
         self.exclude_methods = exclude_methods if exclude_methods is not None else []
         self.async_methods = async_methods if async_methods is not None else []
+        self.rename_methods = rename_methods if rename_methods is not None else {}
 
     def _keep_sync(self, name: str) -> bool:
+        # Renamed methods (e.g. __enter__ -> __aenter__) have no counterpart on
+        # AsyncQdrantBase, so the async_methods lookup would keep them sync.
+        if name in self.rename_methods.values():
+            return False
         return name in self.keep_sync or name not in self.async_methods
 
     def visit_FunctionDef(self, sync_node: ast.FunctionDef) -> ast.AST | None:
         if sync_node.name in self.exclude_methods:
             return None
 
+        if sync_node.name in self.rename_methods:
+            sync_node.name = self.rename_methods[sync_node.name]
+            # A string return annotation ("QdrantClient") is a Constant, so the
+            # name-based transformers leave it alone.
+            if isinstance(sync_node.returns, ast.Constant) and isinstance(
+                sync_node.returns.value, str
+            ):
+                sync_node.returns.value = self.class_replace_map.get(
+                    sync_node.returns.value, sync_node.returns.value
+                )
+
         return super().visit_FunctionDef(sync_node)