111 lines
3.7 KiB
Python
111 lines
3.7 KiB
Python
import subprocess
|
|
import time
|
|
import requests
|
|
import os
|
|
import concurrent.futures
|
|
from rich.console import Console
|
|
from rich.table import Table
|
|
from rich.progress import Progress, SpinnerColumn, TextColumn
|
|
from rich.panel import Panel
|
|
import sys
|
|
|
|
MIRRORS = {
|
|
"Tsinghua Mirror": "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple",
|
|
"PyPI Official": "https://pypi.org/simple"
|
|
}
|
|
|
|
console = Console()
|
|
|
|
FAST_THRESHOLD = 3000 # ms
|
|
SLOW_THRESHOLD = 5000 # ms
|
|
|
|
def get_optimal_thread_count():
|
|
try:
|
|
cpu_count = os.cpu_count()
|
|
return max(cpu_count - 1, 1)
|
|
except:
|
|
return 2
|
|
|
|
def test_mirror_speed(name, url):
|
|
try:
|
|
start_time = time.time()
|
|
response = requests.get(url, timeout=5)
|
|
end_time = time.time()
|
|
if response.status_code == 200:
|
|
speed = (end_time - start_time) * 1000
|
|
return name, speed
|
|
else:
|
|
return name, float('inf')
|
|
except requests.RequestException:
|
|
return name, float('inf')
|
|
|
|
def set_pip_mirror(url):
|
|
try:
|
|
subprocess.run([sys.executable, "-m", "pip", "config", "set", "global.index-url", url],
|
|
check=True,
|
|
capture_output=True)
|
|
return True
|
|
except subprocess.CalledProcessError as e:
|
|
print(f"Failed to set pip mirror: {e}")
|
|
return False
|
|
|
|
def get_current_pip_mirror():
|
|
try:
|
|
result = subprocess.run([sys.executable, "-m", "pip", "config", "get", "global.index-url"],
|
|
capture_output=True, text=True, check=True)
|
|
return result.stdout.strip()
|
|
except subprocess.CalledProcessError:
|
|
return None
|
|
|
|
def main():
|
|
console.print(Panel.fit("🚀 PyPI Mirror Speed Test", style="bold cyan"))
|
|
|
|
# Test all mirrors simultaneously
|
|
speeds = {}
|
|
with Progress(
|
|
SpinnerColumn(),
|
|
TextColumn("[cyan]Testing mirrors...[/cyan]"),
|
|
) as progress:
|
|
progress.add_task("", total=None) # Indeterminate spinner
|
|
|
|
with concurrent.futures.ThreadPoolExecutor(max_workers=get_optimal_thread_count()) as executor:
|
|
future_to_mirror = {executor.submit(test_mirror_speed, name, url): name
|
|
for name, url in MIRRORS.items()}
|
|
|
|
for future in concurrent.futures.as_completed(future_to_mirror):
|
|
name = future_to_mirror[future]
|
|
try:
|
|
name, speed = future.result()
|
|
if speed != float('inf'):
|
|
speeds[name] = speed
|
|
except Exception as exc:
|
|
print(f'{name} generated an exception: {exc}')
|
|
|
|
# Results display
|
|
table = Table(show_header=False)
|
|
table.add_column(style="cyan")
|
|
table.add_column(justify="right", style="magenta")
|
|
|
|
for name, speed in sorted(speeds.items(), key=lambda x: x[1]):
|
|
table.add_row(name, f"{speed:.0f}ms")
|
|
|
|
console.print(table)
|
|
|
|
if speeds:
|
|
fastest_mirror = min(speeds, key=speeds.get)
|
|
fastest_url = MIRRORS[fastest_mirror]
|
|
|
|
if set_pip_mirror(fastest_url):
|
|
current_mirror = get_current_pip_mirror()
|
|
if current_mirror == fastest_url:
|
|
console.print(f"✅ Switched to {fastest_mirror}\n🔗 {fastest_url}", style="green")
|
|
else:
|
|
console.print(f"❌ Switch failed\nExpected: {fastest_url}\nCurrent: {current_mirror}\n💡 Try running with admin privileges", style="red")
|
|
else:
|
|
console.print(f"❌ Failed to switch mirror\n💡 Check permissions and try again", style="red")
|
|
else:
|
|
console.print("❌ All mirrors unreachable\n💡 Check network connection", style="red")
|
|
|
|
if __name__ == "__main__":
|
|
main()
|