From 8c471f827bae5e7e155d05dff3126499ca130f72 Mon Sep 17 00:00:00 2001
From: jack1142 <6032823+jack1142@users.noreply.github.com>
Date: Tue, 9 Feb 2021 00:25:53 +0100
Subject: [PATCH] Don't stop the loop in `run_until_complete()` on `SystemExit`
 and `KeyboardInterrupt` exceptions (#337)

* Don't stop the loop in run_until_complete on SystemExit & KeyboardInterrupt
* Add unit test (based on CPython's equivalent)
* Fix the test to resume its coverage.

Co-authored-by: Victor Stinner <victor.stinner@gmail.com>
Co-authored-by: Fantix King <fantix.king@gmail.com>
---
 tests/test_base.py | 20 ++++++++++++++++++++
 uvloop/loop.pyx    |  9 ++++++++-
 2 files changed, 28 insertions(+), 1 deletion(-)

diff --git a/tests/test_base.py b/tests/test_base.py
index e46178a..6d87f54 100644
--- a/tests/test_base.py
+++ b/tests/test_base.py
@@ -361,6 +361,26 @@ class _TestBase:
         # done-callback for the previous future.
         self.loop.run_until_complete(foo(0.2))
 
+    def test_run_until_complete_keyboard_interrupt(self):
+        # Issue #336: run_until_complete() must not schedule a pending
+        # call to stop() if the future raised a KeyboardInterrupt
+        async def raise_keyboard_interrupt():
+            raise KeyboardInterrupt
+
+        self.loop._process_events = mock.Mock()
+
+        with self.assertRaises(KeyboardInterrupt):
+            self.loop.run_until_complete(raise_keyboard_interrupt())
+
+        def func():
+            self.loop.stop()
+            func.called = True
+
+        func.called = False
+        self.loop.call_later(0.01, func)
+        self.loop.run_forever()
+        self.assertTrue(func.called)
+
     def test_debug_slow_callbacks(self):
         logger = logging.getLogger('asyncio')
         self.loop.set_debug(True)
diff --git a/uvloop/loop.pyx b/uvloop/loop.pyx
index a7b24d6..9ebf3a6 100644
--- a/uvloop/loop.pyx
+++ b/uvloop/loop.pyx
@@ -1431,7 +1431,14 @@ cdef class Loop:
             # is no need to log the "destroy pending task" message
             future._log_destroy_pending = False
 
-        done_cb = lambda fut: self.stop()
+        def done_cb(fut):
+            if not fut.cancelled():
+                exc = fut.exception()
+                if isinstance(exc, (SystemExit, KeyboardInterrupt)):
+                    # Issue #336: run_forever() already finished,
+                    # no need to stop it.
+                    return
+            self.stop()
 
         future.add_done_callback(done_cb)
         try: