Created
March 7, 2019 16:45
-
-
Save antirais/17f5d5c1b7b5ee7a1150bbea1dd9a802 to your computer and use it in GitHub Desktop.
uvicorn server header override
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
#!/usr/bin/env python3 | |
# coding: utf-8 | |
# PYTHON_ARGCOMPLETE_OK | |
from email.utils import formatdate | |
from src.main import app | |
from uvicorn.config import Config | |
from uvicorn.main import Server | |
import argparse | |
import argcomplete | |
import uvicorn | |
import time | |
import types | |
async def on_tick(self, counter) -> bool: | |
# Update the default headers, once per second. | |
if counter % 10 == 0: | |
current_time = time.time() | |
current_date = formatdate(current_time, usegmt=True).encode() | |
self.server_state.default_headers = [(b"date", current_date)] | |
# Callback to `callback_notify` once every `timeout_notify` seconds. | |
if self.config.callback_notify is not None: | |
if counter % (10 * self.config.timeout_notify) == 0: | |
await self.config.callback_notify() | |
# Determine if we should exit. | |
if self.should_exit: | |
return True | |
if self.config.limit_max_requests is not None: | |
return self.server_state.total_requests >= self.config.limit_max_requests | |
return False | |
def monkey_patch_header(server, method): | |
server.on_tick = types.MethodType(method, server) | |
def run(app, **kwargs): | |
config = Config(app, **kwargs) | |
server = Server(config=config) | |
monkey_patch_header(server, on_tick) | |
server.run() | |
if __name__ == "__main__": | |
parser = argparse.ArgumentParser(add_help=False) | |
parser.add_argument('-h', '--host', default='0.0.0.0', required=False) | |
parser.add_argument('-p', '--port', default=8000, type=int, required=False) | |
parser.add_argument('-l', '--log_level', default='info', required=False) | |
parser.add_argument('-r', '--reload', default=False, action='store_true') | |
argcomplete.autocomplete(parser) | |
args = parser.parse_args() | |
run(app, host=args.host, port=args.port, log_level=args.log_level, reload=args.reload) |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment