simple_repository_server.py
1 """PEP 700-compliant Simple Repository API server for serving wheels in tests. 2 3 This replaces the plain ``http.server`` approach so that uv's ``exclude-newer`` 4 can filter packages by upload time when resolving the local dev wheel. 5 """ 6 7 from __future__ import annotations 8 9 import hashlib 10 import json 11 import threading 12 from http.server import HTTPServer, SimpleHTTPRequestHandler 13 from pathlib import Path 14 from typing import Any 15 16 from typing_extensions import Self 17 18 19 def _make_handler(wheel_dir: Path) -> type[SimpleHTTPRequestHandler]: 20 class Handler(SimpleHTTPRequestHandler): 21 def do_GET(self) -> None: 22 path = self.path.split("?")[0].split("#")[0].rstrip("/") 23 if path in ("", "/simple"): 24 self._serve_index() 25 elif path == "/simple/mlflow": 26 self._serve_project() 27 else: 28 super().do_GET() 29 30 def _wants_json(self) -> bool: 31 accept = self.headers.get("Accept", "") 32 return "application/vnd.pypi.simple.v1+json" in accept 33 34 def _serve_index(self) -> None: 35 if self._wants_json(): 36 body = json.dumps({ 37 "meta": {"api-version": "1.1"}, 38 "projects": [{"name": "mlflow"}], 39 }).encode() 40 content_type = "application/vnd.pypi.simple.v1+json" 41 else: 42 body = b"<html><body><a href='/simple/mlflow/'>mlflow</a></body></html>" 43 content_type = "text/html" 44 self._respond(200, content_type, body) 45 46 def _serve_project(self) -> None: 47 files = [] 48 versions: set[str] = set() 49 for f in sorted(wheel_dir.iterdir()): 50 if f.suffix != ".whl": 51 continue 52 sha256 = hashlib.sha256(f.read_bytes()).hexdigest() 53 version = f.stem.split("-")[1] 54 versions.add(version) 55 files.append({ 56 "filename": f.name, 57 "url": f"/mlflow/{f.name}#sha256={sha256}", 58 "hashes": {"sha256": sha256}, 59 "size": f.stat().st_size, 60 # A date safely before any exclude-newer cutoff so the dev 61 # wheel is always resolvable. 62 "upload-time": "2020-01-01T00:00:00Z", 63 }) 64 65 if self._wants_json(): 66 body = json.dumps({ 67 "meta": {"api-version": "1.1"}, 68 "name": "mlflow", 69 "versions": sorted(versions), 70 "files": files, 71 }).encode() 72 content_type = "application/vnd.pypi.simple.v1+json" 73 else: 74 links = "".join( 75 f'<a href="{f["url"]}" data-dist-info-metadata="false">{f["filename"]}</a>' 76 for f in files 77 ) 78 body = f"<html><body>{links}</body></html>".encode() 79 content_type = "text/html" 80 self._respond(200, content_type, body) 81 82 def _respond(self, status: int, content_type: str, body: bytes) -> None: 83 self.send_response(status) 84 self.send_header("Content-Type", content_type) 85 self.send_header("Content-Length", str(len(body))) 86 self.end_headers() 87 self.wfile.write(body) 88 89 def translate_path(self, path: str) -> str: 90 path = path.split("?")[0].split("#")[0] 91 return str(wheel_dir.parent / path.lstrip("/")) 92 93 def log_message(self, format: str, *args: Any) -> None: 94 pass 95 96 return Handler 97 98 99 class SimpleRepositoryServer: 100 def __init__(self, wheel_dir: Path, port: int) -> None: 101 handler = _make_handler(wheel_dir) 102 self._server = HTTPServer(("", port), handler) 103 self._thread = threading.Thread( 104 target=self._server.serve_forever, daemon=True, name="simple-repository-server" 105 ) 106 107 @property 108 def url(self) -> str: 109 _, port = self._server.server_address 110 return f"http://localhost:{port}/simple" 111 112 def start(self) -> None: 113 self._thread.start() 114 115 def shutdown(self) -> None: 116 self._server.shutdown() 117 118 def __enter__(self) -> Self: 119 self.start() 120 return self 121 122 def __exit__(self, *args: Any) -> None: 123 self.shutdown() 124 self._thread.join()