pyhOn/pyhon/connection/handler/auth.py

39 lines
1.3 KiB
Python
Raw Permalink Normal View History

2023-04-15 14:22:04 +02:00
import logging
from collections.abc import AsyncIterator
from contextlib import asynccontextmanager
2023-06-28 20:25:52 +02:00
from typing import Optional, List, Tuple, Any, Dict
2023-04-15 14:22:04 +02:00
import aiohttp
2023-06-28 20:25:52 +02:00
from yarl import URL
2023-04-15 14:22:04 +02:00
from pyhon import const
from pyhon.connection.handler.base import ConnectionHandler
2023-06-28 19:02:11 +02:00
from pyhon.typedefs import Callback
2023-04-15 14:22:04 +02:00
_LOGGER = logging.getLogger(__name__)
class HonAuthConnectionHandler(ConnectionHandler):
_HEADERS = {"user-agent": const.USER_AGENT}
def __init__(self, session: Optional[aiohttp.ClientSession] = None) -> None:
super().__init__(session)
self._called_urls: List[Tuple[int, str]] = []
@property
def called_urls(self) -> List[Tuple[int, str]]:
return self._called_urls
@called_urls.setter
def called_urls(self, called_urls: List[Tuple[int, str]]) -> None:
self._called_urls = called_urls
@asynccontextmanager
async def _intercept(
2023-06-28 20:25:52 +02:00
self, method: Callback, url: str | URL, *args: Any, **kwargs: Dict[str, Any]
2023-06-28 19:02:11 +02:00
) -> AsyncIterator[aiohttp.ClientResponse]:
2023-04-15 14:22:04 +02:00
kwargs["headers"] = kwargs.pop("headers", {}) | self._HEADERS
2023-06-28 20:25:52 +02:00
async with method(url, *args, **kwargs) as response:
2023-06-28 19:02:11 +02:00
self._called_urls.append((response.status, str(response.request_info.url)))
2023-04-15 14:22:04 +02:00
yield response