Skip to content

Commit

Permalink
fix: fix aiohttp warning. (#834)
Browse files Browse the repository at this point in the history
  • Loading branch information
overcat committed Dec 4, 2023
1 parent 80bfc55 commit 7d2f4de
Showing 1 changed file with 23 additions and 14 deletions.
37 changes: 23 additions & 14 deletions stellar_sdk/client/aiohttp_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -96,11 +96,7 @@ def __init__(
self.backoff_factor: Optional[float] = backoff_factor
self.request_timeout: float = request_timeout
self.post_timeout: float = post_timeout
# init session
if pool_size is None:
connector = aiohttp.TCPConnector()
else:
connector = aiohttp.TCPConnector(limit=pool_size)
self.__kwargs = kwargs

self.user_agent: Optional[str] = USER_AGENT
if user_agent:
Expand All @@ -114,14 +110,7 @@ def __init__(
if custom_headers:
self.headers = {**self.headers, **custom_headers}

session = aiohttp.ClientSession(
headers=self.headers.copy(),
connector=connector,
timeout=aiohttp.ClientTimeout(total=request_timeout),
**kwargs,
)

self._session: aiohttp.ClientSession = session
self._session: Optional[aiohttp.ClientSession] = None
self._sse_session: Optional[aiohttp.ClientSession] = None

async def get(self, url: str, params: Dict[str, str] = None) -> Response:
Expand All @@ -132,6 +121,8 @@ async def get(self, url: str, params: Dict[str, str] = None) -> Response:
:return: the response from server
:raise: :exc:`ConnectionError <stellar_sdk.exceptions.ConnectionError>`
"""
await self.__init_session()
assert self._session is not None
try:
response = await self._session.get(url, params=params)
return Response(
Expand All @@ -154,6 +145,8 @@ async def post(
:return: the response from server
:raise: :exc:`ConnectionError <stellar_sdk.exceptions.ConnectionError>`
"""
await self.__init_session()
assert self._session is not None
try:
response = await self._session.post(
url,
Expand Down Expand Up @@ -235,6 +228,21 @@ async def stream(
)
await asyncio.sleep(retry)

async def __init_session(self):
# init session
if self._session is None:
if self.pool_size is None:
connector = aiohttp.TCPConnector()
else:
connector = aiohttp.TCPConnector(limit=self.pool_size)

self._session = aiohttp.ClientSession(
headers=self.headers.copy(),
connector=connector,
timeout=aiohttp.ClientTimeout(total=self.request_timeout),
**self.__kwargs,
)

async def __aenter__(self) -> "AiohttpClient":
return self

Expand All @@ -246,7 +254,8 @@ async def close(self) -> None:
Release all acquired resources.
"""
await self._session.__aexit__(None, None, None)
if self._session is not None:
await self._session.__aexit__(None, None, None)
if self._sse_session is not None:
await self._sse_session.__aexit__(None, None, None)

Expand Down

0 comments on commit 7d2f4de

Please sign in to comment.