Skip to content

Commit 788554a

Browse files
CopilotArzelaAscoIicursoragent
authored
fix: add backward compatibility for pyrate-limiter 4.0.0+ (#296)
* Initial plan * Add backward compatibility for pyrate-limiter 3.x and 4.x Co-authored-by: ArzelaAscoIi <37148029+ArzelaAscoIi@users.noreply.github.com> * Refactor: Extract rate limiting logic and add error handling Co-authored-by: ArzelaAscoIi <37148029+ArzelaAscoIi@users.noreply.github.com> * Address code review feedback: improve error handling and documentation Co-authored-by: ArzelaAscoIi <37148029+ArzelaAscoIi@users.noreply.github.com> * fix: resolve mypy and formatting failures in pyrate-limiter compat - Replace inspect.signature() version detection with try/except on init - Use **kwargs for try_acquire() to avoid mypy call-arg error - Fix trailing whitespace formatting issues Co-authored-by: Cursor <cursoragent@cursor.com> --------- Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Co-authored-by: ArzelaAscoIi <37148029+ArzelaAscoIi@users.noreply.github.com> Co-authored-by: ArzelaAscoIi <kristof.herrmann@rwth-aachen.de> Co-authored-by: Cursor <cursoragent@cursor.com>
1 parent 3b5d228 commit 788554a

1 file changed

Lines changed: 19 additions & 8 deletions

File tree

deepset_cloud_sdk/_s3/upload.py

Lines changed: 19 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@
77
from http import HTTPStatus
88
from pathlib import Path
99
from types import TracebackType
10-
from typing import Any, Coroutine, List, Optional, Sequence, Type, Union
10+
from typing import Any, Coroutine, Dict, List, Optional, Sequence, Type, Union
1111

1212
import aiofiles
1313
import aiohttp
@@ -94,7 +94,16 @@ def __init__(
9494
"""
9595
self.connector = aiohttp.TCPConnector(limit=concurrency)
9696
self.semaphore = asyncio.BoundedSemaphore(concurrency)
97-
self.limiter = Limiter(rate_limit, raise_when_fail=False, max_delay=Duration.SECOND * 1)
97+
98+
try:
99+
# pyrate-limiter 3.x
100+
self.limiter = Limiter(rate_limit, raise_when_fail=False, max_delay=Duration.SECOND * 1)
101+
self._try_acquire_kwargs: Dict[str, Any] = {}
102+
except TypeError:
103+
# pyrate-limiter 4.0.0+ removed raise_when_fail and max_delay
104+
self.limiter = Limiter(rate_limit)
105+
self._try_acquire_kwargs = {"blocking": False}
106+
98107
self.max_attempts = max_attempts
99108

100109
async def __aenter__(self) -> "S3":
@@ -110,15 +119,15 @@ async def __aexit__(
110119
"""Exit the context manager."""
111120
await self.connector.close()
112121

113-
# Handle limiter cleanup based on available methods
114-
# Support both older and newer versions of pyrate_limiter
115-
# In version 3.7.0, the dispose method was added to the Limiter class
116-
# See diff here: https://github.com/vutran1710/PyrateLimiter/compare/v3.6.2...master
117122
try:
118123
list(map(self.limiter.dispose, self.limiter.buckets()))
119124
except AttributeError:
120125
pass
121126

127+
def _rate_limit_acquire(self) -> None:
128+
"""Acquire a rate limit token with backward compat for pyrate-limiter 3.x and 4.x."""
129+
self.limiter.try_acquire("", **self._try_acquire_kwargs)
130+
122131
async def _upload_file_with_retries(
123132
self,
124133
file_name: str,
@@ -158,7 +167,8 @@ async def _upload_file(
158167
file_data = self._build_file_data(content, aws_safe_name, aws_config)
159168

160169
try:
161-
self.limiter.try_acquire("") # rate limit requests
170+
self._rate_limit_acquire()
171+
162172
async with client_session.post(
163173
aws_config.url,
164174
data=file_data,
@@ -172,7 +182,8 @@ async def _upload_file(
172182
# for example during automatic redirects. See https://github.com/aio-libs/aiohttp/issues/5577
173183
redirect_url = response.headers["Location"]
174184
file_data = self._build_file_data(content, aws_safe_name, aws_config)
175-
self.limiter.try_acquire("") # rate limit requests
185+
self._rate_limit_acquire()
186+
176187
async with client_session.post(
177188
redirect_url,
178189
json=file_data,

0 commit comments

Comments
 (0)