Python mini HTTP server to serve a single request and shutdown

server.py

import http.server
import socketserver
import threading
import urllib.parse


def wait_for_oauth_callback(port: int = 9090) -> str:
    code: str | None = None

    class CallbackHandler(http.server.SimpleHTTPRequestHandler):
        def do_GET(self):
            nonlocal code
            code = urllib.parse.parse_qs(urllib.parse.urlparse(self.path).query).get("code")[0]

            self.send_response(200)
            self.send_header("Content-type", "text/plain")
            self.end_headers()
            self.wfile.write(b"OAuth callback received. You can close this window.")

            # signal the server to shut down (must be run in a separate thread)
            shutdown_thread = threading.Thread(target=self.server.shutdown)
            shutdown_thread.start()

    with socketserver.TCPServer(("", port), CallbackHandler) as httpd:
        httpd.serve_forever()  # won't actually run forever

    return code


if __name__ == "__main__":
    code = wait_for_oauth_callback(port=9090)
    print("got", code)