Fix #1106 (#1119)

* Fix #1106

* flake8

* flake8

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

8 files changed, 72 insertions, 54 deletions

diff --git a/archinstall/lib/menu/global_menu.py b/archinstall/lib/menu/global_menu.py
index afccd45b..63b8c03e 100644
--- a/archinstall/lib/menu/global_menu.py
+++ b/archinstall/lib/menu/global_menu.py
@@ -8,7 +8,6 @@ from ..general import SysCommand, secret
 from ..hardware import has_uefi
 from ..models import NetworkConfiguration
 from ..storage import storage
-from ..output import log
 from ..profiles import is_desktop_profile
 from ..disk import encrypted_partitions
 
@@ -277,11 +276,12 @@ class GlobalMenu(GeneralMenu):
 		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():
-					log(' * Profile\'s preparation requirements was not fulfilled.', fg='red')
-					exit(1)
+				if imported._prep_function():
+					return profile
+				else:
+					return self._select_profile()
 
-		return profile
+		return self._data_store.get('profile', None)
 
 	def _create_superuser_account(self):
 		superusers = ask_for_superuser_account(str(_('Manage superuser accounts: ')))
diff --git a/archinstall/lib/user_interaction/partitioning_conf.py b/archinstall/lib/user_interaction/partitioning_conf.py
index af1d224f..db52d447 100644
--- a/archinstall/lib/user_interaction/partitioning_conf.py
+++ b/archinstall/lib/user_interaction/partitioning_conf.py
@@ -215,7 +215,7 @@ def manage_new_and_existing_partitions(block_device: 'BlockDevice') -> Dict[str,
 					continue
 
 			block_device_struct.update(suggest_single_disk_layout(block_device)[block_device.path])
-						
+
 		elif task is None:
 			return block_device_struct
 		else:
diff --git a/archinstall/lib/user_interaction/system_conf.py b/archinstall/lib/user_interaction/system_conf.py
index 0120fd8a..5d4c80dc 100644
--- a/archinstall/lib/user_interaction/system_conf.py
+++ b/archinstall/lib/user_interaction/system_conf.py
@@ -60,7 +60,7 @@ def select_harddrives(preset: List[str] = []) -> List[str]:
 	return []
 
 
-def select_driver(options: Dict[str, Any] = AVAILABLE_GFX_DRIVERS, force_ask: bool = False) -> str:
+def select_driver(options: Dict[str, Any] = AVAILABLE_GFX_DRIVERS) -> str:
 	"""
 	Some what convoluted function, whose job is simple.
 	Select a graphics driver from a pre-defined set of popular options.
@@ -88,9 +88,8 @@ def select_driver(options: Dict[str, Any] = AVAILABLE_GFX_DRIVERS, force_ask: bo
 				'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) 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()
+		title += _('\n\nSelect a graphics driver or leave blank to install all open-source drivers')
+		arguments['gfx_driver'] = Menu(title, drivers).run()
 
 		if arguments.get('gfx_driver', None) is None:
 			arguments['gfx_driver'] = _("All open-source (default)")
diff --git a/profiles/desktop.py b/profiles/desktop.py
index bd3353e8..977fe2de 100644
--- a/profiles/desktop.py
+++ b/profiles/desktop.py
@@ -1,5 +1,6 @@
 # A desktop environment selector.
 import archinstall
+from archinstall import log
 
 is_top_level_profile = True
 
@@ -38,29 +39,33 @@ __supported__ = [
 ]
 
 
-def _prep_function(*args, **kwargs):
+def _prep_function(*args, **kwargs) -> bool:
 	"""
 	Magic function called by the importing installer
 	before continuing any further. It also avoids executing any
 	other code in this stage. So it's a safe way to ask the user
 	for more input before any other installer steps start.
 	"""
-	desktop = archinstall.Menu('Select your desired desktop environment', __supported__, skip=False).run()
+	desktop = archinstall.Menu('Select your desired desktop environment', __supported__).run()
 
-	# Temporarily store the selected desktop profile
-	# in a session-safe location, since this module will get reloaded
-	# the next time it gets executed.
-	if not archinstall.storage.get('_desktop_profile', None):
-		archinstall.storage['_desktop_profile'] = desktop
-	if not archinstall.arguments.get('desktop-environment', None):
-		archinstall.arguments['desktop-environment'] = desktop
-	profile = archinstall.Profile(None, desktop)
-	# Loading the instructions with a custom namespace, ensures that a __name__ comparison is never triggered.
-	with profile.load_instructions(namespace=f"{desktop}.py") as imported:
-		if hasattr(imported, '_prep_function'):
-			return imported._prep_function()
-		else:
-			print(f"Deprecated (??): {desktop} profile has no _prep_function() anymore")
+	if desktop:
+		# Temporarily store the selected desktop profile
+		# in a session-safe location, since this module will get reloaded
+		# the next time it gets executed.
+		if not archinstall.storage.get('_desktop_profile', None):
+			archinstall.storage['_desktop_profile'] = desktop
+		if not archinstall.arguments.get('desktop-environment', None):
+			archinstall.arguments['desktop-environment'] = desktop
+		profile = archinstall.Profile(None, desktop)
+		# Loading the instructions with a custom namespace, ensures that a __name__ comparison is never triggered.
+		with profile.load_instructions(namespace=f"{desktop}.py") as imported:
+			if hasattr(imported, '_prep_function'):
+				return imported._prep_function()
+			else:
+				log(f"Deprecated (??): {desktop} profile has no _prep_function() anymore")
+				exit(1)
+
+	return False
 
 
 if __name__ == 'desktop':
diff --git a/profiles/i3.py b/profiles/i3.py
index 24956209..3283848e 100644
--- a/profiles/i3.py
+++ b/profiles/i3.py
@@ -27,20 +27,21 @@ def _prep_function(*args, **kwargs):
 
 	supported_configurations = ['i3-wm', 'i3-gaps']
 
-	desktop = archinstall.Menu('Select your desired configuration', supported_configurations, skip=False).run()
-
-	# Temporarily store the selected desktop profile
-	# in a session-safe location, since this module will get reloaded
-	# the next time it gets executed.
-	archinstall.storage['_i3_configuration'] = desktop
-
-	# i3 requires a functioning Xorg installation.
-	profile = archinstall.Profile(None, 'xorg')
-	with profile.load_instructions(namespace='xorg.py') as imported:
-		if hasattr(imported, '_prep_function'):
-			return imported._prep_function()
-		else:
-			print('Deprecated (??): xorg profile has no _prep_function() anymore')
+	desktop = archinstall.Menu('Select your desired configuration', supported_configurations).run()
+
+	if desktop:
+		# Temporarily store the selected desktop profile
+		# in a session-safe location, since this module will get reloaded
+		# the next time it gets executed.
+		archinstall.storage['_i3_configuration'] = desktop
+
+		# i3 requires a functioning Xorg installation.
+		profile = archinstall.Profile(None, 'xorg')
+		with profile.load_instructions(namespace='xorg.py') as imported:
+			if hasattr(imported, '_prep_function'):
+				return imported._prep_function()
+			else:
+				print('Deprecated (??): xorg profile has no _prep_function() anymore')
 
 
 if __name__ == 'i3':
diff --git a/profiles/server.py b/profiles/server.py
index c4f35f7b..bbeece12 100644
--- a/profiles/server.py
+++ b/profiles/server.py
@@ -26,15 +26,18 @@ def _prep_function(*args, **kwargs):
 	Magic function called by the importing installer
 	before continuing any further.
 	"""
-	if not archinstall.storage.get('_selected_servers', None):
-		servers = archinstall.Menu(
-			'Choose which servers to install, if none then a minimal installation wil be done', available_servers,
-			multi=True
-		).run()
-
+	servers = archinstall.Menu(
+		'Choose which servers to install, if none then a minimal installation wil be done',
+		available_servers,
+		preset_values=archinstall.storage.get('_selected_servers', []),
+		multi=True
+	).run()
+
+	if servers:
 		archinstall.storage['_selected_servers'] = servers
+		return True
 
-	return True
+	return False
 
 
 if __name__ == 'server':
diff --git a/profiles/sway.py b/profiles/sway.py
index 32d626d7..e9c71b79 100644
--- a/profiles/sway.py
+++ b/profiles/sway.py
@@ -18,7 +18,9 @@ __packages__ = [
 
 
 def _check_driver() -> bool:
-	if "nvidia" in archinstall.storage.get("gfx_driver_packages", None):
+	packages = archinstall.storage.get("gfx_driver_packages", [])
+
+	if packages and "nvidia" in packages:
 		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':
@@ -34,11 +36,15 @@ 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(force_ask=True)
-	if not _check_driver():
-		return _prep_function(args, kwargs)
+	driver = archinstall.select_driver()
 
-	return True
+	if driver:
+		archinstall.storage["gfx_driver_packages"] = driver
+		if not _check_driver():
+			return _prep_function(args, kwargs)
+		return True
+
+	return False
 
 
 # Ensures that this code only gets executed if executed
diff --git a/profiles/xorg.py b/profiles/xorg.py
index 33d2aa4c..237e4350 100644
--- a/profiles/xorg.py
+++ b/profiles/xorg.py
@@ -25,12 +25,16 @@ def _prep_function(*args, **kwargs):
 	for more input before any other installer steps start.
 	"""
 
-	archinstall.storage["gfx_driver_packages"] = archinstall.select_driver()
+	driver = archinstall.select_driver()
+
+	if driver:
+		archinstall.storage["gfx_driver_packages"] = driver
+		return True
 
 	# TODO: Add language section and/or merge it with the locale selected
 	#       earlier in for instance guided.py installer.
 
-	return True
+	return False
 
 
 # Ensures that this code only gets executed if executed