tushare

http_timeout.py

@@ -0,0 +1,77 @@
+"""
+统一的外部请求超时控制工具
+
+- install_default_requests_timeout: 给所有基于 requests 的外部请求
+  （akshare、tushare、pywencai 等）注入默认超时，调用方未指定 timeout 时生效。
+- call_with_timeout: 在独立线程中执行阻塞调用，超过时限立即放弃等待，
+  防止个别环节长时间卡死整个程序。
+"""
+
+import os
+import threading
+
+# 默认连接超时 / 读取超时（单位：秒），可通过环境变量覆盖
+DEFAULT_CONNECT_TIMEOUT = float(os.getenv("HTTP_CONNECT_TIMEOUT", "6"))
+DEFAULT_READ_TIMEOUT = float(os.getenv("HTTP_READ_TIMEOUT", "20"))
+DEFAULT_CALL_TIMEOUT = float(os.getenv("HTTP_CALL_TIMEOUT", "30"))
+
+_installed = False
+_install_lock = threading.Lock()
+
+
+def install_default_requests_timeout(connect_timeout=None, read_timeout=None):
+    """给 requests 库注入默认超时（仅当调用方未显式指定 timeout 时生效）。"""
+    global _installed
+    with _install_lock:
+        if _installed:
+            return
+
+        connect = DEFAULT_CONNECT_TIMEOUT if connect_timeout is None else connect_timeout
+        read = DEFAULT_READ_TIMEOUT if read_timeout is None else read_timeout
+
+        try:
+            import requests.sessions
+            original_request = requests.sessions.Session.request
+
+            def request_with_timeout(self, method, url, **kwargs):
+                if kwargs.get("timeout") is None:
+                    kwargs["timeout"] = (connect, read)
+                return original_request(self, method, url, **kwargs)
+
+            requests.sessions.Session.request = request_with_timeout
+            _installed = True
+            print(f"✅ 已为外部请求注入默认超时（连接 {connect}s / 读取 {read}s）")
+        except Exception as e:
+            print(f"⚠️ 注入 requests 默认超时失败: {e}")
+
+
+def call_with_timeout(func, timeout=None, *args, **kwargs):
+    """在线程中执行函数，超过 timeout 秒则放弃等待并抛出 TimeoutError。
+
+    注意：超时后线程会继续在后台运行直至结束（守护线程），不会影响主程序。
+    """
+    if timeout is None:
+        timeout = DEFAULT_CALL_TIMEOUT
+    if timeout <= 0:
+        return func(*args, **kwargs)
+
+    result_box = {}
+
+    def _run():
+        try:
+            result_box["ok"] = True
+            result_box["value"] = func(*args, **kwargs)
+        except BaseException as e:  # noqa: BLE001
+            result_box["ok"] = False
+            result_box["error"] = e
+
+    worker = threading.Thread(target=_run, daemon=True)
+    worker.start()
+    worker.join(timeout)
+
+    if worker.is_alive():
+        func_name = getattr(func, "__name__", "function")
+        raise TimeoutError(f"调用 {func_name} 超过 {timeout}s 未返回，已放弃等待")
+    if result_box.get("ok"):
+        return result_box["value"]
+    raise result_box.get("error", RuntimeError("未知错误"))