From 89af7c100a8e253aa7fea92bd523d8649d60fbc3 Mon Sep 17 00:00:00 2001 From: Jennifer Taylor Date: Sun, 16 May 2021 19:39:44 +0000 Subject: [PATCH] Better multi-threaded Ctrl+C handling. --- bemani/format/afp/blend.py | 15 ++++++++++++++- 1 file changed, 14 insertions(+), 1 deletion(-) diff --git a/bemani/format/afp/blend.py b/bemani/format/afp/blend.py index 925fecf..bb679e9 100644 --- a/bemani/format/afp/blend.py +++ b/bemani/format/afp/blend.py @@ -1,6 +1,7 @@ import multiprocessing +import signal from PIL import Image # type: ignore -from typing import List, Sequence, Tuple +from typing import Any, List, Sequence, Tuple from .types.generic import Color, Matrix, Point @@ -225,6 +226,14 @@ def affine_composite( work: multiprocessing.Queue = multiprocessing.Queue() results: multiprocessing.Queue = multiprocessing.Queue() expected: int = 0 + interrupted: bool = False + + def ctrlc(sig: Any, frame: Any) -> None: + nonlocal interrupted + interrupted = True + + original_handler = signal.getsignal(signal.SIGINT) + signal.signal(signal.SIGINT, ctrlc) for _ in range(cores): proc = multiprocessing.Process( @@ -270,6 +279,10 @@ def affine_composite( for proc in procs: proc.join() + signal.signal(signal.SIGINT, original_handler) + if interrupted: + raise KeyboardInterrupt() + img = Image.frombytes('RGBA', (imgwidth, imgheight), b''.join(lines)) return img