Migrate old input to new menu (#874)

* Migrate old input to new menu

* Fix imports

* Remove imports

* Update

* Fixed import by changing 'import archinstall', to 'from ..menu import Menu' and use Menu() directly

* Converted archinstall.<thing> to from ..where import <thing>. This enables us to use archinstall as a module, a git repository in testing and other things without having to install archinstall as an actual module.

Co-authored-by: Daniel Girtler <girtler.daniel@gmail.com>
Co-authored-by: Anton Hvornum <anton@hvornum.se>

6 files changed, 116 insertions, 67 deletions

diff --git a/archinstall/lib/disk/user_guides.py b/archinstall/lib/disk/user_guides.py
index 25db14ea..8acb8cd2 100644
--- a/archinstall/lib/disk/user_guides.py
+++ b/archinstall/lib/disk/user_guides.py
@@ -1,6 +1,7 @@
 from __future__ import annotations
 import logging
 from typing import Optional, Dict, Any, List, TYPE_CHECKING
+
 # https://stackoverflow.com/a/39757388/929999
 if TYPE_CHECKING:
 	from .blockdevice import BlockDevice
@@ -8,6 +9,7 @@ if TYPE_CHECKING:
 from .helpers import sort_block_devices_based_on_performance, select_largest_device, select_disk_larger_than_or_close_to
 from ..hardware import has_uefi
 from ..output import log
+from ..menu import Menu
 
 def suggest_single_disk_layout(block_device :BlockDevice,
 	default_filesystem :Optional[str] = None,
@@ -22,7 +24,9 @@ def suggest_single_disk_layout(block_device :BlockDevice,
 	using_home_partition = False
 
 	if default_filesystem == 'btrfs':
-		using_subvolumes = input('Would you like to use BTRFS subvolumes with a default structure? (Y/n): ').strip().lower() in ('', 'y', 'yes')
+		prompt = 'Would you like to use BTRFS subvolumes with a default structure?'
+		choice = Menu(prompt, ['yes', 'no'], skip=False, default_option='yes').run()
+		using_subvolumes = choice == 'yes'
 
 	layout = {
 		block_device.path : {
@@ -76,7 +80,9 @@ def suggest_single_disk_layout(block_device :BlockDevice,
 		layout[block_device.path]['partitions'][-1]['start'] = '513MiB'
 
 	if not using_subvolumes and block_device.size >= MIN_SIZE_TO_ALLOW_HOME_PART:
-		using_home_partition = input('Would you like to create a separate partition for /home? (Y/n): ').strip().lower() in ('', 'y', 'yes')
+		prompt = 'Would you like to create a separate partition for /home?'
+		choice = Menu(prompt, ['yes', 'no'], skip=False, default_option='yes').run()
+		using_home_partition = choice == 'yes'
 
 	# Set a size for / (/root)
 	if using_subvolumes or block_device.size < MIN_SIZE_TO_ALLOW_HOME_PART or not using_home_partition:
diff --git a/archinstall/lib/menu/selection_menu.py b/archinstall/lib/menu/selection_menu.py
index 8128fefc..940941be 100644
--- a/archinstall/lib/menu/selection_menu.py
+++ b/archinstall/lib/menu/selection_menu.py
@@ -1,8 +1,32 @@
 import sys
 
-import archinstall
-from archinstall import Menu
-
+from .menu import Menu
+from ..general import SysCommand
+from ..storage import storage
+from ..output import log
+from ..profiles import is_desktop_profile
+from ..disk import encrypted_partitions
+from ..locale_helpers import set_keyboard_language
+from ..user_interaction import get_password
+from ..user_interaction import ask_ntp
+from ..user_interaction import ask_for_swap
+from ..user_interaction import ask_for_bootloader
+from ..user_interaction import ask_hostname
+from ..user_interaction import ask_for_audio_selection
+from ..user_interaction import ask_additional_packages_to_install
+from ..user_interaction import ask_to_configure_network
+from ..user_interaction import ask_for_a_timezone
+from ..user_interaction import ask_for_superuser_account
+from ..user_interaction import ask_for_additional_users
+from ..user_interaction import select_language
+from ..user_interaction import select_mirror_regions
+from ..user_interaction import select_locale_lang
+from ..user_interaction import select_locale_enc
+from ..user_interaction import select_disk_layout
+from ..user_interaction import select_kernel
+from ..user_interaction import select_encrypted_partitions
+from ..user_interaction import select_harddrives
+from ..user_interaction import select_profile
 
 class Selector:
 	def __init__(
@@ -109,17 +133,17 @@ class GlobalMenu:
 
 	def _setup_selection_menu_options(self):
 		self._menu_options['keyboard-layout'] = \
-			Selector('Select keyboard layout', lambda: archinstall.select_language('us'), default='us')
+			Selector('Select keyboard layout', lambda: select_language('us'), default='us')
 		self._menu_options['mirror-region'] = \
 			Selector(
 				'Select mirror region',
-				lambda: archinstall.select_mirror_regions(),
+				lambda: select_mirror_regions(),
 				display_func=lambda x: list(x.keys()) if x else '[]',
 				default={})
 		self._menu_options['sys-language'] = \
-			Selector('Select locale language', lambda: archinstall.select_locale_lang('en_US'), default='en_US')
+			Selector('Select locale language', lambda: select_locale_lang('en_US'), default='en_US')
 		self._menu_options['sys-encoding'] = \
-			Selector('Select locale encoding', lambda: archinstall.select_locale_enc('utf-8'), default='utf-8')
+			Selector('Select locale encoding', lambda: select_locale_enc('utf-8'), default='utf-8')
 		self._menu_options['harddrives'] = \
 			Selector(
 				'Select harddrives',
@@ -127,28 +151,28 @@ class GlobalMenu:
 		self._menu_options['disk_layouts'] = \
 			Selector(
 				'Select disk layout',
-				lambda: archinstall.select_disk_layout(
-					archinstall.arguments['harddrives'],
-					archinstall.arguments.get('advanced', False)
+				lambda: select_disk_layout(
+					storage['arguments'].get('harddrives', []),
+					storage['arguments'].get('advanced', False)
 				),
 				dependencies=['harddrives'])
 		self._menu_options['!encryption-password'] = \
 			Selector(
 				'Set encryption password',
-				lambda: archinstall.get_password(prompt='Enter disk encryption password (leave blank for no encryption): '),
+				lambda: get_password(prompt='Enter disk encryption password (leave blank for no encryption): '),
 				display_func=lambda x: self._secret(x) if x else 'None',
 				dependencies=['harddrives'])
 		self._menu_options['swap'] = \
 			Selector(
 				'Use swap',
-				lambda: archinstall.ask_for_swap(),
+				lambda: ask_for_swap(),
 				default=True)
 		self._menu_options['bootloader'] = \
 			Selector(
 				'Select bootloader',
-				lambda: archinstall.ask_for_bootloader(archinstall.arguments.get('advanced', False)),)
+				lambda: ask_for_bootloader(storage['arguments'].get('advanced', False)),)
 		self._menu_options['hostname'] = \
-			Selector('Specify hostname', lambda: archinstall.ask_hostname())
+			Selector('Specify hostname', lambda: ask_hostname())
 		self._menu_options['!root-password'] = \
 			Selector(
 				'Set root password',
@@ -174,29 +198,29 @@ class GlobalMenu:
 		self._menu_options['audio'] = \
 			Selector(
 				'Select audio',
-				lambda: archinstall.ask_for_audio_selection(archinstall.is_desktop_profile(archinstall.arguments.get('profile', None))))
+				lambda: ask_for_audio_selection(is_desktop_profile(storage['arguments'].get('profile', None))))
 		self._menu_options['kernels'] = \
 			Selector(
 				'Select kernels',
-				lambda: archinstall.select_kernel(),
+				lambda: select_kernel(),
 				default=['linux'])
 		self._menu_options['packages'] = \
 			Selector(
 				'Additional packages to install',
-				lambda: archinstall.ask_additional_packages_to_install(archinstall.arguments.get('packages', None)),
+				lambda: ask_additional_packages_to_install(storage['arguments'].get('packages', None)),
 				default=[])
 		self._menu_options['nic'] = \
 			Selector(
 				'Configure network',
-				lambda: archinstall.ask_to_configure_network(),
+				lambda: ask_to_configure_network(),
 				display_func=lambda x: x if x else 'Not configured, unavailable unless setup manually',
 				default={})
 		self._menu_options['timezone'] = \
-			Selector('Select timezone', lambda: archinstall.ask_for_a_timezone())
+			Selector('Select timezone', lambda: ask_for_a_timezone())
 		self._menu_options['ntp'] = \
 			Selector(
 				'Set automatic time sync (NTP)',
-				lambda: archinstall.ask_ntp(),
+				lambda: self._select_ntp(),
 				default=True)
 		self._menu_options['install'] = \
 			Selector(
@@ -205,7 +229,7 @@ class GlobalMenu:
 		self._menu_options['abort'] = Selector('Abort', enabled=True)
 
 	def enable(self, selector_name, omit_if_set=False):
-		arg = archinstall.arguments.get(selector_name, None)
+		arg = storage['arguments'].get(selector_name, None)
 
 		# don't display the menu option if it was defined already
 		if arg is not None and omit_if_set:
@@ -239,8 +263,8 @@ class GlobalMenu:
 					self._process_selection(selection)
 		for key in self._menu_options:
 			sel = self._menu_options[key]
-			if key not in archinstall.arguments:
-				archinstall.arguments[key] = sel._current_selection
+			if key not in storage['arguments']:
+				storage['arguments'][key] = sel._current_selection
 		self._post_processing()
 
 	def _process_selection(self, selection):
@@ -254,7 +278,7 @@ class GlobalMenu:
 		selector = option[0][1]
 		result = selector.func()
 		self._menu_options[selector_name].set_current_selection(result)
-		archinstall.arguments[selector_name] = result
+		storage['arguments'][selector_name] = result
 
 		self._update_install()
 
@@ -263,12 +287,12 @@ class GlobalMenu:
 		self._menu_options.get('install').update_description(text)
 
 	def _post_processing(self):
-		if archinstall.arguments.get('harddrives', None) and archinstall.arguments.get('!encryption-password', None):
+		if storage['arguments'].get('harddrives', None) and storage['arguments'].get('!encryption-password', None):
 			# If no partitions was marked as encrypted, but a password was supplied and we have some disks to format..
 			# Then we need to identify which partitions to encrypt. This will default to / (root).
-			if len(list(archinstall.encrypted_partitions(archinstall.arguments['disk_layouts']))) == 0:
-				archinstall.arguments['disk_layouts'] = archinstall.select_encrypted_partitions(
-					archinstall.arguments['disk_layouts'], archinstall.arguments['!encryption-password'])
+			if len(list(encrypted_partitions(storage['arguments'].get('disk_layouts', [])))) == 0:
+				storage['arguments']['disk_layouts'] = select_encrypted_partitions(
+					storage['arguments']['disk_layouts'], storage['arguments']['!encryption-password'])
 
 	def _install_text(self):
 		missing = self._missing_configs()
@@ -301,27 +325,37 @@ class GlobalMenu:
 
 	def _set_root_password(self):
 		prompt = 'Enter root password (leave blank to disable root & create superuser): '
-		password = archinstall.get_password(prompt=prompt)
+		password = get_password(prompt=prompt)
 
+		# TODO: Do we really wanna wipe the !superusers and !users if root password is set?
+		# What if they set a superuser first, but then decides to set a root password?
 		if password is not None:
 			self._menu_options.get('!superusers').set_current_selection(None)
-			archinstall.arguments['!users'] = {}
-			archinstall.arguments['!superusers'] = {}
+			storage['arguments']['!users'] = {}
+			storage['arguments']['!superusers'] = {}
 
 		return password
 
+	def _select_ntp(self) -> bool:
+		ntp = ask_ntp()
+
+		value = str(ntp).lower()
+		SysCommand(f'timedatectl set-ntp {value}')
+
+		return ntp
+
 	def _select_harddrives(self):
-		old_haddrives = archinstall.arguments.get('harddrives')
-		harddrives = archinstall.select_harddrives()
+		old_haddrives = storage['arguments'].get('harddrives', [])
+		harddrives = select_harddrives()
 
 		# in case the harddrives got changed we have to reset the disk layout as well
 		if old_haddrives != harddrives:
 			self._menu_options.get('disk_layouts').set_current_selection(None)
-			archinstall.arguments['disk_layouts'] = {}
+			storage['arguments']['disk_layouts'] = {}
 
 		if not harddrives:
 			prompt = 'You decided to skip harddrive selection\n'
-			prompt += f"and will use whatever drive-setup is mounted at {archinstall.storage['MOUNT_POINT']} (experimental)\n"
+			prompt += f"and will use whatever drive-setup is mounted at {storage['MOUNT_POINT']} (experimental)\n"
 			prompt += "WARNING: Archinstall won't check the suitability of this setup\n"
 
 			prompt += 'Do you wish to continue?'
@@ -336,36 +370,33 @@ class GlobalMenu:
 		return '*' * len(x)
 
 	def _select_profile(self):
-		profile = archinstall.select_profile()
+		profile = select_profile()
 
 		# 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 not imported._prep_function():
-					archinstall.log(' * Profile\'s preparation requirements was not fulfilled.', fg='red')
+					log(' * Profile\'s preparation requirements was not fulfilled.', fg='red')
 					exit(1)
 
 		return profile
 
 	def _create_superuser_account(self):
-		superuser = archinstall.ask_for_superuser_account('Create a required super-user with sudo privileges: ', forced=True)
+		superuser = ask_for_superuser_account('Create a required super-user with sudo privileges: ', forced=True)
 		return superuser
 
 	def _create_user_account(self):
-		users, superusers = archinstall.ask_for_additional_users('Enter a username to create an additional user: ')
-		if not archinstall.arguments.get('!superusers', None):
-			archinstall.arguments['!superusers'] = superusers
-		else:
-			archinstall.arguments['!superusers'] = {**archinstall.arguments['!superusers'], **superusers}
+		users, superusers = ask_for_additional_users('Enter a username to create an additional user: ')
+		storage['arguments']['!superusers'] = {**storage['arguments'].get('!superusers', {}), **superusers}
 
 		return users
 
 	def _set_kb_language(self):
 		# Before continuing, set the preferred keyboard layout/language in the current terminal.
 		# This will just help the user with the next following questions.
-		if archinstall.arguments.get('keyboard-layout', None) and len(archinstall.arguments['keyboard-layout']):
-			archinstall.set_keyboard_language(archinstall.arguments['keyboard-layout'])
+		if len(storage['arguments'].get('keyboard-layout', [])):
+			set_keyboard_language(storage['arguments']['keyboard-layout'])
 
 	def _verify_selection_enabled(self, selection_name):
 		if selection := self._menu_options.get(selection_name, None):
diff --git a/archinstall/lib/output.py b/archinstall/lib/output.py
index 3ef1e234..ffe00370 100644
--- a/archinstall/lib/output.py
+++ b/archinstall/lib/output.py
@@ -106,4 +106,4 @@ def log(*args :str, **kwargs :Union[str, int, Dict[str, Union[str, int]]]) -> No
 	# We use sys.stdout.write()+flush() instead of print() to try and
 	# fix issue #94
 	sys.stdout.write(f"{string}\n")
-	sys.stdout.flush()
\ No newline at end of file
+	sys.stdout.flush()
diff --git a/archinstall/lib/user_interaction.py b/archinstall/lib/user_interaction.py
index 34ce5534..8cc7de0a 100644
--- a/archinstall/lib/user_interaction.py
+++ b/archinstall/lib/user_interaction.py
@@ -79,8 +79,9 @@ def do_countdown() -> bool:
 			print(".", end='')
 
 		if SIG_TRIGGER:
-			abort = input('\nDo you really want to abort (y/n)? ')
-			if abort.strip() != 'n':
+			prompt = 'Do you really want to abort'
+			choice = Menu(prompt, ['yes', 'no'], skip=False).run()
+			if choice == 'yes':
 				exit(0)
 
 			if SIG_TRIGGER is False:
@@ -271,7 +272,7 @@ def ask_for_swap(prompt='Would you like to use swap on zram?', forced=False):
 	return False if choice == 'no' else True
 
 
-def ask_ntp():
+def ask_ntp() -> bool:
 	prompt = 'Would you like to use automatic time synchronization (NTP) with the default time servers?'
 	prompt += 'Hardware time and other post-configuration steps might be required in order for NTP to work. For more information, please check the Arch wiki'
 	choice = Menu(prompt, ['yes', 'no'], skip=False, default_option='yes').run()
@@ -859,7 +860,7 @@ def select_harddrives() -> Optional[str]:
 	return []
 
 
-def select_driver(options :Dict[str, Any] = AVAILABLE_GFX_DRIVERS) -> str:
+def select_driver(options :Dict[str, Any] = AVAILABLE_GFX_DRIVERS, force_ask :bool = False) -> str:
 	"""
 	Some what convoluted function, whose job is simple.
 	Select a graphics driver from a pre-defined set of popular options.
@@ -881,7 +882,7 @@ def select_driver(options :Dict[str, Any] = AVAILABLE_GFX_DRIVERS) -> str:
 		if has_nvidia_graphics():
 			title += 'For the best compatibility with your Nvidia hardware, you may want to use the Nvidia proprietary driver.\n'
 
-		if not arguments.get('gfx_driver', None):
+		if not arguments.get('gfx_driver', None) or force_ask:
 			title += '\n\nSelect a graphics driver or leave blank to install all open-source drivers'
 			arguments['gfx_driver'] = Menu(title, drivers).run()
 
diff --git a/examples/guided.py b/examples/guided.py
index 45e213df..9cc39c86 100644
--- a/examples/guided.py
+++ b/examples/guided.py
@@ -58,15 +58,16 @@ def ask_user_questions():
 		will we continue with the actual installation steps.
 	"""
 
+	# ref: https://github.com/archlinux/archinstall/pull/831
+	# we'll set NTP to true by default since this is also
+	# the default value specified in the menu options; in
+	# case it will be changed by the user we'll also update
+	# the system immediately
+	archinstall.SysCommand('timedatectl set-ntp true')
+
 	global_menu = archinstall.GlobalMenu()
 	global_menu.enable('keyboard-layout')
 
-	if not archinstall.arguments.get('ntp', False):
-		archinstall.arguments['ntp'] = input("Would you like to use automatic time synchronization (NTP) with the default time servers? [Y/n]: ").strip().lower() in ('y', 'yes', '')
-		if archinstall.arguments['ntp']:
-			archinstall.log("Hardware time and other post-configuration steps might be required in order for NTP to work. For more information, please check the Arch wiki.", fg="yellow")
-			archinstall.SysCommand('timedatectl set-ntp true')
-
 	# Set which region to download packages from during the installation
 	global_menu.enable('mirror-region')
 
@@ -293,8 +294,9 @@ def perform_installation(mountpoint):
 
 		installation.log("For post-installation tips, see https://wiki.archlinux.org/index.php/Installation_guide#Post-installation", fg="yellow")
 		if not archinstall.arguments.get('silent'):
-			choice = input("Would you like to chroot into the newly created installation and perform post-installation configuration? [Y/n] ")
-			if choice.lower() in ("y", ""):
+			prompt = 'Would you like to chroot into the newly created installation and perform post-installation configuration?'
+			choice = archinstall.Menu(prompt, ['yes', 'no'], default_option='yes').run()
+			if choice == 'yes':
 				try:
 					installation.drop_to_shell()
 				except:
diff --git a/profiles/sway.py b/profiles/sway.py
index b0ff8c19..27905e33 100644
--- a/profiles/sway.py
+++ b/profiles/sway.py
@@ -1,5 +1,4 @@
 # A desktop environment using "Sway"
-
 import archinstall
 
 is_top_level_profile = False
@@ -18,6 +17,16 @@ __packages__ = [
 ]
 
 
+def _check_driver() -> bool:
+	if "nvidia" in archinstall.storage.get("gfx_driver_packages", None):
+		prompt = 'The proprietary Nvidia driver is not supported by Sway. It is likely that you will run into issues, are you okay with that?'
+		choice = archinstall.Menu(prompt, ['yes', 'no'], default_option='no').run()
+		if choice == 'no':
+			return False
+
+	return True
+
+
 def _prep_function(*args, **kwargs):
 	"""
 	Magic function called by the importing installer
@@ -25,7 +34,9 @@ def _prep_function(*args, **kwargs):
 	other code in this stage. So it's a safe way to ask the user
 	for more input before any other installer steps start.
 	"""
-	archinstall.storage["gfx_driver_packages"] = archinstall.select_driver()
+	archinstall.storage["gfx_driver_packages"] = archinstall.select_driver(force_ask=True)
+	if not _check_driver():
+		return _prep_function(args, kwargs)
 
 	return True
 
@@ -34,10 +45,8 @@ def _prep_function(*args, **kwargs):
 # through importlib.util.spec_from_file_location("sway", "/somewhere/sway.py")
 # or through conventional import sway
 if __name__ == "sway":
-	if "nvidia" in archinstall.storage.get("gfx_driver_packages", None):
-		choice = input("The proprietary Nvidia driver is not supported by Sway. It is likely that you will run into issues. Continue anyways? [y/N] ")
-		if choice.lower() in ("n", ""):
-			raise archinstall.lib.exceptions.HardwareIncompatibilityError("Sway does not support the proprietary nvidia drivers.")
+	if not _check_driver():
+		raise archinstall.lib.exceptions.HardwareIncompatibilityError("Sway does not support the proprietary nvidia drivers.")
 
 	# Install the Sway packages
 	archinstall.storage['installation_session'].add_additional_packages(__packages__)