From 9f5c2bb70b0a4551eaa871164a3c9d71c1e65086 Mon Sep 17 00:00:00 2001
From: codefiles <11915375+codefiles@users.noreply.github.com>
Date: Fri, 29 Sep 2023 10:09:28 -0400
Subject: Add support for ESP partition flag (#2133)

---
 archinstall/lib/disk/device_model.py | 18 +++++++++++-------
 1 file changed, 11 insertions(+), 7 deletions(-)

(limited to 'archinstall/lib/disk/device_model.py')

diff --git a/archinstall/lib/disk/device_model.py b/archinstall/lib/disk/device_model.py
index b1f012f7..6992bccb 100644
--- a/archinstall/lib/disk/device_model.py
+++ b/archinstall/lib/disk/device_model.py
@@ -658,7 +658,8 @@ class PartitionModification:
 	partuuid: Optional[str] = None
 	uuid: Optional[str] = None
 
-	_boot_indicator_flags = [PartitionFlag.Boot, PartitionFlag.XBOOTLDR]
+	_efi_indicator_flags = (PartitionFlag.Boot, PartitionFlag.ESP)
+	_boot_indicator_flags = (PartitionFlag.Boot, PartitionFlag.XBOOTLDR)
 
 	def __post_init__(self):
 		# needed to use the object as a dictionary key due to hash func
@@ -728,6 +729,13 @@ class PartitionModification:
 
 		raise ValueError('Mountpoint is not specified')
 
+	def is_efi(self) -> bool:
+		return (
+			any(set(self.flags) & set(self._efi_indicator_flags))
+			and self.fs_type == FilesystemType.Fat32
+			and PartitionFlag.XBOOTLDR not in self.flags
+		)
+
 	def is_boot(self) -> bool:
 		"""
 		Returns True if any of the boot indicator flags are found in self.flags
@@ -828,9 +836,8 @@ class DeviceModification:
 	def get_efi_partition(self) -> Optional[PartitionModification]:
 		"""
 		Similar to get_boot_partition() but excludes XBOOTLDR partitions from it's candidates.
-		Also works with ESP flag.
 		"""
-		filtered = filter(lambda x: (x.is_boot() or PartitionFlag.ESP in x.flags) and x.fs_type == FilesystemType.Fat32 and PartitionFlag.XBOOTLDR not in x.flags, self.partitions)
+		filtered = filter(lambda x: x.is_efi() and x.mountpoint, self.partitions)
 		return next(filtered, None)
 
 	def get_boot_partition(self) -> Optional[PartitionModification]:
@@ -843,10 +850,7 @@ class DeviceModification:
 			filtered = filter(lambda x: x.is_boot() and x != efi_partition and x.mountpoint, self.partitions)
 			if boot_partition := next(filtered, None):
 				return boot_partition
-			if efi_partition.is_boot():
-				return efi_partition
-			else:
-				return None
+			return efi_partition
 		else:
 			filtered = filter(lambda x: x.is_boot() and x.mountpoint, self.partitions)
 			return next(filtered, None)
-- 
cgit v1.2.3-70-g09d2