refactor, re-add max-workers

1 files changed, 21 insertions, 13 deletions

diff --git a/syntax-highlighting-server.py b/syntax-highlighting-server.py
index d270c4c..9f534c6 100755
--- a/syntax-highlighting-server.py
+++ b/syntax-highlighting-server.py
@@ -17,7 +17,7 @@ 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, style):
+def do_highlight(filename, data, formatter):
     try:
         lexer = guess_lexer_for_filename(filename, data)
     except ClassNotFound:
@@ -28,7 +28,6 @@ def do_highlight(filename, data, style):
                 lexer = TextLexer()
         except ClassNotFound:
             lexer = TextLexer()
-    formatter = HtmlFormatter(style=style, nobackground=True)
     return ''.join([
         f'<!-- Pygments {pygments.__version__}: {lexer.name} ({lexer.__class__.__name__}) -->',
         '<style>', formatter.get_style_defs('.highlight'), '</style>',
@@ -46,7 +45,9 @@ def parse_args():
     parser.add_argument('--style', type=str, default='pastie',
             help='pygments formatting style')
     parser.add_argument('--preload', type=bool, default=True,
-            help='preload lexers and formatters to reduce fork memory usage')
+            help='preload lexers to reduce fork memory usage')
+    parser.add_argument('--max-workers', type=int, default=0,
+            help='number of workers, 0 is one per cpu')
     return parser.parse_args()
 
 async def handle_highlight(request):
@@ -56,24 +57,31 @@ async def handle_highlight(request):
     text = await request.text()
     result = await loop.run_in_executor(
             request.app['pool'], do_highlight,
-            request.query['filename'], text, request.app['style'])
+            request.query['filename'], text, request.app['formatter'])
     return web.Response(text=result)
 
+def run(args, pool=None):
+    from aiohttp import web
+    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 main():
     args = parse_args()
 
-    from aiohttp import web
-    from concurrent.futures import ProcessPoolExecutor
-
     if args.preload:
         guess_lexer('')
 
-    with ProcessPoolExecutor() as pool:
-        app = web.Application()
-        app['pool'] = pool
-        app['style'] = args.style
-        app.add_routes([web.post('/highlight', handle_highlight)])
-        web.run_app(app, host=args.host, port=args.port)
+    if args.max_workers == 0:
+        max_workers = None
+    else:
+        max_workers = args.max_workers
+
+    from concurrent.futures import ProcessPoolExecutor
+    with ProcessPoolExecutor(max_workers=max_workers) as pool:
+        run(args, pool)
 
 if __name__ == '__main__':
     main()