From a02213484d7461b67f5d83e095ccffb54c54a2fe Mon Sep 17 00:00:00 2001
From: codefiles <11915375+codefiles@users.noreply.github.com>
Date: Fri, 17 Jan 2025 20:20:53 -0500
Subject: [PATCH] disk: add support for creating swap partitions

---
 archinstall/lib/disk/device_handler.py    | 22 +++++++++++++++++++---
 archinstall/lib/disk/device_model.py      |  9 +++++++++
 archinstall/lib/disk/encryption_menu.py   |  5 ++++-
 archinstall/lib/disk/partitioning_menu.py |  4 +++-
 archinstall/lib/installer.py              | 10 +++++++---
 5 files changed, 42 insertions(+), 8 deletions(-)

diff --git a/archinstall/lib/disk/device_handler.py b/archinstall/lib/disk/device_handler.py
index 4d54314d56..d3e658cca0 100644
--- a/archinstall/lib/disk/device_handler.py
+++ b/archinstall/lib/disk/device_handler.py
@@ -151,6 +151,8 @@ def _determine_fs_type(
 	) -> FilesystemType | None:
 		try:
 			if partition.fileSystem:
+				if partition.fileSystem.type == FilesystemType.LinuxSwap.parted_value:
+					return FilesystemType.LinuxSwap
 				return FilesystemType(partition.fileSystem.type)
 			elif lsblk_info is not None:
 				return FilesystemType(lsblk_info.fstype) if lsblk_info.fstype else None
@@ -259,6 +261,7 @@ def format(
 		additional_parted_options: list[str] = []
 	) -> None:
 		mkfs_type = fs_type.value
+		command = None
 		options = []
 
 		match fs_type:
@@ -277,10 +280,15 @@ def format(
 				options.append('--fast')
 			case FilesystemType.Reiserfs:
 				pass
+			case FilesystemType.LinuxSwap:
+				command = "mkswap"
 			case _:
 				raise UnknownFilesystemFormat(f'Filetype "{fs_type.value}" is not supported')
 
-		cmd = [f'mkfs.{mkfs_type}', *options, *additional_parted_options, str(path)]
+		if not command:
+			command = f'mkfs.{mkfs_type}'
+
+		cmd = [command, *options, *additional_parted_options, str(path)]
 
 		debug('Formatting filesystem:', ' '.join(cmd))
 
@@ -536,7 +544,8 @@ def _setup_partition(
 			length=length_sector.value
 		)
 
-		filesystem = FileSystem(type=part_mod.safe_fs_type.value, geometry=geometry)
+		fs_value = part_mod.safe_fs_type.parted_value
+		filesystem = FileSystem(type=fs_value, geometry=geometry)
 
 		partition = Partition(
 			disk=disk,
@@ -549,7 +558,7 @@ def _setup_partition(
 			partition.setFlag(flag.flag_id)
 
 		debug(f'\tType: {part_mod.type.value}')
-		debug(f'\tFilesystem: {part_mod.safe_fs_type.value}')
+		debug(f'\tFilesystem: {fs_value}')
 		debug(f'\tGeometry: {start_sector.value} start sector, {length_sector.value} length')
 
 		try:
@@ -723,6 +732,13 @@ def partition(
 
 		disk.commit()
 
+	@staticmethod
+	def swapon(path: Path) -> None:
+		try:
+			SysCommand(['swapon', str(path)])
+		except SysCallError as err:
+			raise DiskError(f'Could not enable swap {path}:\n{err.message}')
+
 	def mount(
 		self,
 		dev_path: Path,
diff --git a/archinstall/lib/disk/device_model.py b/archinstall/lib/disk/device_model.py
index 610798a3d4..36d46f2215 100644
--- a/archinstall/lib/disk/device_model.py
+++ b/archinstall/lib/disk/device_model.py
@@ -728,6 +728,7 @@ class PartitionFlag(PartitionFlagDataMixin, Enum):
 	XBOOTLDR = parted.PARTITION_BLS_BOOT, "bls_boot"
 	ESP = parted.PARTITION_ESP
 	LINUX_HOME = parted.PARTITION_LINUX_HOME, "linux-home"
+	SWAP = parted.PARTITION_SWAP
 
 	@property
 	def description(self) -> str:
@@ -767,6 +768,7 @@ class FilesystemType(Enum):
 	Ntfs = 'ntfs'
 	Reiserfs = 'reiserfs'
 	Xfs = 'xfs'
+	LinuxSwap = 'linux-swap'
 
 	# this is not a FS known to parted, so be careful
 	# with the usage from this enum
@@ -785,6 +787,10 @@ def fs_type_mount(self) -> str:
 			case _:
 				return self.value
 
+	@property
+	def parted_value(self) -> str:
+		return self.value + '(v1)' if self == FilesystemType.LinuxSwap else self.value
+
 	@property
 	def installation_pkg(self) -> str | None:
 		match self:
@@ -969,6 +975,9 @@ def is_home(self) -> bool:
 			or PartitionFlag.LINUX_HOME in self.flags
 		)
 
+	def is_swap(self) -> bool:
+		return self.fs_type == FilesystemType.LinuxSwap
+
 	def is_modify(self) -> bool:
 		return self.status == ModificationStatus.Modify
 
diff --git a/archinstall/lib/disk/encryption_menu.py b/archinstall/lib/disk/encryption_menu.py
index 1d62935ef8..0f602795c4 100644
--- a/archinstall/lib/disk/encryption_menu.py
+++ b/archinstall/lib/disk/encryption_menu.py
@@ -284,7 +284,10 @@ def select_partitions_to_encrypt(
 
 	# do not allow encrypting the boot partition
 	for mod in modification:
-		partitions += [p for p in mod.partitions if p.mountpoint != Path('/boot')]
+		partitions += [
+			p for p in mod.partitions
+			if p.mountpoint != Path('/boot') and not p.is_swap()
+		]
 
 	# do not allow encrypting existing partitions that are not marked as wipe
 	avail_partitions = [p for p in partitions if not p.exists()]
diff --git a/archinstall/lib/disk/partitioning_menu.py b/archinstall/lib/disk/partitioning_menu.py
index e7bf4e2cb4..08ece3e863 100644
--- a/archinstall/lib/disk/partitioning_menu.py
+++ b/archinstall/lib/disk/partitioning_menu.py
@@ -470,7 +470,7 @@ def _create_new_partition(self, free_space: FreeSpace) -> PartitionModification:
 		fs_type = self._prompt_partition_fs_type()
 
 		mountpoint = None
-		if fs_type != FilesystemType.Btrfs:
+		if fs_type not in (FilesystemType.Btrfs, FilesystemType.LinuxSwap):
 			mountpoint = self._prompt_mountpoint()
 
 		partition = PartitionModification(
@@ -486,6 +486,8 @@ def _create_new_partition(self, free_space: FreeSpace) -> PartitionModification:
 			partition.set_flag(PartitionFlag.BOOT)
 			if self._using_gpt:
 				partition.set_flag(PartitionFlag.ESP)
+		elif partition.is_swap():
+			partition.set_flag(PartitionFlag.SWAP)
 
 		return partition
 
diff --git a/archinstall/lib/installer.py b/archinstall/lib/installer.py
index 6e41096c9a..9b523f9712 100644
--- a/archinstall/lib/installer.py
+++ b/archinstall/lib/installer.py
@@ -315,17 +315,21 @@ def _prepare_luks_lvm(
 		}
 
 	def _mount_partition(self, part_mod: disk.PartitionModification) -> None:
+		if not part_mod.dev_path:
+			return
+
 		# it would be none if it's btrfs as the subvolumes will have the mountpoints defined
-		if part_mod.mountpoint and part_mod.dev_path:
+		if part_mod.mountpoint:
 			target = self.target / part_mod.relative_mountpoint
 			disk.device_handler.mount(part_mod.dev_path, target, options=part_mod.mount_options)
-
-		if part_mod.fs_type == disk.FilesystemType.Btrfs and part_mod.dev_path:
+		elif part_mod.fs_type == disk.FilesystemType.Btrfs:
 			self._mount_btrfs_subvol(
 				part_mod.dev_path,
 				part_mod.btrfs_subvols,
 				part_mod.mount_options
 			)
+		elif part_mod.is_swap():
+			disk.device_handler.swapon(part_mod.dev_path)
 
 	def _mount_lvm_vol(self, volume: disk.LvmVolume) -> None:
 		if volume.fs_type != disk.FilesystemType.Btrfs: