athena-intelligence 0.1.87__py3-none-any.whl → 0.1.89__py3-none-any.whl

athena/core/http_client.py

@@ -5,12 +5,18 @@ import email.utils
 import re
 import time
 import typing
+import urllib.parse
 from contextlib import asynccontextmanager, contextmanager
-from functools import wraps
 from random import random
 
 import httpx
 
+from .file import File, convert_file_dict_to_httpx_tuples
+from .jsonable_encoder import jsonable_encoder
+from .query_encoder import encode_query
+from .remove_none_from_dict import remove_none_from_dict
+from .request_options import RequestOptions
+
 INITIAL_RETRY_DELAY_SECONDS = 0.5
 MAX_RETRY_DELAY_SECONDS = 10
 MAX_RETRY_DELAY_SECONDS_FROM_HEADER = 30
@@ -82,49 +88,356 @@ def _should_retry(response: httpx.Response) -> bool:
     return response.status_code >= 500 or response.status_code in retriable_400s
 
 
+def remove_omit_from_dict(
+    original: typing.Dict[str, typing.Optional[typing.Any]], omit: typing.Optional[typing.Any]
+) -> typing.Dict[str, typing.Any]:
+    if omit is None:
+        return original
+    new: typing.Dict[str, typing.Any] = {}
+    for key, value in original.items():
+        if value is not omit:
+            new[key] = value
+    return new
+
+
+def maybe_filter_request_body(
+    data: typing.Optional[typing.Any],
+    request_options: typing.Optional[RequestOptions],
+    omit: typing.Optional[typing.Any],
+) -> typing.Optional[typing.Any]:
+    if data is None:
+        return None
+    elif not isinstance(data, typing.Mapping):
+        data_content = jsonable_encoder(data)
+    else:
+        data_content = {
+            **(jsonable_encoder(remove_omit_from_dict(data, omit))),  # type: ignore
+            **(
+                jsonable_encoder(request_options.get("additional_body_parameters", {}))
+                if request_options is not None
+                else {}
+            ),
+        }
+    return data_content
+
+
 class HttpClient:
-    def __init__(self, *, httpx_client: httpx.Client):
+    def __init__(
+        self,
+        *,
+        httpx_client: httpx.Client,
+        base_timeout: typing.Optional[float],
+        base_headers: typing.Dict[str, str],
+        base_url: typing.Optional[str] = None,
+    ):
+        self.base_url = base_url
+        self.base_timeout = base_timeout
+        self.base_headers = base_headers
         self.httpx_client = httpx_client
 
-    # Ensure that the signature of the `request` method is the same as the `httpx.Client.request` method
-    @wraps(httpx.Client.request)
+    def get_base_url(self, maybe_base_url: typing.Optional[str]) -> str:
+        base_url = self.base_url if maybe_base_url is None else maybe_base_url
+        if base_url is None:
+            raise ValueError("A base_url is required to make this request, please provide one and try again.")
+        return base_url
+
     def request(
-        self, *args: typing.Any, max_retries: int = 0, retries: int = 0, **kwargs: typing.Any
+        self,
+        path: typing.Optional[str] = None,
+        *,
+        method: str,
+        base_url: typing.Optional[str] = None,
+        params: typing.Optional[typing.Dict[str, typing.Any]] = None,
+        json: typing.Optional[typing.Any] = None,
+        data: typing.Optional[typing.Any] = None,
+        content: typing.Optional[typing.Union[bytes, typing.Iterator[bytes], typing.AsyncIterator[bytes]]] = None,
+        files: typing.Optional[typing.Dict[str, typing.Optional[typing.Union[File, typing.List[File]]]]] = None,
+        headers: typing.Optional[typing.Dict[str, typing.Any]] = None,
+        request_options: typing.Optional[RequestOptions] = None,
+        retries: int = 0,
+        omit: typing.Optional[typing.Any] = None,
     ) -> httpx.Response:
-        response = self.httpx_client.request(*args, **kwargs)
+        base_url = self.get_base_url(base_url)
+        timeout = (
+            request_options.get("timeout_in_seconds")
+            if request_options is not None and request_options.get("timeout_in_seconds") is not None
+            else self.base_timeout
+        )
+
+        response = self.httpx_client.request(
+            method=method,
+            url=urllib.parse.urljoin(f"{base_url}/", path),
+            headers=jsonable_encoder(
+                remove_none_from_dict(
+                    {
+                        **self.base_headers,
+                        **(headers if headers is not None else {}),
+                        **(request_options.get("additional_headers", {}) if request_options is not None else {}),
+                    }
+                )
+            ),
+            params=encode_query(
+                jsonable_encoder(
+                    remove_none_from_dict(
+                        remove_omit_from_dict(
+                            {
+                                **(params if params is not None else {}),
+                                **(
+                                    request_options.get("additional_query_parameters", {})
+                                    if request_options is not None
+                                    else {}
+                                ),
+                            },
+                            omit,
+                        )
+                    )
+                )
+            ),
+            json=maybe_filter_request_body(json, request_options, omit),
+            data=maybe_filter_request_body(data, request_options, omit),
+            content=content,
+            files=convert_file_dict_to_httpx_tuples(remove_none_from_dict(files)) if files is not None else None,
+            timeout=timeout,
+        )
+
+        max_retries: int = request_options.get("max_retries", 0) if request_options is not None else 0
         if _should_retry(response=response):
             if max_retries > retries:
                 time.sleep(_retry_timeout(response=response, retries=retries))
-                return self.request(max_retries=max_retries, retries=retries + 1, *args, **kwargs)
+                return self.request(
+                    path=path,
+                    method=method,
+                    base_url=base_url,
+                    params=params,
+                    json=json,
+                    content=content,
+                    files=files,
+                    headers=headers,
+                    request_options=request_options,
+                    retries=retries + 1,
+                    omit=omit,
+                )
+
         return response
 
-    @wraps(httpx.Client.stream)
     @contextmanager
-    def stream(self, *args: typing.Any, max_retries: int = 0, retries: int = 0, **kwargs: typing.Any) -> typing.Any:
-        with self.httpx_client.stream(*args, **kwargs) as stream:
+    def stream(
+        self,
+        path: typing.Optional[str] = None,
+        *,
+        method: str,
+        base_url: typing.Optional[str] = None,
+        params: typing.Optional[typing.Dict[str, typing.Any]] = None,
+        json: typing.Optional[typing.Any] = None,
+        data: typing.Optional[typing.Any] = None,
+        content: typing.Optional[typing.Union[bytes, typing.Iterator[bytes], typing.AsyncIterator[bytes]]] = None,
+        files: typing.Optional[typing.Dict[str, typing.Optional[typing.Union[File, typing.List[File]]]]] = None,
+        headers: typing.Optional[typing.Dict[str, typing.Any]] = None,
+        request_options: typing.Optional[RequestOptions] = None,
+        retries: int = 0,
+        omit: typing.Optional[typing.Any] = None,
+    ) -> typing.Iterator[httpx.Response]:
+        base_url = self.get_base_url(base_url)
+        timeout = (
+            request_options.get("timeout_in_seconds")
+            if request_options is not None and request_options.get("timeout_in_seconds") is not None
+            else self.base_timeout
+        )
+
+        with self.httpx_client.stream(
+            method=method,
+            url=urllib.parse.urljoin(f"{base_url}/", path),
+            headers=jsonable_encoder(
+                remove_none_from_dict(
+                    {
+                        **self.base_headers,
+                        **(headers if headers is not None else {}),
+                        **(request_options.get("additional_headers", {}) if request_options is not None else {}),
+                    }
+                )
+            ),
+            params=encode_query(
+                jsonable_encoder(
+                    remove_none_from_dict(
+                        remove_omit_from_dict(
+                            {
+                                **(params if params is not None else {}),
+                                **(
+                                    request_options.get("additional_query_parameters", {})
+                                    if request_options is not None
+                                    else {}
+                                ),
+                            },
+                            omit,
+                        )
+                    )
+                )
+            ),
+            json=maybe_filter_request_body(json, request_options, omit),
+            data=maybe_filter_request_body(data, request_options, omit),
+            content=content,
+            files=convert_file_dict_to_httpx_tuples(remove_none_from_dict(files)) if files is not None else None,
+            timeout=timeout,
+        ) as stream:
             yield stream
 
 
 class AsyncHttpClient:
-    def __init__(self, *, httpx_client: httpx.AsyncClient):
+    def __init__(
+        self,
+        *,
+        httpx_client: httpx.AsyncClient,
+        base_timeout: typing.Optional[float],
+        base_headers: typing.Dict[str, str],
+        base_url: typing.Optional[str] = None,
+    ):
+        self.base_url = base_url
+        self.base_timeout = base_timeout
+        self.base_headers = base_headers
         self.httpx_client = httpx_client
 
-    # Ensure that the signature of the `request` method is the same as the `httpx.Client.request` method
-    @wraps(httpx.AsyncClient.request)
+    def get_base_url(self, maybe_base_url: typing.Optional[str]) -> str:
+        base_url = self.base_url if maybe_base_url is None else maybe_base_url
+        if base_url is None:
+            raise ValueError("A base_url is required to make this request, please provide one and try again.")
+        return base_url
+
     async def request(
-        self, *args: typing.Any, max_retries: int = 0, retries: int = 0, **kwargs: typing.Any
+        self,
+        path: typing.Optional[str] = None,
+        *,
+        method: str,
+        base_url: typing.Optional[str] = None,
+        params: typing.Optional[typing.Dict[str, typing.Any]] = None,
+        json: typing.Optional[typing.Any] = None,
+        data: typing.Optional[typing.Any] = None,
+        content: typing.Optional[typing.Union[bytes, typing.Iterator[bytes], typing.AsyncIterator[bytes]]] = None,
+        files: typing.Optional[typing.Dict[str, typing.Optional[typing.Union[File, typing.List[File]]]]] = None,
+        headers: typing.Optional[typing.Dict[str, typing.Any]] = None,
+        request_options: typing.Optional[RequestOptions] = None,
+        retries: int = 0,
+        omit: typing.Optional[typing.Any] = None,
     ) -> httpx.Response:
-        response = await self.httpx_client.request(*args, **kwargs)
+        base_url = self.get_base_url(base_url)
+        timeout = (
+            request_options.get("timeout_in_seconds")
+            if request_options is not None and request_options.get("timeout_in_seconds") is not None
+            else self.base_timeout
+        )
+
+        # Add the input to each of these and do None-safety checks
+        response = await self.httpx_client.request(
+            method=method,
+            url=urllib.parse.urljoin(f"{base_url}/", path),
+            headers=jsonable_encoder(
+                remove_none_from_dict(
+                    {
+                        **self.base_headers,
+                        **(headers if headers is not None else {}),
+                        **(request_options.get("additional_headers", {}) if request_options is not None else {}),
+                    }
+                )
+            ),
+            params=encode_query(
+                jsonable_encoder(
+                    remove_none_from_dict(
+                        remove_omit_from_dict(
+                            {
+                                **(params if params is not None else {}),
+                                **(
+                                    request_options.get("additional_query_parameters", {})
+                                    if request_options is not None
+                                    else {}
+                                ),
+                            },
+                            omit,
+                        )
+                    )
+                )
+            ),
+            json=maybe_filter_request_body(json, request_options, omit),
+            data=maybe_filter_request_body(data, request_options, omit),
+            content=content,
+            files=convert_file_dict_to_httpx_tuples(remove_none_from_dict(files)) if files is not None else None,
+            timeout=timeout,
+        )
+
+        max_retries: int = request_options.get("max_retries", 0) if request_options is not None else 0
         if _should_retry(response=response):
             if max_retries > retries:
                 await asyncio.sleep(_retry_timeout(response=response, retries=retries))
-                return await self.request(max_retries=max_retries, retries=retries + 1, *args, **kwargs)
+                return await self.request(
+                    path=path,
+                    method=method,
+                    base_url=base_url,
+                    params=params,
+                    json=json,
+                    content=content,
+                    files=files,
+                    headers=headers,
+                    request_options=request_options,
+                    retries=retries + 1,
+                    omit=omit,
+                )
         return response
 
-    @wraps(httpx.AsyncClient.stream)
     @asynccontextmanager
     async def stream(
-        self, *args: typing.Any, max_retries: int = 0, retries: int = 0, **kwargs: typing.Any
-    ) -> typing.Any:
-        async with self.httpx_client.stream(*args, **kwargs) as stream:
+        self,
+        path: typing.Optional[str] = None,
+        *,
+        method: str,
+        base_url: typing.Optional[str] = None,
+        params: typing.Optional[typing.Dict[str, typing.Any]] = None,
+        json: typing.Optional[typing.Any] = None,
+        data: typing.Optional[typing.Any] = None,
+        content: typing.Optional[typing.Union[bytes, typing.Iterator[bytes], typing.AsyncIterator[bytes]]] = None,
+        files: typing.Optional[typing.Dict[str, typing.Optional[typing.Union[File, typing.List[File]]]]] = None,
+        headers: typing.Optional[typing.Dict[str, typing.Any]] = None,
+        request_options: typing.Optional[RequestOptions] = None,
+        retries: int = 0,
+        omit: typing.Optional[typing.Any] = None,
+    ) -> typing.AsyncIterator[httpx.Response]:
+        base_url = self.get_base_url(base_url)
+        timeout = (
+            request_options.get("timeout_in_seconds")
+            if request_options is not None and request_options.get("timeout_in_seconds") is not None
+            else self.base_timeout
+        )
+
+        async with self.httpx_client.stream(
+            method=method,
+            url=urllib.parse.urljoin(f"{base_url}/", path),
+            headers=jsonable_encoder(
+                remove_none_from_dict(
+                    {
+                        **self.base_headers,
+                        **(headers if headers is not None else {}),
+                        **(request_options.get("additional_headers", {}) if request_options is not None else {}),
+                    }
+                )
+            ),
+            params=encode_query(
+                jsonable_encoder(
+                    remove_none_from_dict(
+                        remove_omit_from_dict(
+                            {
+                                **(params if params is not None else {}),
+                                **(
+                                    request_options.get("additional_query_parameters", {})
+                                    if request_options is not None
+                                    else {}
+                                ),
+                            },
+                            omit=omit,
+                        )
+                    )
+                )
+            ),
+            json=maybe_filter_request_body(json, request_options, omit),
+            data=maybe_filter_request_body(data, request_options, omit),
+            content=content,
+            files=convert_file_dict_to_httpx_tuples(remove_none_from_dict(files)) if files is not None else None,
+            timeout=timeout,
+        ) as stream:
             yield stream

athena/core/remove_none_from_dict.py

@@ -1,9 +1,9 @@
 # This file was auto-generated by Fern from our API Definition.
 
-from typing import Any, Dict, Optional
+from typing import Any, Dict, Mapping, Optional
 
 
-def remove_none_from_dict(original: Dict[str, Optional[Any]]) -> Dict[str, Any]:
+def remove_none_from_dict(original: Mapping[str, Optional[Any]]) -> Dict[str, Any]:
     new: Dict[str, Any] = {}
     for key, value in original.items():
         if value is not None:

athena/dataset/client.py

@@ -1,15 +1,11 @@
 # This file was auto-generated by Fern from our API Definition.
 
 import typing
-import urllib.parse
 from json.decoder import JSONDecodeError
 
 from ..core.api_error import ApiError
 from ..core.client_wrapper import AsyncClientWrapper, SyncClientWrapper
-from ..core.jsonable_encoder import jsonable_encoder
 from ..core.pydantic_utilities import pydantic_v1
-from ..core.query_encoder import encode_query
-from ..core.remove_none_from_dict import remove_none_from_dict
 from ..core.request_options import RequestOptions
 from ..errors.unprocessable_entity_error import UnprocessableEntityError
 from ..types.get_datasets_response import GetDatasetsResponse
@@ -25,7 +21,7 @@ class DatasetClient:
         *,
         page: typing.Optional[int] = None,
         page_size: typing.Optional[int] = None,
-        request_options: typing.Optional[RequestOptions] = None,
+        request_options: typing.Optional[RequestOptions] = None
     ) -> GetDatasetsResponse:
         """
         Parameters
@@ -54,36 +50,10 @@ class DatasetClient:
         client.dataset.get()
         """
         _response = self._client_wrapper.httpx_client.request(
+            "api/v0/datasets",
             method="GET",
-            url=urllib.parse.urljoin(f"{self._client_wrapper.get_base_url()}/", "api/v0/datasets"),
-            params=encode_query(
-                jsonable_encoder(
-                    remove_none_from_dict(
-                        {
-                            "page": page,
-                            "page_size": page_size,
-                            **(
-                                request_options.get("additional_query_parameters", {})
-                                if request_options is not None
-                                else {}
-                            ),
-                        }
-                    )
-                )
-            ),
-            headers=jsonable_encoder(
-                remove_none_from_dict(
-                    {
-                        **self._client_wrapper.get_headers(),
-                        **(request_options.get("additional_headers", {}) if request_options is not None else {}),
-                    }
-                )
-            ),
-            timeout=request_options.get("timeout_in_seconds")
-            if request_options is not None and request_options.get("timeout_in_seconds") is not None
-            else self._client_wrapper.get_timeout(),
-            retries=0,
-            max_retries=request_options.get("max_retries") if request_options is not None else 0,  # type: ignore
+            params={"page": page, "page_size": page_size},
+            request_options=request_options,
         )
         if 200 <= _response.status_code < 300:
             return pydantic_v1.parse_obj_as(GetDatasetsResponse, _response.json())  # type: ignore
@@ -107,7 +77,7 @@ class AsyncDatasetClient:
         *,
         page: typing.Optional[int] = None,
         page_size: typing.Optional[int] = None,
-        request_options: typing.Optional[RequestOptions] = None,
+        request_options: typing.Optional[RequestOptions] = None
     ) -> GetDatasetsResponse:
         """
         Parameters
@@ -136,36 +106,10 @@ class AsyncDatasetClient:
         await client.dataset.get()
         """
         _response = await self._client_wrapper.httpx_client.request(
+            "api/v0/datasets",
             method="GET",
-            url=urllib.parse.urljoin(f"{self._client_wrapper.get_base_url()}/", "api/v0/datasets"),
-            params=encode_query(
-                jsonable_encoder(
-                    remove_none_from_dict(
-                        {
-                            "page": page,
-                            "page_size": page_size,
-                            **(
-                                request_options.get("additional_query_parameters", {})
-                                if request_options is not None
-                                else {}
-                            ),
-                        }
-                    )
-                )
-            ),
-            headers=jsonable_encoder(
-                remove_none_from_dict(
-                    {
-                        **self._client_wrapper.get_headers(),
-                        **(request_options.get("additional_headers", {}) if request_options is not None else {}),
-                    }
-                )
-            ),
-            timeout=request_options.get("timeout_in_seconds")
-            if request_options is not None and request_options.get("timeout_in_seconds") is not None
-            else self._client_wrapper.get_timeout(),
-            retries=0,
-            max_retries=request_options.get("max_retries") if request_options is not None else 0,  # type: ignore
+            params={"page": page, "page_size": page_size},
+            request_options=request_options,
         )
         if 200 <= _response.status_code < 300:
             return pydantic_v1.parse_obj_as(GetDatasetsResponse, _response.json())  # type: ignore