From f05edb2b43c347d4929efd52c8e2b4e08839f542 Mon Sep 17 00:00:00 2001 From: Aman Sharma <210100011@iitb.ac.in> Date: Tue, 17 Dec 2024 14:00:33 +0530 Subject: [PATCH] Fixed #36016 -- Prevented traceback when quitting makemigrations with Ctrl-C. --- django/db/migrations/questioner.py | 12 +++++++++--- tests/migrations/test_questioner.py | 18 +++++++++++++++++- 2 files changed, 26 insertions(+), 4 deletions(-) diff --git a/django/db/migrations/questioner.py b/django/db/migrations/questioner.py index 2e61195581..445d4410e6 100644 --- a/django/db/migrations/questioner.py +++ b/django/db/migrations/questioner.py @@ -111,17 +111,19 @@ class InteractiveMigrationQuestioner(MigrationQuestioner): for i, choice in enumerate(choices): self.prompt_output.write(" %s) %s" % (i + 1, choice)) self.prompt_output.write("Select an option: ", ending="") - result = input() while True: try: + result = input() value = int(result) except ValueError: pass + except KeyboardInterrupt: + self.prompt_output.write("\nCancelled.") + sys.exit(1) else: if 0 < value <= len(choices): return value self.prompt_output.write("Please select a valid option: ", ending="") - result = input() def _ask_default(self, default=""): """ @@ -148,7 +150,11 @@ class InteractiveMigrationQuestioner(MigrationQuestioner): else: prompt = ">>> " self.prompt_output.write(prompt, ending="") - code = input() + try: + code = input() + except KeyboardInterrupt: + self.prompt_output.write("\nCancelled.") + sys.exit(1) if not code and default: code = default if not code: diff --git a/tests/migrations/test_questioner.py b/tests/migrations/test_questioner.py index ec1013923b..5c737274e4 100644 --- a/tests/migrations/test_questioner.py +++ b/tests/migrations/test_questioner.py @@ -85,8 +85,9 @@ class QuestionerHelperMethodsTests(SimpleTestCase): @mock.patch("builtins.input", side_effect=[KeyboardInterrupt()]) def test_questioner_no_default_keyboard_interrupt(self, mock_input): - with self.assertRaises(KeyboardInterrupt): + with self.assertRaises(SystemExit): self.questioner._ask_default() + self.assertIn("Cancelled.\n", self.prompt.getvalue()) @mock.patch("builtins.input", side_effect=["", "n"]) def test_questioner_no_default_no_user_entry_boolean(self, mock_input): @@ -105,3 +106,18 @@ class QuestionerHelperMethodsTests(SimpleTestCase): expected_msg = f"{question}\n" f" 1) a\n" f" 2) b\n" f" 3) c\n" self.assertIn(expected_msg, self.prompt.getvalue()) self.assertEqual(value, 1) + + @mock.patch("builtins.input", side_effect=[KeyboardInterrupt()]) + def test_questioner_no_choice_keyboard_interrupt(self, mock_input): + question = "Make a choice:" + with self.assertRaises(SystemExit): + self.questioner._choice_input(question, choices="abc") + expected_msg = ( + f"{question}\n" + f" 1) a\n" + f" 2) b\n" + f" 3) c\n" + f"Select an option: \n" + f"Cancelled.\n" + ) + self.assertIn(expected_msg, self.prompt.getvalue())