From fdf0d7096e2e66c4e69cbf9d99004907e98b47d6 Mon Sep 17 00:00:00 2001 From: Stanley Phu Date: Wed, 4 Nov 2026 22:45:16 -0800 Subject: [PATCH] fix(client): Keep the public client's type in `with WorkOSClient() as workos` WorkOSClient.__enter__ and AsyncWorkOSClient.__aenter__ were annotated with the base classes in _base_client.py, so type checkers saw the target of a with statement as a client without any resources or helpers. They return Self now. Co-Authored-By: Claude Opus 5.5 --- src/workos/_base_client.py | 10 ++++++++-- tests/test_generated_client.py | 16 ++++++++++++++++ 2 files changed, 24 insertions(+), 2 deletions(-) diff --git a/src/workos/_base_client.py b/src/workos/_base_client.py index 4f743d42..bffea76e 100644 --- a/src/workos/_base_client.py +++ b/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: diff --git a/tests/test_generated_client.py b/tests/test_generated_client.py index c42234bc..f0a8bd2c 100644 --- a/tests/test_generated_client.py +++ b/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(