Introduce ctrl+c and other bug fixes (#1152)

* Intergrate ctrl+c

* stash

* Update

* Fix profile reset

* flake8

Co-authored-by: Daniel Girtler <girtler.daniel@gmail.com>

5 files changed, 231 insertions, 106 deletions

diff --git a/archinstall/lib/menu/global_menu.py b/archinstall/lib/menu/global_menu.py
index b3c5c6a2..23a8ec11 100644
--- a/archinstall/lib/menu/global_menu.py
+++ b/archinstall/lib/menu/global_menu.py
@@ -2,13 +2,15 @@ from __future__ import annotations
 
 from typing import Any, List, Optional, Union
 
+import archinstall
+
 from ..menu import Menu
 from ..menu.selection_menu import Selector, GeneralMenu
 from ..general import SysCommand, secret
 from ..hardware import has_uefi
 from ..models import NetworkConfiguration
 from ..storage import storage
-from ..profiles import is_desktop_profile
+from ..profiles import is_desktop_profile, Profile
 from ..disk import encrypted_partitions
 
 from ..user_interaction import get_password, ask_for_a_timezone, save_config
@@ -41,28 +43,38 @@ class GlobalMenu(GeneralMenu):
 		self._menu_options['archinstall-language'] = \
 			Selector(
 				_('Select Archinstall language'),
-				lambda x: self._select_archinstall_language('English'),
+				lambda x: self._select_archinstall_language(x),
 				default='English')
 		self._menu_options['keyboard-layout'] = \
-			Selector(_('Select keyboard layout'), lambda preset: select_language('us',preset), default='us')
+			Selector(
+				_('Select keyboard layout'),
+				lambda preset: select_language(preset),
+				default='us')
 		self._menu_options['mirror-region'] = \
 			Selector(
 				_('Select mirror region'),
-				select_mirror_regions,
+				lambda preset: select_mirror_regions(preset),
 				display_func=lambda x: list(x.keys()) if x else '[]',
 				default={})
 		self._menu_options['sys-language'] = \
-			Selector(_('Select locale language'), lambda preset: select_locale_lang('en_US',preset), default='en_US')
+			Selector(
+				_('Select locale language'),
+				lambda preset: select_locale_lang(preset),
+				default='en_US')
 		self._menu_options['sys-encoding'] = \
-			Selector(_('Select locale encoding'), lambda preset: select_locale_enc('utf-8',preset), default='utf-8')
+			Selector(
+				_('Select locale encoding'),
+				lambda preset: select_locale_enc(preset),
+				default='UTF-8')
 		self._menu_options['harddrives'] = \
 			Selector(
 				_('Select harddrives'),
-				self._select_harddrives)
+				lambda preset: self._select_harddrives(preset))
 		self._menu_options['disk_layouts'] = \
 			Selector(
 				_('Select disk layout'),
-				lambda x: select_disk_layout(
+				lambda preset: select_disk_layout(
+					preset,
 					storage['arguments'].get('harddrives', []),
 					storage['arguments'].get('advanced', False)
 				),
@@ -112,7 +124,7 @@ class GlobalMenu(GeneralMenu):
 		self._menu_options['profile'] = \
 			Selector(
 				_('Specify profile'),
-				lambda x: self._select_profile(),
+				lambda preset: self._select_profile(preset),
 				display_func=lambda x: x if x else 'None')
 		self._menu_options['audio'] = \
 			Selector(
@@ -247,41 +259,73 @@ class GlobalMenu(GeneralMenu):
 		return ntp
 
 	def _select_harddrives(self, old_harddrives : list) -> list:
-		# old_haddrives = storage['arguments'].get('harddrives', [])
 		harddrives = select_harddrives(old_harddrives)
 
-		# in case the harddrives got changed we have to reset the disk layout as well
-		if old_harddrives != harddrives:
-			self._menu_options.get('disk_layouts').set_current_selection(None)
-			storage['arguments']['disk_layouts'] = {}
-
-		if not harddrives:
+		if len(harddrives) == 0:
 			prompt = _(
 				"You decided to skip harddrive selection\nand will use whatever drive-setup is mounted at {} (experimental)\n"
 				"WARNING: Archinstall won't check the suitability of this setup\n"
 				"Do you wish to continue?"
 			).format(storage['MOUNT_POINT'])
 
-			choice = Menu(prompt, Menu.yes_no(), default_option=Menu.yes()).run()
+			choice = Menu(prompt, Menu.yes_no(), default_option=Menu.yes(), skip=False).run()
 
-			if choice == Menu.no():
-				exit(1)
+			if choice.value == Menu.no():
+				return self._select_harddrives(old_harddrives)
+
+		# in case the harddrives got changed we have to reset the disk layout as well
+		if old_harddrives != harddrives:
+			self._menu_options.get('disk_layouts').set_current_selection(None)
+			storage['arguments']['disk_layouts'] = {}
 
 		return harddrives
 
-	def _select_profile(self):
-		profile = select_profile()
+	def _select_profile(self, preset):
+		profile = select_profile(preset)
+		ret = None
+
+		if profile is None:
+			if any([
+				archinstall.storage.get('profile_minimal', False),
+				archinstall.storage.get('_selected_servers', None),
+				archinstall.storage.get('_desktop_profile', None),
+				archinstall.arguments.get('desktop-environment', None),
+				archinstall.arguments.get('gfx_driver_packages', None)
+			]):
+				return preset
+			else:  # ctrl+c was actioned and all profile settings have been reset
+				return None
+
+		servers = archinstall.storage.get('_selected_servers', [])
+		desktop = archinstall.storage.get('_desktop_profile', None)
+		desktop_env = archinstall.arguments.get('desktop-environment', None)
+		gfx_driver = archinstall.arguments.get('gfx_driver_packages', None)
 
 		# Check the potentially selected profiles preparations to get early checks if some additional questions are needed.
 		if profile and profile.has_prep_function():
 			namespace = f'{profile.namespace}.py'
 			with profile.load_instructions(namespace=namespace) as imported:
-				if imported._prep_function():
-					return profile
+				if imported._prep_function(servers=servers, desktop=desktop, desktop_env=desktop_env, gfx_driver=gfx_driver):
+					ret: Profile = profile
+
+					match ret.name:
+						case 'minimal':
+							reset = ['_selected_servers', '_desktop_profile', 'desktop-environment', 'gfx_driver_packages']
+						case 'server':
+							reset = ['_desktop_profile', 'desktop-environment']
+						case 'desktop':
+							reset = ['_selected_servers']
+						case 'xorg':
+							reset = ['_selected_servers', '_desktop_profile', 'desktop-environment']
+
+					for r in reset:
+						archinstall.storage[r] = None
 				else:
-					return self._select_profile()
+					return self._select_profile(preset)
+		elif profile:
+			ret = profile
 
-		return self._data_store.get('profile', None)
+		return ret
 
 	def _create_superuser_account(self):
 		superusers = ask_for_superuser_account(str(_('Manage superuser accounts: ')))
diff --git a/archinstall/lib/menu/list_manager.py b/archinstall/lib/menu/list_manager.py
index 377d30f2..f3927e6f 100644
--- a/archinstall/lib/menu/list_manager.py
+++ b/archinstall/lib/menu/list_manager.py
@@ -86,7 +86,7 @@ The contents in the base class of this methods serve for a very basic usage, and
 """
 
 from .text_input import TextInput
-from .menu import Menu
+from .menu import Menu, MenuSelectionType
 from os import system
 from copy import copy
 from typing import Union, Any, TYPE_CHECKING, Dict
@@ -167,6 +167,7 @@ class ListManager:
 			options += self.bottom_list
 
 			system('clear')
+
 			target = Menu(
 				self._prompt,
 				options,
@@ -174,27 +175,31 @@ class ListManager:
 				clear_screen=False,
 				clear_menu_on_exit=False,
 				header=self.header,
-				skip_empty_entries=True).run()
+				skip_empty_entries=True
+			).run()
+
+			if target.type_ == MenuSelectionType.Esc:
+				return self.run()
 
-			if not target or target in self.bottom_list:
+			if not target.value or target.value in self.bottom_list:
 				self.action = target
 				break
 
-			if target and target in self._default_action:
-				self.action = target
+			if target.value and target.value in self._default_action:
+				self.action = target.value
 				self.target = None
 				self.exec_action(self._data)
 				continue
 
 			if isinstance(self._data,dict):
-				data_key = data_formatted[target]
+				data_key = data_formatted[target.value]
 				key = self._data[data_key]
 				self.target = {data_key: key}
 			else:
-				self.target = self._data[data_formatted[target]]
+				self.target = self._data[data_formatted[target.value]]
 
 			# Possible enhacement. If run_actions returns false a message line indicating the failure
-			self.run_actions(target)
+			self.run_actions(target.value)
 
 		if not target or target == self.cancel_action:  # TODO dubious
 			return self.base_data  # return the original list
@@ -204,14 +209,18 @@ class ListManager:
 	def run_actions(self,prompt_data=None):
 		options = self.action_list() + self.bottom_item
 		prompt = _("Select an action for < {} >").format(prompt_data if prompt_data else self.target)
-		self.action = Menu(
+		choice = Menu(
 			prompt,
 			options,
 			sort=False,
 			clear_screen=False,
 			clear_menu_on_exit=False,
 			preset_values=self.bottom_item,
-			show_search_hint=False).run()
+			show_search_hint=False
+		).run()
+
+		self.action = choice.value
+
 		if not self.action or self.action == self.cancel_action:
 			return False
 		else:
diff --git a/archinstall/lib/menu/menu.py b/archinstall/lib/menu/menu.py
index b2f4146d..c34814eb 100644
--- a/archinstall/lib/menu/menu.py
+++ b/archinstall/lib/menu/menu.py
@@ -1,4 +1,6 @@
-from typing import Dict, List, Union, Any, TYPE_CHECKING
+from dataclasses import dataclass
+from enum import Enum, auto
+from typing import Dict, List, Union, Any, TYPE_CHECKING, Optional
 
 from archinstall.lib.menu.simple_menu import TerminalMenu
 
@@ -13,6 +15,18 @@ if TYPE_CHECKING:
 	_: Any
 
 
+class MenuSelectionType(Enum):
+	Selection = auto()
+	Esc = auto()
+	Ctrl_c = auto()
+
+
+@dataclass
+class MenuSelection:
+	type_: MenuSelectionType
+	value: Optional[Union[str, List[str]]] = None
+
+
 class Menu(TerminalMenu):
 
 	@classmethod
@@ -33,14 +47,16 @@ class Menu(TerminalMenu):
 		p_options :Union[List[str], Dict[str, Any]],
 		skip :bool = True,
 		multi :bool = False,
-		default_option :str = None,
+		default_option : Optional[str] = None,
 		sort :bool = True,
 		preset_values :Union[str, List[str]] = None,
-		cursor_index :int = None,
+		cursor_index : Optional[int] = None,
 		preview_command=None,
 		preview_size=0.75,
 		preview_title='Info',
 		header :Union[List[str],str] = None,
+		explode_on_interrupt :bool = False,
+		explode_warning :str = '',
 		**kwargs
 	):
 		"""
@@ -80,9 +96,15 @@ class Menu(TerminalMenu):
 		:param preview_title: Title of the preview window
 		:type preview_title: str
 
-		param: header one or more header lines for the menu
+		param header: one or more header lines for the menu
 		type param: string or list
 
+		param explode_on_interrupt: This will explicitly handle a ctrl+c instead and return that specific state
+		type param: bool
+
+		param explode_warning: If explode_on_interrupt is True and this is non-empty, there will be a warning with a user confirmation displayed
+		type param: str
+
 		:param kwargs : any SimpleTerminal parameter
 		"""
 		# we guarantee the inmutability of the options outside the class.
@@ -99,6 +121,8 @@ class Menu(TerminalMenu):
 			log(f"invalid parameter at Menu() call was at <{sys._getframe(1).f_code.co_name}>",level=logging.WARNING)
 			raise RequirementError("Menu() requires an iterable as option.")
 
+		self._default_str = str(_('(default)'))
+
 		if isinstance(p_options,dict):
 			options = list(p_options.keys())
 		else:
@@ -117,27 +141,40 @@ class Menu(TerminalMenu):
 		if sort:
 			options = sorted(options)
 
-		self.menu_options = options
-		self.skip = skip
-		self.default_option = default_option
-		self.multi = multi
+		self._menu_options = options
+		self._skip = skip
+		self._default_option = default_option
+		self._multi = multi
+		self._explode_on_interrupt = explode_on_interrupt
+		self._explode_warning = explode_warning
+
 		menu_title = f'\n{title}\n\n'
+
 		if header:
-			separator = '\n  '
 			if not isinstance(header,(list,tuple)):
-				header = [header,]
-			if skip:
-				menu_title += str(_("Use ESC to skip\n"))
-			menu_title += separator + separator.join(header)
-		elif skip:
-			menu_title += str(_("Use ESC to skip\n\n"))
+				header = [header]
+			header = '\n'.join(header)
+			menu_title += f'\n{header}\n'
+
+		action_info = ''
+		if skip:
+			action_info += str(_("Use ESC to skip"))
+
+		if self._explode_on_interrupt:
+			if len(action_info) > 0:
+				action_info += '\n'
+			action_info += str(_('Use CTRL+C to reset current selection\n\n'))
+
+		menu_title += action_info
+
 		if default_option:
 			# if a default value was specified we move that one
 			# to the top of the list and mark it as default as well
-			default = f'{default_option} (default)'
-			self.menu_options = [default] + [o for o in self.menu_options if default_option != o]
+			default = f'{default_option} {self._default_str}'
+			self._menu_options = [default] + [o for o in self._menu_options if default_option != o]
+
+		self._preselection(preset_values,cursor_index)
 
-		self.preselection(preset_values,cursor_index)
 		cursor = "> "
 		main_menu_cursor_style = ("fg_cyan", "bold")
 		main_menu_style = ("bg_blue", "fg_gray")
@@ -145,8 +182,9 @@ class Menu(TerminalMenu):
 		kwargs['clear_screen'] = kwargs.get('clear_screen',True)
 		kwargs['show_search_hint'] = kwargs.get('show_search_hint',True)
 		kwargs['cycle_cursor'] = kwargs.get('cycle_cursor',True)
+
 		super().__init__(
-			menu_entries=self.menu_options,
+			menu_entries=self._menu_options,
 			title=menu_title,
 			menu_cursor=cursor,
 			menu_cursor_style=main_menu_cursor_style,
@@ -160,32 +198,46 @@ class Menu(TerminalMenu):
 			preview_command=preview_command,
 			preview_size=preview_size,
 			preview_title=preview_title,
+			explode_on_interrupt=self._explode_on_interrupt,
 			multi_select_select_on_accept=False,
 			**kwargs,
 		)
 
-	def _show(self):
-		idx = self.show()
+	def _show(self) -> MenuSelection:
+		try:
+			idx = self.show()
+		except KeyboardInterrupt:
+			return MenuSelection(type_=MenuSelectionType.Ctrl_c)
+
+		def check_default(elem):
+			if self._default_option is not None and f'{self._default_option} {self._default_str}' in elem:
+				return self._default_option
+			else:
+				return elem
+
 		if idx is not None:
 			if isinstance(idx, (list, tuple)):
-				return [self.default_option if ' (default)' in self.menu_options[i] else self.menu_options[i] for i in idx]
+				results = []
+				for i in idx:
+					option = check_default(self._menu_options[i])
+					results.append(option)
+				return MenuSelection(type_=MenuSelectionType.Selection, value=results)
 			else:
-				selected = self.menu_options[idx]
-				if ' (default)' in selected and self.default_option:
-					return self.default_option
-				return selected
+				result = check_default(self._menu_options[idx])
+				return MenuSelection(type_=MenuSelectionType.Selection, value=result)
 		else:
-			if self.default_option:
-				if self.multi:
-					return [self.default_option]
-				else:
-					return self.default_option
-			return None
-
-	def run(self):
+			return MenuSelection(type_=MenuSelectionType.Esc)
+
+	def run(self) -> MenuSelection:
 		ret = self._show()
 
-		if ret is None and not self.skip:
+		if ret.type_ == MenuSelectionType.Ctrl_c:
+			if self._explode_on_interrupt and len(self._explode_warning) > 0:
+				response = Menu(self._explode_warning, Menu.yes_no(), skip=False).run()
+				if response.value == Menu.no():
+					return self.run()
+
+		if ret.type_ is not MenuSelectionType.Selection and not self._skip:
 			return self.run()
 
 		return ret
@@ -200,15 +252,15 @@ class Menu(TerminalMenu):
 		pos = self._menu_entries.index(value)
 		self.set_cursor_pos(pos)
 
-	def preselection(self,preset_values :list = [],cursor_index :int = None):
+	def _preselection(self,preset_values :Union[str, List[str]] = [], cursor_index : Optional[int] = None):
 		def from_preset_to_cursor():
 			if preset_values:
 				# if the value is not extant return 0 as cursor index
 				try:
 					if isinstance(preset_values,str):
-						self.cursor_index = self.menu_options.index(self.preset_values)
+						self.cursor_index = self._menu_options.index(self.preset_values)
 					else:  # should return an error, but this is smoother
-						self.cursor_index = self.menu_options.index(self.preset_values[0])
+						self.cursor_index = self._menu_options.index(self.preset_values[0])
 				except ValueError:
 					self.cursor_index = 0
 
@@ -218,13 +270,13 @@ class Menu(TerminalMenu):
 			return
 
 		self.preset_values = preset_values
-		if self.default_option:
-			if isinstance(preset_values,str) and self.default_option == preset_values:
-				self.preset_values = f"{preset_values} (default)"
-			elif isinstance(preset_values,(list,tuple)) and self.default_option in preset_values:
-				idx = preset_values.index(self.default_option)
-				self.preset_values[idx] = f"{preset_values[idx]} (default)"
-		if cursor_index is None or not self.multi:
+		if self._default_option:
+			if isinstance(preset_values,str) and self._default_option == preset_values:
+				self.preset_values = f"{preset_values} {self._default_str}"
+			elif isinstance(preset_values,(list,tuple)) and self._default_option in preset_values:
+				idx = preset_values.index(self._default_option)
+				self.preset_values[idx] = f"{preset_values[idx]} {self._default_str}"
+		if cursor_index is None or not self._multi:
 			from_preset_to_cursor()
-		if not self.multi: # Not supported by the infraestructure
+		if not self._multi: # Not supported by the infraestructure
 			self.preset_values = None
diff --git a/archinstall/lib/menu/selection_menu.py b/archinstall/lib/menu/selection_menu.py
index e18d92c2..54b1441b 100644
--- a/archinstall/lib/menu/selection_menu.py
+++ b/archinstall/lib/menu/selection_menu.py
@@ -4,7 +4,7 @@ import logging
 import sys
 from typing import Callable, Any, List, Iterator, Tuple, Optional, Dict, TYPE_CHECKING
 
-from .menu import Menu
+from .menu import Menu, MenuSelectionType
 from ..locale_helpers import set_keyboard_language
 from ..output import log
 from ..translation import Translation
@@ -12,13 +12,15 @@ from ..translation import Translation
 if TYPE_CHECKING:
 	_: Any
 
-def select_archinstall_language(default='English'):
+
+def select_archinstall_language(preset_value: str) -> Optional[str]:
 	"""
 	copied from user_interaction/general_conf.py as a temporary measure
 	"""
 	languages = Translation.get_all_names()
-	language = Menu(_('Select Archinstall language'), languages, default_option=default).run()
-	return language
+	language = Menu(_('Select Archinstall language'), languages, preset_values=preset_value).run()
+	return language.value
+
 
 class Selector:
 	def __init__(
@@ -307,22 +309,27 @@ class GeneralMenu:
 				skip_empty_entries=True
 			).run()
 
-			if selection and self.auto_cursor:
-				cursor_pos = menu_options.index(selection) + 1  # before the strip otherwise fails
+			if selection.type_ == MenuSelectionType.Selection:
+				value = selection.value
+
+				if self.auto_cursor:
+					cursor_pos = menu_options.index(value) + 1  # before the strip otherwise fails
+
+					# in case the new position lands on a "placeholder" we'll skip them as well
+					while True:
+						if cursor_pos >= len(menu_options):
+							cursor_pos = 0
+						if len(menu_options[cursor_pos]) > 0:
+							break
+						cursor_pos += 1
 
-				# in case the new position lands on a "placeholder" we'll skip them as well
-				while True:
-					if cursor_pos >= len(menu_options):
-						cursor_pos = 0
-					if len(menu_options[cursor_pos]) > 0:
+					value = value.strip()
+
+					# if this calls returns false, we exit the menu
+					# we allow for an callback for special processing on realeasing control
+					if not self._process_selection(value):
 						break
-					cursor_pos += 1
 
-				selection = selection.strip()
-			if selection:
-				# if this calls returns false, we exit the menu. We allow for an callback for special processing on realeasing control
-				if not self._process_selection(selection):
-					break
 		if not self.is_context_mgr:
 			self.__exit__()
 
@@ -443,15 +450,17 @@ class GeneralMenu:
 	def mandatory_overview(self) -> Tuple[int, int]:
 		mandatory_fields = 0
 		mandatory_waiting = 0
-		for field in self._menu_options:
-			option = self._menu_options[field]
+		for field, option in self._menu_options.items():
 			if option.is_mandatory():
 				mandatory_fields += 1
 				if not option.has_selection():
 					mandatory_waiting += 1
 		return mandatory_fields, mandatory_waiting
 
-	def _select_archinstall_language(self, default_lang):
-		language = select_archinstall_language(default_lang)
-		self._translation.activate(language)
-		return language
+	def _select_archinstall_language(self, preset_value: str) -> str:
+		language = select_archinstall_language(preset_value)
+		if language is not None:
+			self._translation.activate(language)
+			return language
+
+		return preset_value
diff --git a/archinstall/lib/menu/simple_menu.py b/archinstall/lib/menu/simple_menu.py
index a0a241bd..947259eb 100644
--- a/archinstall/lib/menu/simple_menu.py
+++ b/archinstall/lib/menu/simple_menu.py
@@ -596,7 +596,8 @@ class TerminalMenu:
         status_bar: Optional[Union[str, Iterable[str], Callable[[str], str]]] = None,
         status_bar_below_preview: bool = DEFAULT_STATUS_BAR_BELOW_PREVIEW,
         status_bar_style: Optional[Iterable[str]] = DEFAULT_STATUS_BAR_STYLE,
-        title: Optional[Union[str, Iterable[str]]] = None
+        title: Optional[Union[str, Iterable[str]]] = None,
+        explode_on_interrupt: bool = False
     ):
         def extract_shortcuts_menu_entries_and_preview_arguments(
             entries: Iterable[str],
@@ -718,6 +719,7 @@ class TerminalMenu:
         self._search_case_sensitive = search_case_sensitive
         self._search_highlight_style = tuple(search_highlight_style) if search_highlight_style is not None else ()
         self._search_key = search_key
+        self._explode_on_interrupt = explode_on_interrupt
         self._shortcut_brackets_highlight_style = (
             tuple(shortcut_brackets_highlight_style) if shortcut_brackets_highlight_style is not None else ()
         )
@@ -1538,7 +1540,9 @@ class TerminalMenu:
                         # Only append `next_key` if it is a printable character and the first character is not the
                         # `search_start` key
                         self._search.search_text += next_key
-        except KeyboardInterrupt:
+        except KeyboardInterrupt as e:
+            if self._explode_on_interrupt:
+                raise e
             menu_was_interrupted = True
         finally:
             reset_signal_handling()
@@ -1842,6 +1846,12 @@ def get_argumentparser() -> argparse.ArgumentParser:
     )
     parser.add_argument("-t", "--title", action="store", dest="title", help="menu title")
     parser.add_argument(
+        "--explode-on-interrupt",
+        action="store_true",
+        dest="explode_on_interrupt",
+        help="Instead of quitting the menu, this will raise the KeyboardInterrupt Exception",
+    )
+    parser.add_argument(
         "-V", "--version", action="store_true", dest="print_version", help="print the version number and exit"
     )
     parser.add_argument("entries", action="store", nargs="*", help="the menu entries to show")
@@ -1971,6 +1981,7 @@ def main() -> None:
             status_bar_below_preview=args.status_bar_below_preview,
             status_bar_style=args.status_bar_style,
             title=args.title,
+            explode_on_interrupt=args.explode_on_interrupt,
         )
     except (InvalidParameterCombinationError, InvalidStyleError, UnknownMenuEntryError) as e:
         print(str(e), file=sys.stderr)