blackd.py

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