retry using a decorator

Author: denniswittich

learning_loop_node/loop_communication.py

@@ -11,6 +11,23 @@
 logging.basicConfig(level=logging.INFO)
 
 SLEEP_TIME_ON_429 = 5
+MAX_RETRIES_ON_429 = 20
+
+
+def retry_on_429(func: Callable[..., Awaitable]) -> Callable[..., Awaitable]:
+    """Decorator that retries requests that receive a 429 status code."""
+    async def wrapper(*args, **kwargs) -> httpx.Response:
+        retries = 0
+        while retries < MAX_RETRIES_ON_429:
+            response = await func(*args, **kwargs)
+            if response.status_code != 429:
+                return response
+
+            retries += 1
+            await asyncio.sleep(SLEEP_TIME_ON_429)
+
+        return response
+    return wrapper
 
 
 class LoopCommunicationException(Exception):
@@ -98,19 +115,17 @@ async def get(self, path: str, requires_login: bool = True, api_prefix: str = '/
             return await self.retry_on_401(self._get, path, api_prefix)
         return await self._get(path, api_prefix)
 
+    @retry_on_429
     async def _get(self, path: str, api_prefix: str) -> httpx.Response:
-        response = await self.async_client.get(api_prefix+path)
-        while response.status_code == 429:
-            await asyncio.sleep(SLEEP_TIME_ON_429)
-            response = await self.async_client.get(api_prefix+path)
-        return response
+        return await self.async_client.get(api_prefix+path)
 
     async def put(self, path: str, files: Optional[List[str]] = None, requires_login: bool = True, api_prefix: str = '/api', **kwargs) -> httpx.Response:
         if requires_login:
             await self.ensure_login()
             return await self.retry_on_401(self._put, path, files, api_prefix, **kwargs)
         return await self._put(path, files, api_prefix, **kwargs)
 
+    @retry_on_429
     async def _put(self, path: str, files: Optional[List[str]], api_prefix: str, **kwargs) -> httpx.Response:
         if files is None:
             return await self.async_client.put(api_prefix+path, **kwargs)
@@ -127,9 +142,6 @@ async def _put(self, path: str, files: Optional[List[str]], api_prefix: str, **k
         try:
             file_list = [('files', fh) for fh in file_handles]  # Use file handles
             response = await self.async_client.put(api_prefix+path, files=file_list)
-            while response.status_code == 429:
-                await asyncio.sleep(SLEEP_TIME_ON_429)
-                response = await self.async_client.put(api_prefix+path, files=file_list)
         finally:
             for fh in file_handles:
                 fh.close()  # Ensure all files are closed
@@ -142,22 +154,16 @@ async def post(self, path: str, requires_login: bool = True, api_prefix: str = '
             return await self.retry_on_401(self._post, path, api_prefix, **kwargs)
         return await self._post(path, api_prefix, **kwargs)
 
+    @retry_on_429
     async def _post(self, path, api_prefix='/api', **kwargs) -> httpx.Response:
-        response = await self.async_client.post(api_prefix+path, **kwargs)
-        while response.status_code == 429:
-            await asyncio.sleep(SLEEP_TIME_ON_429)
-            response = await self.async_client.post(api_prefix+path, **kwargs)
-        return response
+        return await self.async_client.post(api_prefix+path, **kwargs)
 
     async def delete(self, path: str, requires_login: bool = True, api_prefix: str = '/api', **kwargs) -> httpx.Response:
         if requires_login:
             await self.ensure_login()
             return await self.retry_on_401(self._delete, path, api_prefix, **kwargs)
         return await self._delete(path, api_prefix, **kwargs)
 
+    @retry_on_429
     async def _delete(self, path, api_prefix, **kwargs) -> httpx.Response:
-        response = await self.async_client.delete(api_prefix+path, **kwargs)
-        while response.status_code == 429:
-            await asyncio.sleep(SLEEP_TIME_ON_429)
-            response = await self.async_client.delete(api_prefix+path, **kwargs)
-        return response
+        return await self.async_client.delete(api_prefix+path, **kwargs)