From 296af3192e10718648069deb29c66175bb0fbc1b Mon Sep 17 00:00:00 2001
From: "Steven B." <51370195+sdb9696@users.noreply.github.com>
Date: Fri, 20 Dec 2024 13:21:38 +0000
Subject: [PATCH] Handle KeyboardInterrupts in the cli better (#1391)

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.
---
 kasa/cli/common.py | 17 +++++++++++++++++
 kasa/discover.py   |  2 +-
 2 files changed, 18 insertions(+), 1 deletion(-)

diff --git a/kasa/cli/common.py b/kasa/cli/common.py
index 649df065..5114f7af 100644
--- 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 @@ def CatchAllExceptions(cls):
             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
index 77ef80be..b696c370 100755
--- a/kasa/discover.py
+++ b/kasa/discover.py
@@ -498,7 +498,7 @@ class 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