rewrite to starlette

2 files changed, 21 insertions, 64 deletions

diff --git a/syntax-highlighting-server.py b/syntax-highlighting-server.py
index 6a13393..244e9d4 100755
--- a/syntax-highlighting-server.py
+++ b/syntax-highlighting-server.py
@@ -1,19 +1,22 @@
 #!/usr/bin/env python3
 
-import asyncio
-import argparse
-import gc
-import os
+from starlette.applications import Starlette
+from starlette.responses import HTMLResponse
+from starlette.routing import Route
 import pygments
-from aiohttp import web
-from concurrent.futures import ProcessPoolExecutor
-from pygments import highlight
 from pygments.formatters import HtmlFormatter
 from pygments.lexers import guess_lexer, guess_lexer_for_filename
 from pygments.lexers.special import TextLexer
 from pygments.util import ClassNotFound
 
-def do_highlight(filename, data, formatter):
+formatter = None
+
+async def startup_task():
+    guess_lexer('')
+    global formatter
+    formatter = HtmlFormatter(style='pastie', nobackground=True)
+
+def do_highlight(filename, data):
     try:
         lexer = guess_lexer_for_filename(filename, data)
     except ClassNotFound:
@@ -27,61 +30,15 @@ def do_highlight(filename, data, formatter):
     return ''.join([
         f'<!-- Pygments {pygments.__version__}: {lexer.name} ({lexer.__class__.__name__}) -->',
         '<style>', formatter.get_style_defs('.highlight'), '</style>',
-        highlight(data, lexer, formatter)
+        pygments.highlight(data, lexer, formatter)
     ])
 
-def parse_args():
-    parser = argparse.ArgumentParser(description='syntax highlighting server',
-            formatter_class=argparse.ArgumentDefaultsHelpFormatter)
-    parser.add_argument('--host', type=str, default='127.0.0.1',
-            help='the host to listen on')
-    parser.add_argument('--port', type=int, default=4872,
-            help='the port to listen on')
-    parser.add_argument('--style', type=str, default='pastie',
-            help='pygments formatting style')
-    parser.add_argument('--preload', type=bool, default=True,
-            help='preload lexers to reduce fork memory usage')
-    parser.add_argument('--workers', type=int, default=0,
-            help='number of workers, 0 is one per cpu')
-    return parser.parse_args()
-
-async def handle_highlight(request):
-    loop = asyncio.get_running_loop()
-    text = await request.text()
-    result = await loop.run_in_executor(
-            request.app['pool'], do_highlight,
-            request.query['filename'], text, request.app['formatter'])
-    return web.Response(text=result)
-
-def run(args, pool=None):
-    app = web.Application()
-    app['pool'] = pool
-    app['formatter'] = HtmlFormatter(style=args.style, nobackground=True)
-    app.add_routes([web.post('/highlight', handle_highlight)])
-    web.run_app(app, host=args.host, port=args.port)
-
-def noop(*args, **kws):
-    pass
-
-def main():
-    args = parse_args()
-
-    if args.preload:
-        guess_lexer('')
-
-    if args.workers == 0:
-        workers = len(os.sched_getaffinity(0))
-    else:
-        workers = args.workers
-
-    if workers == 1:
-        run(args)
-    else:
-        with ProcessPoolExecutor(max_workers=workers) as pool:
-            gc.collect()
-            gc.freeze()
-            pool.map(noop, [None] * workers)
-            run(args, pool)
+async def run_highlight(request):
+    return HTMLResponse(do_highlight(request.query_params['filename'], await request.body()))
 
-if __name__ == '__main__':
-    main()
+app = Starlette(
+        routes=[
+            Route('/highlight', run_highlight, methods=["POST"]),
+            ],
+        on_startup=[startup_task],
+        )
diff --git a/syntax-highlighting.service.in b/syntax-highlighting.service.in
index 02b5cf1..c65e163 100644
--- a/syntax-highlighting.service.in
+++ b/syntax-highlighting.service.in
@@ -3,7 +3,7 @@ Description=syntax highlighting server
 
 [Service]
 Type=simple
-ExecStart=@libdir@/cgit/syntax-highlighting-server.py
+ExecStart=/usr/bin/uvicorn --app-dir @libdir@/cgit --port 4872 syntax-highlighting-server:app
 
 DynamicUser=yes
 IPAddressAllow=localhost