NeteaseUrl/core/sync_http.py

# -*- coding: utf-8 -*-
"""同步 HTTP 门面（底层 aiohttp）——彻底移除 requests 依赖。

为什么需要它：部分代码天然是同步的（worker 线程里的流式下载、纯函数里的短链展开/
IP 归属地查询）。本模块提供 requests 风格的同步 API，但底层复用 aiohttp 连接池，
不再引入 requests。

流式读取的实现：在后台事件循环起一个读取协程，把块推入有界队列（背压），
同步侧从队列消费——每次搬运只过一次队列，不做「每块一次跨线程调度」。
"""
import asyncio
import queue
import threading
from collections.abc import Iterator, Mapping
from typing import Any

from core import async_bridge
from core.http_client import client as _client

_DEFAULT_TIMEOUT = 30


class RequestException(Exception):
    """请求异常基类（对齐 requests 命名，便于迁移）。"""


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

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


class SyncResponse:
    """requests.Response 的同步子集（底层 aiohttp）。"""

    def __init__(self, status: int, headers: Mapping[str, str], content: bytes | None,
                 url: str, stream_state: dict | None = None):
        self.status_code = status
        self.headers = dict(headers)
        self.url = url
        self._content = content
        self._stream = stream_state

    @property
    def content(self) -> bytes:
        if self._content is None:
            self._content = b''.join(self.iter_content(64 * 1024))
        return self._content

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

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

    def raise_for_status(self) -> None:
        if self.status_code >= 400:
            raise HTTPError(f'HTTP {self.status_code}', self)

    def iter_content(self, chunk_size: int = 64 * 1024) -> Iterator[bytes]:
        """流式读取响应体（大小写与 requests 一致：chunk_size 为字节数）。"""
        if self._content is not None:
            for i in range(0, len(self._content), chunk_size):
                yield self._content[i:i + chunk_size]
            return
        st = self._stream
        assert st is not None
        q: queue.Queue = st['queue']
        while True:
            item = q.get()
            if item is None:                      # 正常结束
                break
            if isinstance(item, BaseException):   # 读取出错
                raise item
            yield item

    def close(self) -> None:
        st = self._stream
        if st and st.get('cancel') is not None:
            try:
                async_bridge.run_sync(_cancel(st))
            except Exception:
                pass

    def __enter__(self) -> 'SyncResponse':
        return self

    def __exit__(self, *exc) -> bool:
        self.close()
        return False


async def _cancel(st: dict) -> None:
    st['cancelled'] = True
    resp = st.get('resp')
    if resp is not None:
        resp.close()


def _stream_reader(st: dict, chunk_size: int) -> None:
    """在后台事件循环读块 → 有界队列（背压），同步侧消费。"""
    async def _run():
        try:
            async for chunk in st['resp'].content.iter_chunked(chunk_size):
                if st.get('cancelled'):
                    break
                while True:  # 背压：队列满则稍等，避免内存无界
                    try:
                        st['queue'].put(chunk, timeout=0.5)
                        break
                    except queue.Full:
                        if st.get('cancelled'):
                            return
        except Exception as e:
            st['queue'].put(e)
        finally:
            try:
                st['resp'].close()
            except Exception:
                pass
            st['queue'].put(None)

    asyncio.run_coroutine_threadsafe(_run(), async_bridge.get_loop())


class Session:
    """requests.Session 的同步子集（连接池由底层 aiohttp 统一维护）。"""

    trust_env = False

    def __init__(self, **kw) -> None:
        # headers：会话级默认头（与 requests.Session().headers.update 等价）
        self.headers: dict[str, str] = dict(kw.pop('headers', {}))

    def _merge(self, kw: dict) -> dict:
        h = dict(self.headers)
        h.update(kw.pop('headers', None) or {})
        if h:
            kw['headers'] = h
        return kw

    def get(self, url: str, **kw) -> SyncResponse:
        return _request('GET', url, **self._merge(kw))

    def post(self, url: str, **kw) -> SyncResponse:
        return _request('POST', url, **self._merge(kw))

    def close(self) -> None:
        pass


def _norm_timeout(timeout):
    """requests 允许 (connect, read) 元组；aiohttp 只接受单一总时长——取较大者。"""
    if isinstance(timeout, (tuple, list)):
        vals = [float(v) for v in timeout if isinstance(v, (int, float))]
        return max(vals) if vals else _DEFAULT_TIMEOUT
    return timeout


def _request(method: str, url: str, *, headers: Mapping[str, str] | None = None,
             cookies: Mapping[str, str] | None = None, params: Mapping[str, Any] | None = None,
             data: Any = None, json: Any = None, timeout: float | None = _DEFAULT_TIMEOUT,
             allow_redirects: bool = True, stream: bool = False) -> SyncResponse:
    timeout = _norm_timeout(timeout)
    if stream:
        return _open_stream(method, url, headers=headers, cookies=cookies, params=params,
                            data=data, json=json, timeout=timeout,
                            allow_redirects=allow_redirects)
    resp = async_bridge.run_sync(_client.request(
        method, url, headers=headers, cookies=cookies, params=params, data=data,
        json_body=json, timeout=timeout, allow_redirects=allow_redirects))
    return SyncResponse(resp.status, resp.headers, resp.text.encode('utf-8'),
                        resp.url)


def _open_stream(method: str, url: str, **kw) -> SyncResponse:
    timeout = _norm_timeout(kw.pop('timeout', _DEFAULT_TIMEOUT))
    # AsyncHTTPClient.open 只接受 headers/cookies/timeout/allow_redirects；
    # 流式场景不支持 params/data/json（调用方直接把查询串拼进 URL）
    for unsupported in ('params', 'data', 'json'):
        kw.pop(unsupported, None)
    resp = async_bridge.run_sync(_client.open(method, url, timeout=timeout, **kw))
    st = {'resp': resp, 'queue': queue.Queue(maxsize=8), 'cancel': object(), 'cancelled': False}
    th = threading.Thread(target=_stream_reader, args=(st, 64 * 1024), daemon=True)
    th.start()
    return SyncResponse(resp.status, resp.headers, None, str(resp.url), stream_state=st)


# 模块级便捷函数（等价于共享 Session 的 get/post）
_default_session = Session()


def get(url: str, **kw) -> SyncResponse:
    return _default_session.get(url, **kw)


def post(url: str, **kw) -> SyncResponse:
    return _default_session.post(url, **kw)