pyhOn/pyhon/connection/handler/auth.py

37 lines
1.1 KiB
Python

import logging
from collections.abc import AsyncIterator
from contextlib import asynccontextmanager
from typing import Optional, Callable, List, Tuple
import aiohttp
from pyhon import const
from pyhon.connection.handler.base import ConnectionHandler
_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(
self, method: Callable, *args, loop: int = 0, **kwargs
) -> AsyncIterator:
kwargs["headers"] = kwargs.pop("headers", {}) | self._HEADERS
async with method(*args, **kwargs) as response:
self._called_urls.append((response.status, response.request_info.url))
yield response