From 2c6e5a6e73c10f4e250d8e736e19eacb909d03f4 Mon Sep 17 00:00:00 2001
From: Kovid Goyal <kovid@kovidgoyal.net>
Date: Sun, 21 Jun 2020 14:47:24 +0530
Subject: [PATCH] Get multiprocessing working in kitty

Monkeypatch the stdlib multiprocessing module to
use kitty as its python interpreter for spawning worker
processes.
---
 kittens/diff/highlight.py | 12 ++-----
 kitty/multiprocessing.py  | 66 +++++++++++++++++++++++++++++++++++++++
 kitty_tests/tui.py        |  4 +++
 3 files changed, 72 insertions(+), 10 deletions(-)
 create mode 100644 kitty/multiprocessing.py

diff --git a/kittens/diff/highlight.py b/kittens/diff/highlight.py
index 376d5965c..55d852ded 100644
--- a/kittens/diff/highlight.py
+++ b/kittens/diff/highlight.py
@@ -5,8 +5,6 @@
 import concurrent
 import os
 import re
-import sys
-from multiprocessing import get_context
 from typing import IO, Dict, Iterable, List, Optional, Tuple, Union, cast
 
 from pygments import highlight  # type: ignore
@@ -14,6 +12,7 @@ from pygments.formatter import Formatter  # type: ignore
 from pygments.lexers import get_lexer_for_filename  # type: ignore
 from pygments.util import ClassNotFound  # type: ignore
 
+from kitty.multiprocessing import get_process_pool_executor
 from kitty.rgb import color_as_sgr, parse_sharp
 
 from .collect import Collection, Segment, data_for_path, lines_for_path
@@ -141,14 +140,7 @@ def highlight_for_diff(path: str, aliases: Dict[str, str]) -> DiffHighlight:
 def highlight_collection(collection: Collection, aliases: Optional[Dict[str, str]] = None) -> Union[str, Dict[str, DiffHighlight]]:
     jobs = {}
     ans: Dict[str, DiffHighlight] = {}
-    if sys.version_info[:2] >= (3, 7):
-        # On macOS as of python 3.8 the default executor is changed to spawn
-        # which causes failures, so use fork, which is also faster
-        ppe = concurrent.futures.ProcessPoolExecutor(mp_context=get_context('fork'))
-    else:
-        ppe = concurrent.futures.ProcessPoolExecutor()
-
-    with ppe as executor:
+    with get_process_pool_executor(prefer_fork=True) as executor:
         for path, item_type, other_path in collection:
             if item_type != 'rename':
                 for p in (path, other_path):
diff --git a/kitty/multiprocessing.py b/kitty/multiprocessing.py
new file mode 100644
index 000000000..a3803dd13
--- /dev/null
+++ b/kitty/multiprocessing.py
@@ -0,0 +1,66 @@
+#!/usr/bin/env python
+# vim:fileencoding=utf-8
+# License: GPLv3 Copyright: 2020, Kovid Goyal <kovid at kovidgoyal.net>
+
+# Monkeypatch the stdlib multiprocessing module to work with the embedded python
+# in kitty, when using the spawn launcher.
+
+
+from concurrent.futures import ProcessPoolExecutor
+from multiprocessing import util  # type: ignore
+from multiprocessing import context, get_all_start_methods, get_context, spawn
+from typing import Any, Callable, List, Optional, Tuple, Union
+
+from .constants import kitty_exe
+
+orig_spawn_passfds = util.spawnv_passfds
+orig_executable = spawn.get_executable()
+
+
+def spawnv_passfds(path: str, args: List[str], passfds: List[int]) -> Any:
+    idx = args.index('-c')
+    patched_args = [spawn.get_executable(), '+runpy'] + args[idx + 1:]
+    return orig_spawn_passfds(kitty_exe(), patched_args, passfds)
+
+
+def monkey_patch_multiprocessing() -> None:
+    # Use kitty to run the worker process used by multiprocessing
+    spawn.set_executable(kitty_exe())
+    util.spawnv_passfds = spawnv_passfds
+
+
+def unmonkey_patch_multiprocessing() -> None:
+    spawn.set_executable(orig_executable)
+    util.spawnv_passfds = orig_spawn_passfds
+
+
+def get_process_pool_executor(
+    prefer_fork: bool = False,
+    max_workers: Optional[int] = None,
+    initializer: Optional[Callable] = None,
+    initargs: Tuple[Any, ...] = ()
+) -> ProcessPoolExecutor:
+    if prefer_fork and 'fork' in get_all_start_methods():
+        ctx: Union[context.DefaultContext, context.ForkContext] = get_context('fork')
+    else:
+        monkey_patch_multiprocessing()
+        ctx = get_context()
+    try:
+        return ProcessPoolExecutor(max_workers=max_workers, initializer=initializer, initargs=initargs, mp_context=ctx)
+    except TypeError:
+        return ProcessPoolExecutor(max_workers=max_workers, initializer=initializer, initargs=initargs)
+
+
+def test_spawn() -> None:
+    monkey_patch_multiprocessing()
+    try:
+        from multiprocessing import get_context
+        ctx = get_context('spawn')
+        q = ctx.Queue()
+        p = ctx.Process(target=q.put, args=('hello',))
+        p.start()
+        x = q.get(timeout=2)
+        assert x == 'hello'
+        p.join()
+    finally:
+        unmonkey_patch_multiprocessing()
diff --git a/kitty_tests/tui.py b/kitty_tests/tui.py
index 213a46364..2a989deb6 100644
--- a/kitty_tests/tui.py
+++ b/kitty_tests/tui.py
@@ -42,3 +42,7 @@ class TestTUI(BaseTest):
         self.ae(le.cursor_pos, 0)
         le.backspace()
         self.assertTrue(le.pending_bell)
+
+    def test_multiprocessing_spawn(self):
+        from kitty.multiprocessing import test_spawn
+        test_spawn()