Make Gfx driver handling saver (#1885)

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

7 files changed, 113 insertions, 66 deletions

diff --git a/archinstall/__init__.py b/archinstall/__init__.py
index 305d4096..ddf81824 100644
--- a/archinstall/__init__.py
+++ b/archinstall/__init__.py
@@ -18,7 +18,7 @@ from .lib import profile
 from .lib import interactions
 from . import default_profiles
 
-from .lib.hardware import SysInfo, AVAILABLE_GFX_DRIVERS
+from .lib.hardware import SysInfo, GfxDriver
 from .lib.installer import Installer, accessibility_tools_in_use
 from .lib.output import FormattedOutput, log, error, debug, warn, info
 from .lib.storage import storage
diff --git a/archinstall/lib/global_menu.py b/archinstall/lib/global_menu.py
index 91ebc6a0..54b30240 100644
--- a/archinstall/lib/global_menu.py
+++ b/archinstall/lib/global_menu.py
@@ -349,7 +349,7 @@ class GlobalMenu(AbstractMenu):
 				output += profile_config.profile.name + '\n'
 
 			if profile_config.gfx_driver:
-				output += str(_('Graphics driver')) + ': ' + profile_config.gfx_driver + '\n'
+				output += str(_('Graphics driver')) + ': ' + profile_config.gfx_driver.value + '\n'
 
 			if profile_config.greeter:
 				output += str(_('Greeter')) + ': ' + profile_config.greeter.value + '\n'
diff --git a/archinstall/lib/hardware.py b/archinstall/lib/hardware.py
index bd153a63..85f903e1 100644
--- a/archinstall/lib/hardware.py
+++ b/archinstall/lib/hardware.py
@@ -1,4 +1,5 @@
 import os
+from enum import Enum
 from functools import cached_property
 from pathlib import Path
 from typing import Optional, Dict, List
@@ -8,43 +9,85 @@ from .general import SysCommand
 from .networking import list_interfaces, enrich_iface_types
 from .output import debug
 
-AVAILABLE_GFX_DRIVERS = {
-	# Sub-dicts are layer-2 options to be selected
-	# and lists are a list of packages to be installed
-	"All open-source (default)": [
-		"mesa",
-		"xf86-video-amdgpu",
-		"xf86-video-ati",
-		"xf86-video-nouveau",
-		"xf86-video-vmware",
-		"libva-mesa-driver",
-		"libva-intel-driver",
-		"intel-media-driver",
-		"vulkan-radeon",
-		"vulkan-intel",
-	],
-	"AMD / ATI (open-source)": [
-		"mesa",
-		"xf86-video-amdgpu",
-		"xf86-video-ati",
-		"libva-mesa-driver",
-		"vulkan-radeon",
-	],
-	"Intel (open-source)": [
-		"mesa",
-		"libva-intel-driver",
-		"intel-media-driver",
-		"vulkan-intel",
-	],
-	"Nvidia (open kernel module for newer GPUs, Turing+)": ["nvidia-open"],
-	"Nvidia (open-source nouveau driver)": [
-		"mesa",
-		"xf86-video-nouveau",
-		"libva-mesa-driver"
-	],
-	"Nvidia (proprietary)": ["nvidia"],
-	"VMware / VirtualBox (open-source)": ["mesa", "xf86-video-vmware"],
-}
+
+class GfxPackage(Enum):
+	IntelMediaDriver = 'intel-media-driver'
+	LibvaIntelDriver = 'libva-intel-driver'
+	LibvaMesaDriver = 'libva-mesa-driver'
+	Mesa = "mesa"
+	Nvidia = 'nvidia'
+	NvidiaOpen = 'nvidia-open'
+	VulkanIntel = 'vulkan-intel'
+	VulkanRadeon = 'vulkan-radeon'
+	Xf86VideoAmdgpu = "xf86-video-amdgpu"
+	Xf86VideoAti = "xf86-video-ati"
+	Xf86VideoNouveau = 'xf86-video-nouveau'
+	Xf86VideoVmware = 'xf86-video-vmware'
+
+
+class GfxDriver(Enum):
+	AllOpenSource = 'All open-source'
+	AmdOpenSource = 'AMD / ATI (open-source)'
+	IntelOpenSource = 'Intel (open-source)'
+	NvidiaOpenKernel = 'Nvidia (open kernel module for newer GPUs, Turing+)'
+	NvidiaOpenSource = 'Nvidia (open-source nouveau driver)'
+	NvidiaProprietary = 'Nvidia (proprietary)'
+	VMOpenSource = 'VMware / VirtualBox (open-source)'
+
+	def is_nvidia(self) -> bool:
+		match self:
+			case GfxDriver.NvidiaProprietary | \
+				GfxDriver.NvidiaOpenSource | \
+				GfxDriver.NvidiaOpenKernel:
+				return True
+			case _:
+				return False
+
+	def packages(self) -> List[GfxPackage]:
+		match self:
+			case GfxDriver.AllOpenSource:
+				return [
+					GfxPackage.Mesa,
+					GfxPackage.Xf86VideoAmdgpu,
+					GfxPackage.Xf86VideoAti,
+					GfxPackage.Xf86VideoNouveau,
+					GfxPackage.Xf86VideoVmware,
+					GfxPackage.LibvaMesaDriver,
+					GfxPackage.LibvaIntelDriver,
+					GfxPackage.IntelMediaDriver,
+					GfxPackage.VulkanRadeon,
+					GfxPackage.VulkanIntel
+				]
+			case GfxDriver.AmdOpenSource:
+				return [
+					GfxPackage.Mesa,
+					GfxPackage.Xf86VideoAmdgpu,
+					GfxPackage.Xf86VideoAti,
+					GfxPackage.LibvaMesaDriver,
+					GfxPackage.VulkanRadeon
+				]
+			case GfxDriver.IntelOpenSource:
+				return [
+					GfxPackage.Mesa,
+					GfxPackage.LibvaIntelDriver,
+					GfxPackage.IntelMediaDriver,
+					GfxPackage.VulkanIntel
+				]
+			case GfxDriver.NvidiaOpenKernel:
+				return [GfxPackage.NvidiaOpen]
+			case GfxDriver.NvidiaOpenSource:
+				return [
+					GfxPackage.Mesa,
+					GfxPackage.Xf86VideoNouveau,
+					GfxPackage.LibvaMesaDriver
+				]
+			case GfxDriver.NvidiaProprietary:
+				return [GfxPackage.Nvidia]
+			case GfxDriver.VMOpenSource:
+				return [
+					GfxPackage.Mesa,
+					GfxPackage.Xf86VideoVmware
+				]
 
 
 class _SysInfo:
diff --git a/archinstall/lib/interactions/system_conf.py b/archinstall/lib/interactions/system_conf.py
index ea7e5989..5b1bc456 100644
--- a/archinstall/lib/interactions/system_conf.py
+++ b/archinstall/lib/interactions/system_conf.py
@@ -1,8 +1,8 @@
 from __future__ import annotations
 
-from typing import List, Any, Dict, TYPE_CHECKING, Optional
+from typing import List, Any, TYPE_CHECKING, Optional
 
-from ..hardware import AVAILABLE_GFX_DRIVERS, SysInfo
+from ..hardware import SysInfo, GfxDriver
 from ..menu import MenuSelectionType, Menu
 from ..models.bootloader import Bootloader
 
@@ -65,7 +65,7 @@ def ask_for_bootloader(preset: Bootloader) -> Bootloader:
 	return preset
 
 
-def select_driver(options: Dict[str, Any] = {}, current_value: Optional[str] = None) -> Optional[str]:
+def select_driver(options: List[GfxDriver] = [], current_value: Optional[GfxDriver] = None) -> Optional[GfxDriver]:
 	"""
 	Some what convoluted function, whose job is simple.
 	Select a graphics driver from a pre-defined set of popular options.
@@ -73,11 +73,10 @@ def select_driver(options: Dict[str, Any] = {}, current_value: Optional[str] = N
 	(The template xorg is for beginner users, not advanced, and should
 	there for appeal to the general public first and edge cases later)
 	"""
-
 	if not options:
-		options = AVAILABLE_GFX_DRIVERS
+		options = [driver for driver in GfxDriver]
 
-	drivers = sorted(list(options.keys()))
+	drivers = sorted([o.value for o in options])
 
 	if drivers:
 		title = ''
@@ -90,13 +89,18 @@ def select_driver(options: Dict[str, Any] = {}, current_value: Optional[str] = N
 
 		title += str(_('\nSelect a graphics driver or leave blank to install all open-source drivers'))
 
-		preset = current_value if current_value else None
-		choice = Menu(title, drivers, preset_values=preset).run()
+		preset = current_value.value if current_value else None
+		choice = Menu(
+			title,
+			drivers,
+			preset_values=preset,
+			default_option=GfxDriver.AllOpenSource.value
+		).run()
 
 		if choice.type_ != MenuSelectionType.Selection:
-			return None
+			return current_value
 
-		return choice.value  # type: ignore
+		return GfxDriver(choice.single_value)
 
 	return current_value
 
diff --git a/archinstall/lib/profile/profile_menu.py b/archinstall/lib/profile/profile_menu.py
index 213466a6..079a9817 100644
--- a/archinstall/lib/profile/profile_menu.py
+++ b/archinstall/lib/profile/profile_menu.py
@@ -4,9 +4,9 @@ from typing import TYPE_CHECKING, Any, Optional, Dict
 
 from archinstall.default_profiles.profile import Profile, GreeterType
 from .profile_model import ProfileConfiguration
-from ..hardware import AVAILABLE_GFX_DRIVERS
 from ..menu import Menu, MenuSelectionType, AbstractSubMenu, Selector
 from ..interactions.system_conf import select_driver
+from ..hardware import GfxDriver
 
 if TYPE_CHECKING:
 	_: Any
@@ -38,7 +38,7 @@ class ProfileMenu(AbstractSubMenu):
 		self._menu_options['gfx_driver'] = Selector(
 			_('Graphics driver'),
 			lambda preset: self._select_gfx_driver(preset),
-			display_func=lambda x: x if x else None,
+			display_func=lambda x: x.value if x else None,
 			dependencies=['profile'],
 			default=self._preset.gfx_driver if self._preset.profile and self._preset.profile.is_graphic_driver_supported() else None,
 			enabled=self._preset.profile.is_graphic_driver_supported() if self._preset.profile else False
@@ -73,7 +73,7 @@ class ProfileMenu(AbstractSubMenu):
 				self._menu_options['gfx_driver'].set_current_selection(None)
 			else:
 				self._menu_options['gfx_driver'].set_enabled(True)
-				self._menu_options['gfx_driver'].set_current_selection('All open-source (default)')
+				self._menu_options['gfx_driver'].set_current_selection(GfxDriver.AllOpenSource)
 
 			if not profile.is_greeter_supported():
 				self._menu_options['greeter'].set_enabled(False)
@@ -87,7 +87,7 @@ class ProfileMenu(AbstractSubMenu):
 
 		return profile
 
-	def _select_gfx_driver(self, preset: Optional[str] = None) -> Optional[str]:
+	def _select_gfx_driver(self, preset: Optional[GfxDriver] = None) -> Optional[GfxDriver]:
 		driver = preset
 		profile: Optional[Profile] = self._menu_options['profile'].current_selection
 
@@ -96,11 +96,8 @@ class ProfileMenu(AbstractSubMenu):
 				driver = select_driver(current_value=preset)
 
 			if driver and 'Sway' in profile.current_selection_names():
-				packages = AVAILABLE_GFX_DRIVERS[driver]
-
-				if packages and "nvidia" in packages:
-					prompt = str(
-						_('The proprietary Nvidia driver is not supported by Sway. It is likely that you will run into issues, are you okay with that?'))
+				if driver.is_nvidia():
+					prompt = str(_('The proprietary Nvidia driver is not supported by Sway. It is likely that you will run into issues, are you okay with that?'))
 					choice = Menu(prompt, Menu.yes_no(), default_option=Menu.no(), skip=False).run()
 
 					if choice.value == Menu.no():
diff --git a/archinstall/lib/profile/profile_model.py b/archinstall/lib/profile/profile_model.py
index ad3015ae..2b52073a 100644
--- a/archinstall/lib/profile/profile_model.py
+++ b/archinstall/lib/profile/profile_model.py
@@ -3,6 +3,7 @@ from __future__ import annotations
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any, Optional, Dict
 
+from ..hardware import GfxDriver
 from archinstall.default_profiles.profile import Profile, GreeterType
 
 if TYPE_CHECKING:
@@ -12,14 +13,14 @@ if TYPE_CHECKING:
 @dataclass
 class ProfileConfiguration:
 	profile: Optional[Profile] = None
-	gfx_driver: Optional[str] = None
+	gfx_driver: Optional[GfxDriver] = None
 	greeter: Optional[GreeterType] = None
 
 	def json(self) -> Dict[str, Any]:
 		from .profiles_handler import profile_handler
 		return {
 			'profile': profile_handler.to_json(self.profile),
-			'gfx_driver': self.gfx_driver,
+			'gfx_driver': self.gfx_driver.value if self.gfx_driver else None,
 			'greeter': self.greeter.value if self.greeter else None
 		}
 
@@ -27,9 +28,10 @@ class ProfileConfiguration:
 	def parse_arg(cls, arg: Dict[str, Any]) -> 'ProfileConfiguration':
 		from .profiles_handler import profile_handler
 		greeter = arg.get('greeter', None)
+		gfx_driver = arg.get('gfx_driver', None)
 
 		return ProfileConfiguration(
 			profile_handler.parse_profile_config(arg['profile']),
-			arg.get('gfx_driver', None),
+			GfxDriver(gfx_driver) if gfx_driver else None,
 			GreeterType(greeter) if greeter else None
 		)
diff --git a/archinstall/lib/profile/profiles_handler.py b/archinstall/lib/profile/profiles_handler.py
index 2cc15d8e..4e7c3d2b 100644
--- a/archinstall/lib/profile/profiles_handler.py
+++ b/archinstall/lib/profile/profiles_handler.py
@@ -11,7 +11,7 @@ from typing import List, TYPE_CHECKING, Any, Optional, Dict, Union
 
 from archinstall.default_profiles.profile import Profile, TProfile, GreeterType
 from .profile_model import ProfileConfiguration
-from ..hardware import AVAILABLE_GFX_DRIVERS
+from ..hardware import GfxDriver, GfxPackage
 from ..menu import MenuSelectionType, Menu, MenuSelection
 from ..networking import list_interfaces, fetch_data_from_url
 from ..output import error, debug, info, warn
@@ -188,17 +188,18 @@ class ProfileHandler:
 		if service:
 			install_session.enable_service(service)
 
-	def install_gfx_driver(self, install_session: 'Installer', driver: str):
+	def install_gfx_driver(self, install_session: 'Installer', driver: Optional[GfxDriver]):
 		try:
-			driver_pkgs = AVAILABLE_GFX_DRIVERS[driver] if driver else []
-			additional_pkg = ' '.join(['xorg-server', 'xorg-xinit'] + driver_pkgs)
+			driver_pkgs = driver.packages() if driver else []
+			pkg_names = [p.value for p in driver_pkgs]
+			additional_pkg = ' '.join(['xorg-server', 'xorg-xinit'] + pkg_names)
 
 			if driver is not None:
 				# Find the intersection between the set of known nvidia drivers
 				# and the selected driver packages. Since valid intesections can
 				# only have one element or none, we iterate and try to take the
 				# first element.
-				if driver_pkg := next(iter({'nvidia','nvidia-open'} & set(driver_pkgs)), None):
+				if driver_pkg := next(iter({GfxPackage.Nvidia, GfxPackage.NvidiaOpen} & set(driver_pkgs)), None):
 					if any(kernel in install_session.base_packages for kernel in ("linux-lts", "linux-zen")):
 						for kernel in install_session.kernels:
 							# Fixes https://github.com/archlinux/archinstall/issues/585