Allow specifying return type for http methods

pull/128/head
stijndcl 2022-08-29 02:04:42 +02:00
parent 14e0472954
commit 994ff01de1
1 changed files with 10 additions and 6 deletions

View File

@ -1,6 +1,6 @@
import logging import logging
from contextlib import asynccontextmanager from contextlib import asynccontextmanager
from typing import AsyncGenerator from typing import AsyncGenerator, Type, TypeVar
from aiohttp import ClientResponse, ClientSession, ContentTypeError from aiohttp import ClientResponse, ClientSession, ContentTypeError
@ -12,6 +12,9 @@ logger = logging.getLogger(__name__)
__all__ = ["ensure_get", "ensure_post"] __all__ = ["ensure_get", "ensure_post"]
T = TypeVar("T", str, dict)
def request_successful(response: ClientResponse) -> bool: def request_successful(response: ClientResponse) -> bool:
"""Check if a request was successful or not""" """Check if a request was successful or not"""
return 200 <= response.status < 300 return 200 <= response.status < 300
@ -19,12 +22,12 @@ def request_successful(response: ClientResponse) -> bool:
@asynccontextmanager @asynccontextmanager
async def ensure_get( async def ensure_get(
http_session: ClientSession, endpoint: str, *, log_exceptions: bool = True http_session: ClientSession, endpoint: str, *, return_type: Type[T] = dict, log_exceptions: bool = True
) -> AsyncGenerator[dict, None]: ) -> AsyncGenerator[T, None]:
"""Context manager that automatically raises an exception if a GET-request fails""" """Context manager that automatically raises an exception if a GET-request fails"""
async with http_session.get(endpoint) as response: async with http_session.get(endpoint) as response:
try: try:
content = await response.json() content = (await response.json()) if return_type == dict else (await response.text())
except ContentTypeError: except ContentTypeError:
content = await response.text() content = await response.text()
@ -43,14 +46,15 @@ async def ensure_post(
endpoint: str, endpoint: str,
payload: dict, payload: dict,
*, *,
return_type: Type[T] = dict,
log_exceptions: bool = True, log_exceptions: bool = True,
expect_return: bool = True expect_return: bool = True
) -> AsyncGenerator[dict, None]: ) -> AsyncGenerator[T, None]:
"""Context manager that automatically raises an exception if a POST-request fails""" """Context manager that automatically raises an exception if a POST-request fails"""
async with http_session.post(endpoint, data=payload) as response: async with http_session.post(endpoint, data=payload) as response:
if not request_successful(response): if not request_successful(response):
try: try:
content = await response.json() content = (await response.json()) if return_type == dict else (await response.text())
except ContentTypeError: except ContentTypeError:
content = await response.text() content = await response.text()