Handle KeyboardInterrupts in the cli better (#1391)

sdb9696 · web-flow · commit 296af3192e10 · 2024-12-20T14:21:38.000+01:00
Addresses an issue with how `asyncclick` deals with `KeyboardInterrupt`
errors. Instead of the `click.main` receiving `KeyboardInterrupt` it
receives `CancelledError` because it's a task running inside the loop.

Also ensures that discovery catches the `CancelledError` and closes the
http clients.
diff --git a/kasa/cli/common.py b/kasa/cli/common.py
@@ -2,12 +2,14 @@
 
 from __future__ import annotations
 
+import asyncio
 import json
 import re
 import sys
 from collections.abc import Callable
 from contextlib import contextmanager
 from functools import singledispatch, update_wrapper, wraps
+from gettext import gettext
 from typing import TYPE_CHECKING, Any, Final
 
 import asyncclick as click
@@ -238,4 +240,19 @@ async def invoke(self, ctx):
             except Exception as exc:
                 _handle_exception(self._debug, exc)
 
+        def __call__(self, *args, **kwargs):
+            """Run the coroutine in the event loop and print any exceptions.
+
+            python click catches KeyboardInterrupt in main, raises Abort()
+            and does sys.exit. asyncclick doesn't properly handle a coroutine
+            receiving CancelledError on a KeyboardInterrupt, so we catch the
+            KeyboardInterrupt here once asyncio.run has re-raised it. This
+            avoids large stacktraces when a user presses Ctrl-C.
+            """
+            try:
+                asyncio.run(self.main(*args, **kwargs))
+            except KeyboardInterrupt:
+                click.echo(gettext("\nAborted!"), file=sys.stderr)
+                sys.exit(1)
+
     return _CommandCls
diff --git a/kasa/discover.py b/kasa/discover.py
@@ -498,7 +498,7 @@ async def discover(
         try:
             _LOGGER.debug("Waiting %s seconds for responses...", discovery_timeout)
             await protocol.wait_for_discovery_to_complete()
-        except KasaException as ex:
+        except (KasaException, asyncio.CancelledError) as ex:
             for device in protocol.discovered_devices.values():
                 await device.protocol.close()
             raise ex
