mirror of
https://github.com/kovidgoyal/kitty.git
synced 2024-09-21 03:27:55 +03:00
Get multiprocessing working in kitty
Monkeypatch the stdlib multiprocessing module to use kitty as its python interpreter for spawning worker processes.
This commit is contained in:
parent
2cb25cf5a8
commit
2c6e5a6e73
@ -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.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):
|
||||
|
66
kitty/multiprocessing.py
Normal file
66
kitty/multiprocessing.py
Normal file
@ -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()
|
@ -42,3 +42,7 @@ def test_line_edit(self):
|
||||
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()
|
||||
|
Loading…
Reference in New Issue
Block a user