Fix multiprocessing support for Windows binary (#632)

[etc/vim.git] / blackd.py

diff --git a/blackd.py b/blackd.py
index f2bbc8a4a319e12088e766f0e7a570bc9b3bfe54..1a3b3690bc3b38b41b034eadf41cb271c203d185 100644 (file)
--- a/blackd.py
+++ b/blackd.py
@@ -2,8 +2,10 @@ import asyncio
 from concurrent.futures import Executor, ProcessPoolExecutor
 from functools import partial
 import logging
+from multiprocessing import freeze_support
 
 from aiohttp import web
+import aiohttp_cors
 import black
 import click
 
@@ -14,8 +16,18 @@ VERSION_HEADER = "X-Protocol-Version"
 LINE_LENGTH_HEADER = "X-Line-Length"
 PYTHON_VARIANT_HEADER = "X-Python-Variant"
 SKIP_STRING_NORMALIZATION_HEADER = "X-Skip-String-Normalization"
+SKIP_NUMERIC_UNDERSCORE_NORMALIZATION_HEADER = "X-Skip-Numeric-Underscore-Normalization"
 FAST_OR_SAFE_HEADER = "X-Fast-Or-Safe"
 
+BLACK_HEADERS = [
+    VERSION_HEADER,
+    LINE_LENGTH_HEADER,
+    PYTHON_VARIANT_HEADER,
+    SKIP_STRING_NORMALIZATION_HEADER,
+    SKIP_NUMERIC_UNDERSCORE_NORMALIZATION_HEADER,
+    FAST_OR_SAFE_HEADER,
+]
+
 
 @click.command(context_settings={"help_option_names": ["-h", "--help"]})
 @click.option(
@@ -34,7 +46,18 @@ def main(bind_host: str, bind_port: int) -> None:
 def make_app() -> web.Application:
     app = web.Application()
     executor = ProcessPoolExecutor()
-    app.add_routes([web.post("/", partial(handle, executor=executor))])
+
+    cors = aiohttp_cors.setup(app)
+    resource = cors.add(app.router.add_resource("/"))
+    cors.add(
+        resource.add_route("POST", partial(handle, executor=executor)),
+        {
+            "*": aiohttp_cors.ResourceOptions(
+                allow_headers=(*BLACK_HEADERS, "Content-Type"), expose_headers="*"
+            )
+        },
+    )
+
     return app
 
 
@@ -69,11 +92,17 @@ async def handle(request: web.Request, executor: Executor) -> web.Response:
         skip_string_normalization = bool(
             request.headers.get(SKIP_STRING_NORMALIZATION_HEADER, False)
         )
+        skip_numeric_underscore_normalization = bool(
+            request.headers.get(SKIP_NUMERIC_UNDERSCORE_NORMALIZATION_HEADER, False)
+        )
         fast = False
         if request.headers.get(FAST_OR_SAFE_HEADER, "safe") == "fast":
             fast = True
         mode = black.FileMode.from_configuration(
-            py36=py36, pyi=pyi, skip_string_normalization=skip_string_normalization
+            py36=py36,
+            pyi=pyi,
+            skip_string_normalization=skip_string_normalization,
+            skip_numeric_underscore_normalization=skip_numeric_underscore_normalization,
         )
         req_bytes = await request.content.read()
         charset = request.charset if request.charset is not None else "utf8"
@@ -101,6 +130,11 @@ async def handle(request: web.Request, executor: Executor) -> web.Response:
         return web.Response(status=500, text=str(e))
 
 
-if __name__ == "__main__":
+def patched_main() -> None:
+    freeze_support()
     black.patch_click()
     main()
+
+
+if __name__ == "__main__":
+    patched_main()