From 994ff01de1355b54d9930221d5108a63c1f07e3a Mon Sep 17 00:00:00 2001
From: stijndcl <declercq.stijn@outlook.com>
Date: Mon, 29 Aug 2022 02:04:42 +0200
Subject: [PATCH] Allow specifying return type for http methods

---
 didier/utils/http/requests.py | 16 ++++++++++------
 1 file changed, 10 insertions(+), 6 deletions(-)

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