copy changes from stainless

Author: axl1313

src/alpha/__init__.py

@@ -3,7 +3,18 @@
 from . import types
 from ._types import NOT_GIVEN, Omit, NoneType, NotGiven, Transport, ProxiesTypes
 from ._utils import file_from_path
-from ._client import Alpha, Client, Stream, Timeout, Transport, AsyncAlpha, AsyncClient, AsyncStream, RequestOptions
+from ._client import (
+    ENVIRONMENTS,
+    Alpha,
+    Client,
+    Stream,
+    Timeout,
+    Transport,
+    AsyncAlpha,
+    AsyncClient,
+    AsyncStream,
+    RequestOptions,
+)
 from ._models import BaseModel
 from ._version import __title__, __version__
 from ._response import APIResponse as APIResponse, AsyncAPIResponse as AsyncAPIResponse
@@ -59,6 +70,7 @@
     "AsyncStream",
     "Alpha",
     "AsyncAlpha",
+    "ENVIRONMENTS",
     "file_from_path",
     "BaseModel",
     "DEFAULT_TIMEOUT",

src/alpha/_client.py

@@ -3,8 +3,8 @@
 from __future__ import annotations
 
 import os
-from typing import Any, Union, Mapping
-from typing_extensions import Self, override
+from typing import Any, Dict, Union, Mapping, cast
+from typing_extensions import Self, Literal, override
 
 import httpx
 
@@ -13,6 +13,7 @@
 from ._types import (
     NOT_GIVEN,
     Omit,
+    Headers,
     Timeout,
     NotGiven,
     Transport,
@@ -26,7 +27,7 @@
 from ._version import __version__
 from .resources import health
 from ._streaming import Stream as Stream, AsyncStream as AsyncStream
-from ._exceptions import AlphaError, APIStatusError
+from ._exceptions import APIStatusError
 from ._base_client import (
     DEFAULT_MAX_RETRIES,
     SyncAPIClient,
@@ -36,7 +37,22 @@
 from .resources.projects import projects
 from .resources.organizations import organizations
 
-__all__ = ["Timeout", "Transport", "ProxiesTypes", "RequestOptions", "Alpha", "AsyncAlpha", "Client", "AsyncClient"]
+__all__ = [
+    "ENVIRONMENTS",
+    "Timeout",
+    "Transport",
+    "ProxiesTypes",
+    "RequestOptions",
+    "Alpha",
+    "AsyncAlpha",
+    "Client",
+    "AsyncClient",
+]
+
+ENVIRONMENTS: Dict[str, str] = {
+    "production": "https://api-alpha-o3gxj3oajfu.cleanlab.ai/",
+    "local": "http://localhost:8080",
+}
 
 
 class Alpha(SyncAPIClient):
@@ -48,17 +64,20 @@ class Alpha(SyncAPIClient):
     with_streaming_response: AlphaWithStreamedResponse
 
     # client options
-    bearer_token: str
-    api_key: str
-    access_key: str
+    bearer_token: str | None
+    api_key: str | None
+    access_key: str | None
+
+    _environment: Literal["production", "local"] | NotGiven
 
     def __init__(
         self,
         *,
         bearer_token: str | None = None,
         api_key: str | None = None,
         access_key: str | None = None,
-        base_url: str | httpx.URL | None = None,
+        environment: Literal["production", "local"] | NotGiven = NOT_GIVEN,
+        base_url: str | httpx.URL | None | NotGiven = NOT_GIVEN,
         timeout: Union[float, Timeout, None, NotGiven] = NOT_GIVEN,
         max_retries: int = DEFAULT_MAX_RETRIES,
         default_headers: Mapping[str, str] | None = None,
@@ -86,32 +105,41 @@ def __init__(
         """
         if bearer_token is None:
             bearer_token = os.environ.get("BEARER_TOKEN")
-        if bearer_token is None:
-            raise AlphaError(
-                "The bearer_token client option must be set either by passing bearer_token to the client or by setting the BEARER_TOKEN environment variable"
-            )
         self.bearer_token = bearer_token
 
         if api_key is None:
             api_key = os.environ.get("X_API_KEY")
-        if api_key is None:
-            raise AlphaError(
-                "The api_key client option must be set either by passing api_key to the client or by setting the X_API_KEY environment variable"
-            )
         self.api_key = api_key
 
         if access_key is None:
             access_key = os.environ.get("X_ACCESS_KEY")
-        if access_key is None:
-            raise AlphaError(
-                "The access_key client option must be set either by passing access_key to the client or by setting the X_ACCESS_KEY environment variable"
-            )
         self.access_key = access_key
 
-        if base_url is None:
-            base_url = os.environ.get("ALPHA_BASE_URL")
-        if base_url is None:
-            base_url = f"https://localhost:8080/test-api"
+        self._environment = environment
+
+        base_url_env = os.environ.get("ALPHA_BASE_URL")
+        if is_given(base_url) and base_url is not None:
+            # cast required because mypy doesn't understand the type narrowing
+            base_url = cast("str | httpx.URL", base_url)  # pyright: ignore[reportUnnecessaryCast]
+        elif is_given(environment):
+            if base_url_env and base_url is not None:
+                raise ValueError(
+                    "Ambiguous URL; The `ALPHA_BASE_URL` env var and the `environment` argument are given. If you want to use the environment, you must pass base_url=None",
+                )
+
+            try:
+                base_url = ENVIRONMENTS[environment]
+            except KeyError as exc:
+                raise ValueError(f"Unknown environment: {environment}") from exc
+        elif base_url_env is not None:
+            base_url = base_url_env
+        else:
+            self._environment = environment = "production"
+
+            try:
+                base_url = ENVIRONMENTS[environment]
+            except KeyError as exc:
+                raise ValueError(f"Unknown environment: {environment}") from exc
 
         super().__init__(
             version=__version__,
@@ -150,16 +178,22 @@ def auth_headers(self) -> dict[str, str]:
     @property
     def _http_bearer(self) -> dict[str, str]:
         bearer_token = self.bearer_token
+        if bearer_token is None:
+            return {}
         return {"Authorization": f"Bearer {bearer_token}"}
 
     @property
     def _authenticated_api_key(self) -> dict[str, str]:
         api_key = self.api_key
+        if api_key is None:
+            return {}
         return {"X-API-Key": api_key}
 
     @property
     def _public_access_key(self) -> dict[str, str]:
         access_key = self.access_key
+        if access_key is None:
+            return {}
         return {"X-Access-Key": access_key}
 
     @property
@@ -171,12 +205,34 @@ def default_headers(self) -> dict[str, str | Omit]:
             **self._custom_headers,
         }
 
+    @override
+    def _validate_headers(self, headers: Headers, custom_headers: Headers) -> None:
+        if self.bearer_token and headers.get("Authorization"):
+            return
+        if isinstance(custom_headers.get("Authorization"), Omit):
+            return
+
+        if self.api_key and headers.get("X-API-Key"):
+            return
+        if isinstance(custom_headers.get("X-API-Key"), Omit):
+            return
+
+        if self.access_key and headers.get("X-Access-Key"):
+            return
+        if isinstance(custom_headers.get("X-Access-Key"), Omit):
+            return
+
+        raise TypeError(
+            '"Could not resolve authentication method. Expected one of bearer_token, api_key or access_key to be set. Or for one of the `Authorization`, `X-API-Key` or `X-Access-Key` headers to be explicitly omitted"'
+        )
+
     def copy(
         self,
         *,
         bearer_token: str | None = None,
         api_key: str | None = None,
         access_key: str | None = None,
+        environment: Literal["production", "local"] | None = None,
         base_url: str | httpx.URL | None = None,
         timeout: float | Timeout | None | NotGiven = NOT_GIVEN,
         http_client: httpx.Client | None = None,
@@ -214,6 +270,7 @@ def copy(
             api_key=api_key or self.api_key,
             access_key=access_key or self.access_key,
             base_url=base_url or self.base_url,
+            environment=environment or self._environment,
             timeout=self.timeout if isinstance(timeout, NotGiven) else timeout,
             http_client=http_client,
             max_retries=max_retries if is_given(max_retries) else self.max_retries,
@@ -269,17 +326,20 @@ class AsyncAlpha(AsyncAPIClient):
     with_streaming_response: AsyncAlphaWithStreamedResponse
 
     # client options
-    bearer_token: str
-    api_key: str
-    access_key: str
+    bearer_token: str | None
+    api_key: str | None
+    access_key: str | None
+
+    _environment: Literal["production", "local"] | NotGiven
 
     def __init__(
         self,
         *,
         bearer_token: str | None = None,
         api_key: str | None = None,
         access_key: str | None = None,
-        base_url: str | httpx.URL | None = None,
+        environment: Literal["production", "local"] | NotGiven = NOT_GIVEN,
+        base_url: str | httpx.URL | None | NotGiven = NOT_GIVEN,
         timeout: Union[float, Timeout, None, NotGiven] = NOT_GIVEN,
         max_retries: int = DEFAULT_MAX_RETRIES,
         default_headers: Mapping[str, str] | None = None,
@@ -307,32 +367,41 @@ def __init__(
         """
         if bearer_token is None:
             bearer_token = os.environ.get("BEARER_TOKEN")
-        if bearer_token is None:
-            raise AlphaError(
-                "The bearer_token client option must be set either by passing bearer_token to the client or by setting the BEARER_TOKEN environment variable"
-            )
         self.bearer_token = bearer_token
 
         if api_key is None:
             api_key = os.environ.get("X_API_KEY")
-        if api_key is None:
-            raise AlphaError(
-                "The api_key client option must be set either by passing api_key to the client or by setting the X_API_KEY environment variable"
-            )
         self.api_key = api_key
 
         if access_key is None:
             access_key = os.environ.get("X_ACCESS_KEY")
-        if access_key is None:
-            raise AlphaError(
-                "The access_key client option must be set either by passing access_key to the client or by setting the X_ACCESS_KEY environment variable"
-            )
         self.access_key = access_key
 
-        if base_url is None:
-            base_url = os.environ.get("ALPHA_BASE_URL")
-        if base_url is None:
-            base_url = f"https://localhost:8080/test-api"
+        self._environment = environment
+
+        base_url_env = os.environ.get("ALPHA_BASE_URL")
+        if is_given(base_url) and base_url is not None:
+            # cast required because mypy doesn't understand the type narrowing
+            base_url = cast("str | httpx.URL", base_url)  # pyright: ignore[reportUnnecessaryCast]
+        elif is_given(environment):
+            if base_url_env and base_url is not None:
+                raise ValueError(
+                    "Ambiguous URL; The `ALPHA_BASE_URL` env var and the `environment` argument are given. If you want to use the environment, you must pass base_url=None",
+                )
+
+            try:
+                base_url = ENVIRONMENTS[environment]
+            except KeyError as exc:
+                raise ValueError(f"Unknown environment: {environment}") from exc
+        elif base_url_env is not None:
+            base_url = base_url_env
+        else:
+            self._environment = environment = "production"
+
+            try:
+                base_url = ENVIRONMENTS[environment]
+            except KeyError as exc:
+                raise ValueError(f"Unknown environment: {environment}") from exc
 
         super().__init__(
             version=__version__,
@@ -371,16 +440,22 @@ def auth_headers(self) -> dict[str, str]:
     @property
     def _http_bearer(self) -> dict[str, str]:
         bearer_token = self.bearer_token
+        if bearer_token is None:
+            return {}
         return {"Authorization": f"Bearer {bearer_token}"}
 
     @property
     def _authenticated_api_key(self) -> dict[str, str]:
         api_key = self.api_key
+        if api_key is None:
+            return {}
         return {"X-API-Key": api_key}
 
     @property
     def _public_access_key(self) -> dict[str, str]:
         access_key = self.access_key
+        if access_key is None:
+            return {}
         return {"X-Access-Key": access_key}
 
     @property
@@ -392,12 +467,34 @@ def default_headers(self) -> dict[str, str | Omit]:
             **self._custom_headers,
         }
 
+    @override
+    def _validate_headers(self, headers: Headers, custom_headers: Headers) -> None:
+        if self.bearer_token and headers.get("Authorization"):
+            return
+        if isinstance(custom_headers.get("Authorization"), Omit):
+            return
+
+        if self.api_key and headers.get("X-API-Key"):
+            return
+        if isinstance(custom_headers.get("X-API-Key"), Omit):
+            return
+
+        if self.access_key and headers.get("X-Access-Key"):
+            return
+        if isinstance(custom_headers.get("X-Access-Key"), Omit):
+            return
+
+        raise TypeError(
+            '"Could not resolve authentication method. Expected one of bearer_token, api_key or access_key to be set. Or for one of the `Authorization`, `X-API-Key` or `X-Access-Key` headers to be explicitly omitted"'
+        )
+
     def copy(
         self,
         *,
         bearer_token: str | None = None,
         api_key: str | None = None,
         access_key: str | None = None,
+        environment: Literal["production", "local"] | None = None,
         base_url: str | httpx.URL | None = None,
         timeout: float | Timeout | None | NotGiven = NOT_GIVEN,
         http_client: httpx.AsyncClient | None = None,
@@ -435,6 +532,7 @@ def copy(
             api_key=api_key or self.api_key,
             access_key=access_key or self.access_key,
             base_url=base_url or self.base_url,
+            environment=environment or self._environment,
             timeout=self.timeout if isinstance(timeout, NotGiven) else timeout,
             http_client=http_client,
             max_retries=max_retries if is_given(max_retries) else self.max_retries,