blackd.py

   1 import asyncio
   2 from concurrent.futures import Executor, ProcessPoolExecutor
   3 from functools import partial
   4 import logging
   5 from multiprocessing import freeze_support
   6
   7 from aiohttp import web
   8 import aiohttp_cors
   9 import black
  10 import click
  11
  12 # This is used internally by tests to shut down the server prematurely
  13 _stop_signal = asyncio.Event()
  14
  15 VERSION_HEADER = "X-Protocol-Version"
  16 LINE_LENGTH_HEADER = "X-Line-Length"
  17 PYTHON_VARIANT_HEADER = "X-Python-Variant"
  18 SKIP_STRING_NORMALIZATION_HEADER = "X-Skip-String-Normalization"
  19 SKIP_NUMERIC_UNDERSCORE_NORMALIZATION_HEADER = "X-Skip-Numeric-Underscore-Normalization"
  20 FAST_OR_SAFE_HEADER = "X-Fast-Or-Safe"
  21
  22 BLACK_HEADERS = [
  23     VERSION_HEADER,
  24     LINE_LENGTH_HEADER,
  25     PYTHON_VARIANT_HEADER,
  26     SKIP_STRING_NORMALIZATION_HEADER,
  27     SKIP_NUMERIC_UNDERSCORE_NORMALIZATION_HEADER,
  28     FAST_OR_SAFE_HEADER,
  29 ]
  30
  31
  32 @click.command(context_settings={"help_option_names": ["-h", "--help"]})
  33 @click.option(
  34     "--bind-host", type=str, help="Address to bind the server to.", default="localhost"
  35 )
  36 @click.option("--bind-port", type=int, help="Port to listen on", default=45484)
  37 @click.version_option(version=black.__version__)
  38 def main(bind_host: str, bind_port: int) -> None:
  39     logging.basicConfig(level=logging.INFO)
  40     app = make_app()
  41     ver = black.__version__
  42     black.out(f"blackd version {ver} listening on {bind_host} port {bind_port}")
  43     web.run_app(app, host=bind_host, port=bind_port, handle_signals=True, print=None)
  44
  45
  46 def make_app() -> web.Application:
  47     app = web.Application()
  48     executor = ProcessPoolExecutor()
  49
  50     cors = aiohttp_cors.setup(app)
  51     resource = cors.add(app.router.add_resource("/"))
  52     cors.add(
  53         resource.add_route("POST", partial(handle, executor=executor)),
  54         {
  55             "*": aiohttp_cors.ResourceOptions(
  56                 allow_headers=(*BLACK_HEADERS, "Content-Type"), expose_headers="*"
  57             )
  58         },
  59     )
  60
  61     return app
  62
  63
  64 async def handle(request: web.Request, executor: Executor) -> web.Response:
  65     try:
  66         if request.headers.get(VERSION_HEADER, "1") != "1":
  67             return web.Response(
  68                 status=501, text="This server only supports protocol version 1"
  69             )
  70         try:
  71             line_length = int(
  72                 request.headers.get(LINE_LENGTH_HEADER, black.DEFAULT_LINE_LENGTH)
  73             )
  74         except ValueError:
  75             return web.Response(status=400, text="Invalid line length header value")
  76         py36 = False
  77         pyi = False
  78         if PYTHON_VARIANT_HEADER in request.headers:
  79             value = request.headers[PYTHON_VARIANT_HEADER]
  80             if value == "pyi":
  81                 pyi = True
  82             else:
  83                 try:
  84                     major, *rest = value.split(".")
  85                     if int(major) == 3 and len(rest) > 0:
  86                         if int(rest[0]) >= 6:
  87                             py36 = True
  88                 except ValueError:
  89                     return web.Response(
  90                         status=400, text=f"Invalid value for {PYTHON_VARIANT_HEADER}"
  91                     )
  92         skip_string_normalization = bool(
  93             request.headers.get(SKIP_STRING_NORMALIZATION_HEADER, False)
  94         )
  95         skip_numeric_underscore_normalization = bool(
  96             request.headers.get(SKIP_NUMERIC_UNDERSCORE_NORMALIZATION_HEADER, False)
  97         )
  98         fast = False
  99         if request.headers.get(FAST_OR_SAFE_HEADER, "safe") == "fast":
 100             fast = True
 101         mode = black.FileMode.from_configuration(
 102             py36=py36,
 103             pyi=pyi,
 104             skip_string_normalization=skip_string_normalization,
 105             skip_numeric_underscore_normalization=skip_numeric_underscore_normalization,
 106         )
 107         req_bytes = await request.content.read()
 108         charset = request.charset if request.charset is not None else "utf8"
 109         req_str = req_bytes.decode(charset)
 110         loop = asyncio.get_event_loop()
 111         formatted_str = await loop.run_in_executor(
 112             executor,
 113             partial(
 114                 black.format_file_contents,
 115                 req_str,
 116                 line_length=line_length,
 117                 fast=fast,
 118                 mode=mode,
 119             ),
 120         )
 121         return web.Response(
 122             content_type=request.content_type, charset=charset, text=formatted_str
 123         )
 124     except black.NothingChanged:
 125         return web.Response(status=204)
 126     except black.InvalidInput as e:
 127         return web.Response(status=400, text=str(e))
 128     except Exception as e:
 129         logging.exception("Exception during handling a request")
 130         return web.Response(status=500, text=str(e))
 131
 132
 133 def patched_main() -> None:
 134     freeze_support()
 135     black.patch_click()
 136     main()
 137
 138
 139 if __name__ == "__main__":
 140     patched_main()