PamGram/modules/apihelper/client/base/httpxrequest.py
2023-04-26 16:48:05 +08:00

45 lines
1.2 KiB
Python

from contextlib import AbstractAsyncContextManager
from types import TracebackType
from typing import Optional, Type
import httpx
__all__ = ("HTTPXRequest",)
from core.config import config
timeout = httpx.Timeout(
timeout=config.timeout,
read=config.read_timeout,
write=config.write_timeout,
connect=config.connect_timeout,
pool=config.pool_timeout,
)
class HTTPXRequest(AbstractAsyncContextManager):
def __init__(self, *args, headers=None, **kwargs):
self._client = httpx.AsyncClient(headers=headers, *args, **kwargs)
async def __aenter__(self):
try:
await self.initialize()
return self
except Exception as exc:
await self.shutdown()
raise exc
async def __aexit__(
self, exc_type: Optional[Type[BaseException]], exc_val: Optional[BaseException], exc_tb: Optional[TracebackType]
) -> None:
await self.initialize()
async def initialize(self):
if self._client.is_closed:
self._client = httpx.AsyncClient(timeout=timeout)
async def shutdown(self):
if self._client.is_closed:
return
await self._client.aclose()