feat: add dynamic PyPI mirror selection to bootstrap script

Author: HuanCheng65

bootstrap.py

@@ -2,16 +2,27 @@
 import sys
 import subprocess
 import venv
+import time
+import urllib.request
 from pathlib import Path
+from concurrent.futures import ThreadPoolExecutor, as_completed
 
 # ================= 配置 =================
 VENV_DIR_NAME = ".venv"
 REQUIREMENTS_FILE = "requirements.txt"
-# 这里的包名必须是 import 时使用的名字 (例如 "tomli" 而不是 "tomli>=2.0")
-REQUIRED_IMPORTS = ["rich", "tomli"] 
-REQUIRED_PACKAGES = ["rich", "tomli"] # 如果需要安装，传给 pip 的包名
+REQUIRED_IMPORTS = ["rich", "tomli"]
+REQUIRED_PACKAGES = ["rich", "tomli"]
+
+# 定义待测速的 PyPI 源列表
+PYPI_MIRRORS = {
+    "Official": "https://pypi.org/simple",
+    "Tsinghua": "https://pypi.tuna.tsinghua.edu.cn/simple",
+    "Aliyun": "https://mirrors.aliyun.com/pypi/simple",
+    "Tencent": "https://mirrors.cloud.tencent.com/pypi/simple",
+}
 # =======================================
 
+
 def get_venv_paths(root_dir: Path):
     """获取虚拟环境的路径"""
     venv_dir = root_dir / VENV_DIR_NAME
@@ -23,62 +34,126 @@ def get_venv_paths(root_dir: Path):
         pip_executable = venv_dir / "bin" / "pip"
     return venv_dir, python_executable, pip_executable
 
+
 def check_venv_integrity(python_executable: Path) -> bool:
-    """
-    检查 venv 里的 Python 是否已经安装了必要的包。
-    通过调用 venv python 执行尝试导入的脚本来验证。
-    """
+    """检查 venv 里的 Python 是否已经安装了必要的包。"""
     if not python_executable.exists():
         return False
-    
-    # 构建检查脚本: "import rich; import tomli;"
+
     check_script = "; ".join([f"import {pkg}" for pkg in REQUIRED_IMPORTS])
-    
     try:
-        # 使用 -c 执行导入检查，stdout/stderr 扔掉，只看返回码
         subprocess.run(
             [str(python_executable), "-c", check_script],
             stdout=subprocess.DEVNULL,
             stderr=subprocess.DEVNULL,
-            check=True
+            check=True,
         )
         return True
     except subprocess.CalledProcessError:
         return False
 
+
 def create_venv_if_missing(venv_dir: Path):
     if not venv_dir.exists():
         print(f"[Bootstrap] Creating virtual environment at: {venv_dir}...", flush=True)
         venv.create(venv_dir, with_pip=True)
 
+
+def test_mirror_latency(name, url, timeout=2):
+    """测试单个镜像源的延迟"""
+    start_time = time.time()
+    try:
+        # 使用 HEAD 请求或者极小的 GET 请求来测试连接
+        # 标准库 urllib 发送请求
+        req = urllib.request.Request(url, method="HEAD")
+        with urllib.request.urlopen(req, timeout=timeout) as response:
+            if response.status == 200:
+                latency = (time.time() - start_time) * 1000  # ms
+                return name, url, latency
+    except Exception:
+        pass
+    return name, url, float("inf")
+
+
+def get_fastest_mirror():
+    """并发测试所有源，返回速度最快的源 URL"""
+    print("[Bootstrap] Selecting the fastest PyPI mirror...", flush=True)
+    
+    fastest_url = None
+    min_latency = float("inf")
+    best_name = "Unknown"
+
+    # 使用线程池并发测速，避免串行等待
+    with ThreadPoolExecutor(max_workers=len(PYPI_MIRRORS)) as executor:
+        futures = [
+            executor.submit(test_mirror_latency, name, url) 
+            for name, url in PYPI_MIRRORS.items()
+        ]
+        
+        for future in as_completed(futures):
+            name, url, latency = future.result()
+            if latency < float("inf"):
+                # print(f"  - {name}: {latency:.2f} ms") # 调试时可打开
+                if latency < min_latency:
+                    min_latency = latency
+                    fastest_url = url
+                    best_name = name
+
+    if fastest_url:
+        print(f"[Bootstrap] Selected {best_name} (Latency: {min_latency:.0f}ms)", flush=True)
+        return fastest_url
+    else:
+        print("[Bootstrap] Warning: All mirrors check failed. Fallback to Official.", flush=True)
+        return PYPI_MIRRORS["Official"]
+
+
 def install_dependencies(root_dir: Path, pip_executable: Path):
     req_path = root_dir / REQUIREMENTS_FILE
-    base_cmd = [str(pip_executable), "install", "--disable-pip-version-check", "-i", "https://pypi.tuna.tsinghua.edu.cn/simple"]
 
-    print(f"[Bootstrap] Installing dependencies...", flush=True)
+    # === 获取最快源 ===
+    mirror_url = get_fastest_mirror()
+    # ================
+
+    base_cmd = [
+        str(pip_executable),
+        "install",
+        "--disable-pip-version-check",
+        "-i",
+        mirror_url, # 使用动态获取的源
+    ]
+
+    # 如果选中的不是官方源，通常需要信任该 host，避免 pip 报 SSL 警告
+    if "pypi.org" not in mirror_url:
+        from urllib.parse import urlparse
+        hostname = urlparse(mirror_url).hostname
+        base_cmd.extend(["--trusted-host", hostname])
+
+    print("[Bootstrap] Installing dependencies...", flush=True)
     try:
         if req_path.exists():
             subprocess.run(base_cmd + ["-r", str(req_path)], check=True)
         else:
             subprocess.run(base_cmd + REQUIRED_PACKAGES, check=True)
     except subprocess.CalledProcessError:
-        print(f"[Bootstrap] Error: Failed to install dependencies.", file=sys.stderr)
+        print("[Bootstrap] Error: Failed to install dependencies.", file=sys.stderr)
         sys.exit(1)
 
+
 def restart_in_venv(python_executable: Path):
     """重启当前脚本到 venv 环境"""
     env = os.environ.copy()
     env["GRADER_BOOTSTRAPPED"] = "1"
     args = [str(python_executable)] + sys.argv
-    
+
     if sys.platform == "win32":
         subprocess.run(args, env=env, check=True)
         sys.exit(0)
     else:
         os.execv(str(python_executable), args)
 
+
 def initialize():
-    # 1. Fast Path: 如果当前环境(无论是否venv)已经能用，直接通过
+    # 1. Fast Path
     try:
         for pkg in REQUIRED_IMPORTS:
             __import__(pkg)
@@ -90,21 +165,23 @@ def initialize():
     venv_dir, venv_python, venv_pip = get_venv_paths(root_dir)
     is_in_venv = os.environ.get("GRADER_BOOTSTRAPPED") == "1"
 
-    # 2. 如果已经在 venv 里但还是缺包，说明环境坏了，强制修复
+    # 2. 修复损坏环境
     if is_in_venv:
         print("[Bootstrap] In venv but dependencies missing. Repairing...", flush=True)
         install_dependencies(root_dir, venv_pip)
         return
 
-    # 3. 如果在系统环境 (用户没激活 venv 直接跑脚本)
-    #    逻辑：先检查 venv 是否存在且完好，如果是，直接重启进去，不要跑 install
-    
+    # 3. 创建与首次安装
     create_venv_if_missing(venv_dir)
 
-    # === 关键优化点 ===
-    # 只有当 venv 里的 python 无法导入指定包时，才运行 pip install
     if not check_venv_integrity(venv_python):
         install_dependencies(root_dir, venv_pip)
-    # =================
-    
+
     restart_in_venv(venv_python)
+
+# 示例用法，实际使用时这里可能是 main()
+if __name__ == "__main__":
+    initialize()
+    # 下面写你的正常业务逻辑
+    import rich
+    print("[Success] Environment is ready!")

grader.py

@@ -2,7 +2,6 @@
 import json
 import os
 import re
-import shutil
 import subprocess
 import sys
 import time
@@ -19,7 +18,7 @@
 # -----------------------------------------
 
 import tomli  # noqa: E402
-from rich.console import Console  # noqa: E402
+from rich.console import Console, RenderableType  # noqa: E402
 from rich.panel import Panel  # noqa: E402
 from rich.progress import Progress, SpinnerColumn, Task, TextColumn  # noqa: E402
 from rich.table import Table  # noqa: E402
@@ -196,7 +195,7 @@ def _resolve_path(self, path: str, test_dir: Path) -> str:
 class StatusSpinnerColumn(SpinnerColumn):
     """Spinner that swaps to a custom status icon when provided."""
 
-    def render(self, task: Task) -> Text:
+    def render(self, task: Task) -> RenderableType:
         icon = task.fields.get("status_icon")
         if icon:
             return Text.from_markup(icon)
@@ -605,7 +604,7 @@ def _execute_single_step(
 
         return self._create_success_result(test, step, score, start_time)
 
-    def _resolve_relative_path(self, path: str, cwd: Path = os.getcwd()) -> str:
+    def _resolve_relative_path(self, path: str, cwd: Path = Path.cwd()) -> str:
         result = path
         if isinstance(path, Path):
             result = str(path.resolve())
@@ -623,7 +622,7 @@ def _resolve_relative_path(self, path: str, cwd: Path = os.getcwd()) -> str:
 
         return result
 
-    def _resolve_path(self, path: str, test_dir: Path, cwd: Path = os.getcwd()) -> str:
+    def _resolve_path(self, path: str, test_dir: Path, cwd: Path = Path.cwd()) -> str:
         build_dir = test_dir / "build"
         build_dir.mkdir(exist_ok=True)
 
@@ -1025,7 +1024,7 @@ def _generate_launch_config(
         else:
             raise ValueError(f"Unsupported debug type: {debug_type}")
 
-    def _generate_tasks_config(self, test_case: TestCase) -> Dict[str, Any]:
+    def _generate_tasks_config(self, test_case: TestCase) -> List[Dict[str, Any]]:
         """Generate tasks configuration for building the test case"""
         return [
             {
@@ -1092,7 +1091,7 @@ def _write_or_merge_json(
         with open(file_path, "w") as f:
             json.dump(config_to_write, f, indent=4)
 
-    def _resolve_relative_path(self, path: str, cwd: Path = os.getcwd()) -> str:
+    def _resolve_relative_path(self, path: str, cwd: Path = Path.cwd()) -> str:
         result = path
         if isinstance(path, Path):
             result = str(path.resolve())
@@ -1110,7 +1109,7 @@ def _resolve_relative_path(self, path: str, cwd: Path = os.getcwd()) -> str:
 
         return result
 
-    def _resolve_path(self, path: str, test_dir: Path, cwd: Path = os.getcwd()) -> str:
+    def _resolve_path(self, path: str, test_dir: Path, cwd: Path = Path.cwd()) -> str:
         build_dir = test_dir / "build"
         build_dir.mkdir(exist_ok=True)