NeteaseUrl/core/aiohttp_compat.py

# -*- coding: utf-8 -*-
"""requests 兼容的异步 HTTP 垫片（基于 aiohttp）。

目的：让既有大量 `requests.post/get + response.json()/raise_for_status()` 风格的
代码**以最小改动**迁移到异步——把 `import requests` 换成
`from core import aiohttp_compat as requests`，加 `await`、加 `async def` 即可，
其余 `.text/.json()/.status_code/.raise_for_status()` 语义保持不变。

与 requests 的语义对齐点（很关键，迁移期不能变）：
- 请求本身**不**因 4xx/5xx 抛异常；只有显式 `raise_for_status()` 才抛 `HTTPError`
  （上游风控 429/460 的重试/不重试判断依赖调用方读取 `status_code`，见 music_api.HTTPClient）。
- 网络异常/超时抛 `RequestException`（含 `Timeout`）。
- `Response` 一次性读全 body（与 requests 一致），暴露 text/json/content/headers/cookies/url。

连接按事件循环隔离（见 core.http_client 的同样设计）。
"""
import asyncio
import json as _json
from collections.abc import Mapping
from typing import Any

import aiohttp

DEFAULT_TIMEOUT = 30
DEFAULT_MAX_BYTES = 8 * 1024 * 1024

# ---------- 异常层次（对齐 requests 命名，便于原样 except） ----------


class RequestException(Exception):
    """请求异常基类（网络/协议层）。"""


class HTTPError(RequestException):
    """raise_for_status() 抛出的 HTTP 状态异常。"""

    def __init__(self, message: str, response: 'Response | None' = None):
        super().__init__(message)
        self.response = response


class Timeout(RequestException):
    """请求超时。"""


class ConnectionError(RequestException):
    """连接错误。"""


class Response:
    """requests.Response 的异步子集：body 已读全，属性同步访问。"""

    def __init__(self, status_code: int, url: str, content: bytes,
                 headers: Mapping[str, str], cookies: Mapping[str, str], reason: str = ''):
        self.status_code = status_code
        self.url = url
        self.content = content
        self.headers = dict(headers)
        self.cookies = dict(cookies)
        self.reason = reason

    @property
    def text(self) -> str:
        return self.content.decode('utf-8', errors='replace')

    def json(self, **kw) -> Any:
        return _json.loads(self.text, **kw)

    def raise_for_status(self) -> None:
        if 400 <= self.status_code:
            raise HTTPError(f"{self.status_code} {self.reason}".strip(), self)


# ---------- 会话（按 loop 隔离的 aiohttp session） ----------


class _SessionState:
    def __init__(self):
        self._sessions: dict[int, aiohttp.ClientSession] = {}
        self._lock: asyncio.Lock | None = None

    def _get_lock(self) -> asyncio.Lock:
        if self._lock is None:
            self._lock = asyncio.Lock()
        return self._lock

    async def session(self) -> aiohttp.ClientSession:
        key = id(asyncio.get_running_loop())
        sess = self._sessions.get(key)
        if sess is not None and not sess.closed:
            return sess
        async with self._get_lock():
            sess = self._sessions.get(key)
            if sess is None or sess.closed:
                sess = aiohttp.ClientSession(
                    connector=aiohttp.TCPConnector(limit=100, limit_per_host=32, ttl_dns_cache=300),
                )
                self._sessions[key] = sess
        return sess

    async def close_all(self) -> None:
        for sess in self._sessions.values():
            if not sess.closed:
                await sess.close()
        self._sessions.clear()


_state = _SessionState()


async def close() -> None:
    """关闭所有连接（应用关闭时调用）。"""
    await _state.close_all()


async def request(
    method: str,
    url: str,
    *,
    params: Any = None,
    data: Any = None,
    json: Any = None,
    headers: Mapping[str, str] | None = None,
    cookies: Mapping[str, str] | None = None,
    timeout: float | None = DEFAULT_TIMEOUT,
    allow_redirects: bool = True,
    max_bytes: int = DEFAULT_MAX_BYTES,
) -> Response:
    """发起单次异步请求，返回 Response（不因 4xx/5xx 抛异常，同 requests）。"""
    sess = await _state.session()
    req_timeout = aiohttp.ClientTimeout(total=timeout) if timeout else None
    try:
        async with sess.request(
            method, url,
            params=params, data=data, json=json,
            headers=dict(headers) if headers else None,
            cookies=dict(cookies) if cookies else None,
            timeout=req_timeout, allow_redirects=allow_redirects,
        ) as resp:
            body = await resp.read()
            if len(body) > max_bytes:
                raise RequestException(f"响应体超过上限 {max_bytes} 字节")
            return Response(
                status_code=resp.status,
                url=str(resp.url),
                content=body,
                headers=dict(resp.headers),
                cookies={k: v.value for k, v in resp.cookies.items()},
                reason=resp.reason or '',
            )
    except aiohttp.ClientError as e:
        raise ConnectionError(str(e)) from e
    except TimeoutError as e:
        raise Timeout('请求超时') from e


async def get(url: str, **kw) -> Response:
    return await request('GET', url, **kw)


async def post(url: str, **kw) -> Response:
    return await request('POST', url, **kw)