fix(client): Keep the public client's type in with WorkOSClient() as workos

src/workos/_base_client.py

@@ -6,6 +6,7 @@
 import math
 import os
 import platform
+import sys
 import time
 import uuid
 import random
@@ -14,6 +15,11 @@
 from typing import Any, Dict, List, Optional, Sequence, Type, Union, cast, overload
 from urllib.parse import quote, urlencode
 
+if sys.version_info >= (3, 11):
+    from typing import Self
+else:
+    from typing_extensions import Self
+
 from ._types import _format_datetime
 from ._errors import (
     APIError,
@@ -546,7 +552,7 @@ def close(self) -> None:
         if self._owns_http_client:
             self._backend.close()
 
-    def __enter__(self) -> "WorkOSClient":
+    def __enter__(self) -> Self:
         return self
 
     def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
@@ -795,7 +801,7 @@ async def close(self) -> None:
         if self._owns_http_client:
             await self._backend.close()
 
-    async def __aenter__(self) -> "AsyncWorkOSClient":
+    async def __aenter__(self) -> Self:
         return self
 
     async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:

tests/test_generated_client.py

@@ -64,6 +64,14 @@ def test_context_manager(self):
         with WorkOSClient(api_key="sk_test_123", client_id="client_test") as client:
             assert client._api_key == "sk_test_123"
 
+    def test_context_manager_yields_the_client_with_its_resources(self):
+        # pyright checks this file, so this also checks that `with` gives the
+        # public client's type, which has the resources and helpers.
+        with WorkOSClient(api_key="sk_test_123") as client:
+            assert client.organizations is not None
+            runtime = client.create_feature_flags_runtime_client()
+            runtime.close()
+
     def test_api_key_only_initializes(self):
         client = WorkOSClient(api_key="sk_test_123")
         assert client._api_key == "sk_test_123"
@@ -632,6 +640,14 @@ async def test_documented_import_surface_exposes_resources(self):
         assert client.widgets is not None
         await client.close()
 
+    async def test_context_manager_yields_the_client_with_its_resources(self):
+        # pyright checks this file, so this also checks that `async with` gives
+        # the public client's type, which has the resources and helpers.
+        async with AsyncWorkOSClient(api_key="sk_test_123") as client:
+            assert client.organizations is not None
+            runtime = client.create_feature_flags_runtime_client()
+            await runtime.close()
+
     async def test_request_raw_preserves_json_dict_response(self, httpx_mock):
         httpx_mock.add_response(json={"ok": True})
         client = AsyncWorkOSClient(