[IMPR] enable KeyboardInterrupt with -async option

Change-Id: I021005b9c7c80e5ef6a14ef5246d006596333ddb
---
M scripts/archivebot.py
1 file changed, 54 insertions(+), 10 deletions(-)

diff --git a/scripts/archivebot.py b/scripts/archivebot.py
index a530410..ca1aac9 100755
--- a/scripts/archivebot.py
+++ b/scripts/archivebot.py
@@ -105,15 +105,16 @@
   -keep           Preserve thread order in archive even if threads are
                   archived later
   -sort           Sort archive by timestamp; should not be used with -keep
-  -async          Run the bot in parallel tasks. This is experimental
-                  and the bot cannot be stopped with KeyboardInterrupt
+  -async          Run the bot in parallel tasks.
 
 .. versionchanged:: 7.6
    Localized variables for "archive" template parameter are supported.
    `User:MiszaBot/config` is the default template. `-keep` option was
    added.
 .. versionchanged:: 7.7
-   `-sort` and `-async` options were added.
+   ``-sort`` and ``-async`` options were added.
+.. versionchanged:: 8.2
+   KeyboardInterrupt was enabled with ``-async`` option.
 """
 #
 # (C) Pywikibot team, 2006-2023
@@ -124,13 +125,15 @@
 import locale
 import os
 import re
+import signal
+import threading
 import time
 from collections import OrderedDict, defaultdict
 from concurrent.futures import ThreadPoolExecutor
 from hashlib import md5
 from math import ceil
 from textwrap import fill
-from typing import Any, Optional, Pattern
+from typing import Any, Optional, Pattern, Union
 from warnings import warn
 
 import pywikibot
@@ -145,6 +148,7 @@
     to_local_digits,
 )
 from pywikibot.time import MW_KEYS, parse_duration, str2timedelta
+from pywikibot.tools import PYTHON_VERSION
 
 
 ShouldArchive = Tuple[str, str]
@@ -877,6 +881,11 @@
 
     :param args: command line arguments
     """
+    def signal_handler(signum, frame):
+        pywikibot.info('\n<<lightyellow>>User quit bot run...')
+        exiting.set()
+
+    exiting = threading.Event()
     filename = None
     pagename = None
     namespace = None
@@ -885,7 +894,7 @@
     calc = None
     keep = False
     sort = False
-    asyncronous = False
+    asynchronous = False
     templates = []
 
     local_args = pywikibot.handle_args(args)
@@ -920,7 +929,7 @@
         elif option == 'sort':
             sort = True
         elif option == 'async':
-            asyncronous = True
+            asynchronous = True
 
     site = pywikibot.Site()
 
@@ -932,6 +941,12 @@
         pywikibot.info('No template was specified, using default {{{{{}}}}}.'
                        .format(templates[0]))
 
+    if asynchronous:
+        signal.signal(signal.SIGINT, signal_handler)
+        context = ThreadPoolExecutor
+    else:
+        context = nullcontext
+
     for template_name in templates:
         tmpl = pywikibot.Page(site, template_name, ns=10)
         if filename:
@@ -949,13 +964,33 @@
                                      content=True)
 
         botargs = tmpl, salt, force, keep, sort
-        context = ThreadPoolExecutor if asyncronous else nullcontext
+        futures = []  # needed for Python < 3.9
         with context() as executor:
             for pg in gen:
-                if asyncronous:
-                    executor.submit(process_page, pg, *botargs)
+                if asynchronous:
+                    future = executor.submit(process_page, pg, *botargs)
+
+                    if PYTHON_VERSION < (3, 9):
+                        futures.append(future)
+
+                    if not exiting.is_set():
+                        continue
+
+                    canceled: Union[str, int] = ''
+                    pywikibot.info(
+                        '<<lightyellow>>Canceling pending Futures... ',
+                        newline=False)
+
+                    if PYTHON_VERSION < (3, 9):
+                        canceled = sum(future.cancel() for future in futures)
+                    else:
+                        executor.shutdown(cancel_futures=True)
+
+                    pywikibot.info(f'{canceled} done')
+                    break
+
                 elif not process_page(pg, *botargs):
-                    return
+                    break
 
 
 if __name__ == '__main__':

