Implement delayed session creation and reuse existing session

Co-authored-by: Joel Payne <[email protected]>

Author: DevilXD

asyncprawcore/auth.py

@@ -5,7 +5,8 @@
 import inspect
 import time
 from abc import ABC, abstractmethod
-from typing import TYPE_CHECKING, Any, Awaitable, Callable
+from contextlib import asynccontextmanager
+from typing import TYPE_CHECKING, Any, AsyncContextManager, Awaitable, Callable
 
 from aiohttp import ClientRequest
 from aiohttp.helpers import BasicAuth
@@ -49,19 +50,20 @@ def __init__(
         self.client_id = client_id
         self.redirect_uri = redirect_uri
 
+    @asynccontextmanager
     async def _post(
         self, url: str, success_status: int = codes["ok"], **data: Any
-    ) -> ClientResponse:
-        response = await self._requestor.request(
+    ) -> Callable[..., AsyncContextManager[ClientResponse]]:
+        async with self._requestor.request(
             "POST",
             url,
             auth=self._auth(),
             data=sorted(data.items()),
             headers={"Connection": "close"},
-        )
-        if response.status != success_status:
-            raise ResponseException(response)
-        return response
+        ) as response:
+            if response.status != success_status:
+                raise ResponseException(response)
+            yield response
 
     def authorize_url(
         self, duration: str, scopes: list[str], state: str, implicit: bool = False
@@ -126,7 +128,8 @@ async def revoke_token(self, token: str, token_type: str | None = None):
         if token_type is not None:
             data["token_type_hint"] = token_type
         url = self._requestor.reddit_url + const.REVOKE_TOKEN_PATH
-        await self._post(url, **data)
+        async with self._post(url, **data) as _:
+            pass  # The response is not used.
 
 
 class BaseAuthorizer:
@@ -152,8 +155,8 @@ def _clear_access_token(self):
     async def _request_token(self, **data: Any):
         url = self._authenticator._requestor.reddit_url + const.ACCESS_TOKEN_PATH
         pre_request_time = time.time()
-        response = await self._authenticator._post(url=url, **data)
-        payload = await response.json()
+        async with self._authenticator._post(url=url, **data) as response:
+            payload = await response.json()
         if "error" in payload:  # Why are these OKAY responses?
             raise OAuthException(
                 response, payload["error"], payload.get("error_description")

asyncprawcore/rate_limit.py

@@ -5,7 +5,8 @@
 import asyncio
 import logging
 import time
-from typing import TYPE_CHECKING, Any, Awaitable, Callable, Mapping
+from contextlib import asynccontextmanager
+from typing import TYPE_CHECKING, Any, AsyncContextManager, Awaitable, Callable, Mapping
 
 if TYPE_CHECKING:
     from aiohttp import ClientResponse
@@ -28,16 +29,15 @@ def __init__(self, *, window_size: int):
         self.used: int | None = None
         self.window_size: int = window_size
 
+    @asynccontextmanager
     async def call(
         self,
-        request_function: Callable[
-            [Any],
-            Awaitable[ClientResponse],
-        ],
+        # async context manager
+        request_function: Callable[..., AsyncContextManager[ClientResponse]],
         set_header_callback: Callable[[], Awaitable[dict[str, str]]],
         *args: Any,
         **kwargs: Any,
-    ) -> ClientResponse:
+    ) -> Callable[..., AsyncContextManager[ClientResponse]]:
         """Rate limit the call to ``request_function``.
 
         :param request_function: A function call that returns an HTTP response object.
@@ -49,9 +49,9 @@ async def call(
         """
         await self.delay()
         kwargs["headers"] = await set_header_callback()
-        response = await request_function(*args, **kwargs)
-        self.update(response.headers)
-        return response
+        async with request_function(*args, **kwargs) as response:
+            self.update(response.headers)
+            yield response
 
     async def delay(self):
         """Sleep for an amount of time to remain under the rate limit."""

asyncprawcore/requestor.py

@@ -2,14 +2,15 @@
 
 from __future__ import annotations
 
-import asyncio
-from typing import TYPE_CHECKING, Any
+from contextlib import asynccontextmanager
+from typing import TYPE_CHECKING, Any, AsyncContextManager, Callable
+from warnings import warn
 
 import aiohttp
 from aiohttp import ClientSession
 
 from .const import TIMEOUT
-from .exceptions import InvalidInvocation, RequestException
+from .exceptions import InvalidInvocation, RequestException, ResponseException
 
 if TYPE_CHECKING:
     from asyncio import AbstractEventLoop
@@ -23,7 +24,7 @@ class Requestor:
     def __getattr__(self, attribute: str) -> Any:  # pragma: no cover
         """Pass all undefined attributes to the ``_http`` attribute."""
         if attribute.startswith("__"):
-            raise AttributeError
+            raise AttributeError(attribute)
         return getattr(self._http, attribute)
 
     def __init__(
@@ -32,7 +33,7 @@ def __init__(
         oauth_url: str = "https://oauth.reddit.com",
         reddit_url: str = "https://www.reddit.com",
         session: ClientSession | None = None,
-        loop: AbstractEventLoop = None,
+        loop: AbstractEventLoop | None = None,
         timeout: float = TIMEOUT,
     ):
         """Create an instance of the Requestor class.
@@ -45,40 +46,77 @@ def __init__(
             ``"https://www.reddit.com"``).
         :param session: A session instance to handle requests, compatible with
             ``aiohttp.ClientSession()`` (default: ``None``).
+        :param loop: The event loop to run the requestor on (default: ``None``).
+
+            .. Deprecated:: 2.5.0
+
+                The ``loop`` argument is deprecated and will be ignored.
+
         :param timeout: How many seconds to wait for the server to send data before
             giving up (default: ``asyncprawcore.const.TIMEOUT``).
 
         """
         # Imported locally to avoid an import cycle, with __init__
         from . import __version__
 
+        if loop is not None:
+            msg = "The loop argument is deprecated and will be ignored."
+            warn(msg, DeprecationWarning, stacklevel=2)
+
         if user_agent is None or len(user_agent) < 7:
             msg = "user_agent is not descriptive"
             raise InvalidInvocation(msg)
 
-        self.loop = loop or asyncio.get_event_loop()
-        self._http = session or aiohttp.ClientSession(
-            loop=self.loop, timeout=aiohttp.ClientTimeout(total=None)
-        )
-        self._http._default_headers["User-Agent"] = (
-            f"{user_agent} asyncprawcore/{__version__}"
-        )
-
+        self.headers = {"User-Agent": f"{user_agent} asyncprawcore/{__version__}"}
         self.oauth_url = oauth_url
         self.reddit_url = reddit_url
         self.timeout = timeout
 
+        self._http = session
+        if self._http is not None and "User-Agent" not in self._http.headers:
+            # ensure user-agent is set
+            self._http.headers.update(self.headers)
+
+    async def _ensure_session(self):
+        """Ensure that the session is open."""
+        if self._http is None or self._http.closed:
+            self._http = aiohttp.ClientSession(
+                headers=self.headers,
+                timeout=aiohttp.ClientTimeout(total=None),
+            )
+
     async def close(self):
         """Call close on the underlying session."""
-        await self._http.close()
+        if self._http is not None and not self._http.closed:
+            await self._http.close()
 
+    @asynccontextmanager
     async def request(
         self, *args: Any, timeout: float | None = None, **kwargs: Any
-    ) -> ClientResponse:
-        """Issue the HTTP request capturing any errors that may occur."""
+    ) -> Callable[..., AsyncContextManager[ClientResponse]]:
+        """Issue the HTTP request capturing any errors that may occur.
+
+        :param args: Positional arguments to pass to ``aiohttp.ClientSession.request``.
+        :param timeout: How many seconds to wait for the server to send data before
+            giving up (default: ``None``).
+        :param kwargs: Keyword arguments to pass to ``aiohttp.ClientSession.request``.
+
+        :returns: The response from the request.
+
+        :raises: RequestException: If an error occurs while issuing the request.
+
+        """
         try:
-            return await self._http.request(
-                *args, timeout=timeout or self.timeout, **kwargs
-            )
+            await self._ensure_session()
+            kwargs_copy = kwargs.copy()
+            async with self._http.request(
+                *args,
+                headers={**self.headers, **kwargs_copy.pop("headers", {})},
+                timeout=timeout or self.timeout,
+                **kwargs_copy,
+            ) as request:
+                yield request
+        except ResponseException as exc:
+            raise exc
         except Exception as exc:  # noqa: BLE001
             raise RequestException(exc, args, kwargs) from None

asyncprawcore/sessions.py

@@ -7,9 +7,10 @@
 import random
 import time
 from abc import ABC, abstractmethod
+from contextlib import asynccontextmanager
 from copy import deepcopy
 from pprint import pformat
-from typing import TYPE_CHECKING, Any, BinaryIO, TextIO
+from typing import TYPE_CHECKING, Any, AsyncContextManager, BinaryIO, Callable, TextIO
 from urllib.parse import urljoin
 
 from aiohttp.web import HTTPRequestTimeout
@@ -179,6 +180,7 @@ async def _do_retry(
             # noqa: E501
         )
 
+    @asynccontextmanager
     async def _make_request(
         self,
         data: list[tuple[str, Any]],
@@ -188,9 +190,11 @@ async def _make_request(
         retry_strategy_state: FiniteRetryStrategy,
         timeout: float,
         url: str,
-    ) -> tuple[ClientResponse, None] | tuple[None, Exception]:
+    ) -> Callable[
+        ..., AsyncContextManager[tuple[ClientResponse | None, Exception | None]]
+    ]:
         try:
-            response = await self._rate_limiter.call(
+            async with self._rate_limiter.call(
                 self._requestor.request,
                 self._set_header_callback,
                 method,
@@ -200,17 +204,17 @@ async def _make_request(
                 json=json,
                 params=params,
                 timeout=timeout,
-            )
-            log.debug(
-                "Response: %s (%s bytes) (rst-%s:rem-%s:used-%s ratelimit) at %s",
-                response.status,
-                response.headers.get("content-length"),
-                response.headers.get("x-ratelimit-reset"),
-                response.headers.get("x-ratelimit-remaining"),
-                response.headers.get("x-ratelimit-used"),
-                time.time(),
-            )
-            return response, None
+            ) as response:
+                log.debug(
+                    "Response: %s (%s bytes) (rst-%s:rem-%s:used-%s ratelimit) at %s",
+                    response.status,
+                    response.headers.get("content-length"),
+                    response.headers.get("x-ratelimit-reset"),
+                    response.headers.get("x-ratelimit-remaining"),
+                    response.headers.get("x-ratelimit-used"),
+                    time.time(),
+                )
+                yield response, None
         except RequestException as exception:
             if (
                 not retry_strategy_state.should_retry_on_failure()
@@ -219,7 +223,7 @@ async def _make_request(
                 )
             ):
                 raise
-            return None, exception.original_exception
+            yield None, exception.original_exception
 
     def _preprocess_data(
         self,
@@ -284,54 +288,54 @@ async def _request_with_retries(
 
         await retry_strategy_state.sleep()
         self._log_request(data, method, params, url)
-        response, saved_exception = await self._make_request(
+        async with self._make_request(
             data,
             json,
             method,
             params,
             retry_strategy_state,
             timeout,
             url,
-        )
-
-        do_retry = False
-        if response is not None and response.status == codes["unauthorized"]:
-            self._authorizer._clear_access_token()
-            if hasattr(self._authorizer, "refresh"):
-                do_retry = True
-
-        if retry_strategy_state.should_retry_on_failure() and (
-            do_retry or response is None or response.status in self.RETRY_STATUSES
-        ):
-            return await self._do_retry(
-                data,
-                json,
-                method,
-                params,
-                response,
-                retry_strategy_state,
-                saved_exception,
-                timeout,
-                url,
-            )
-        if response.status in self.STATUS_EXCEPTIONS:
-            if response.status == codes["media_type"]:
-                # since exception class needs response.json
-                raise self.STATUS_EXCEPTIONS[response.status](
-                    response, await response.json()
+        ) as (response, saved_exception):
+            do_retry = False
+            if response is not None and response.status == codes["unauthorized"]:
+                # noinspection PyProtectedMember
+                self._authorizer._clear_access_token()
+                if hasattr(self._authorizer, "refresh"):
+                    do_retry = True
+
+            if retry_strategy_state.should_retry_on_failure() and (
+                do_retry or response is None or response.status in self.RETRY_STATUSES
+            ):
+                return await self._do_retry(
+                    data,
+                    json,
+                    method,
+                    params,
+                    response,
+                    retry_strategy_state,
+                    saved_exception,
+                    timeout,
+                    url,
                 )
-            raise self.STATUS_EXCEPTIONS[response.status](response)
-        if response.status == codes["no_content"]:
-            return None
-        assert (
-            response.status in self.SUCCESS_STATUSES
-        ), f"Unexpected status code: {response.status}"
-        if response.headers.get("content-length") == "0":
-            return ""
-        try:
-            return await response.json()
-        except ValueError:
-            raise BadJSON(response) from None
+            if response.status in self.STATUS_EXCEPTIONS:
+                if response.status == codes["media_type"]:
+                    # since exception class needs response.json
+                    raise self.STATUS_EXCEPTIONS[response.status](
+                        response, await response.json()
+                    )
+                raise self.STATUS_EXCEPTIONS[response.status](response)
+            if response.status == codes["no_content"]:
+                return None
+            assert (
+                response.status in self.SUCCESS_STATUSES
+            ), f"Unexpected status code: {response.status}"
+            if response.headers.get("content-length") == "0":
+                return ""
+            try:
+                return await response.json()
+            except ValueError:
+                raise BadJSON(response) from None
 
     async def _set_header_callback(self) -> dict[str, str]:
         if not self._authorizer.is_valid() and hasattr(self._authorizer, "refresh"):