Extend the mypy checks (#2120)

* Extend the mypy checks

* Update

* Update

* Update

---------

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

2 files changed, 10 insertions, 9 deletions

diff --git a/archinstall/lib/disk/device_model.py b/archinstall/lib/disk/device_model.py
index 4ac53b0c..cd955851 100644
--- a/archinstall/lib/disk/device_model.py
+++ b/archinstall/lib/disk/device_model.py
@@ -308,9 +308,9 @@ class _PartitionInfo:
 	start: Size
 	length: Size
 	flags: List[PartitionFlag]
-	partn: int
-	partuuid: str
-	uuid: str
+	partn: Optional[int]
+	partuuid: Optional[str]
+	uuid: Optional[str]
 	disk: Disk
 	mountpoints: List[Path]
 	btrfs_subvol_infos: List[_BtrfsSubvolumeInfo] = field(default_factory=list)
@@ -344,9 +344,9 @@ class _PartitionInfo:
 		cls,
 		partition: Partition,
 		fs_type: Optional[FilesystemType],
-		partn: int,
-		partuuid: str,
-		uuid: str,
+		partn: Optional[int],
+		partuuid: Optional[str],
+		uuid: Optional[str],
 		mountpoints: List[Path],
 		btrfs_subvol_infos: List[_BtrfsSubvolumeInfo] = []
 	) -> _PartitionInfo:
diff --git a/archinstall/lib/disk/encryption_menu.py b/archinstall/lib/disk/encryption_menu.py
index 234e3b03..c3a1c32f 100644
--- a/archinstall/lib/disk/encryption_menu.py
+++ b/archinstall/lib/disk/encryption_menu.py
@@ -3,6 +3,7 @@ from typing import Dict, Optional, Any, TYPE_CHECKING, List
 
 from ..disk import (
 	DeviceModification,
+	DiskLayoutConfiguration,
 	PartitionModification,
 	DiskEncryption,
 	EncryptionType
@@ -26,7 +27,7 @@ if TYPE_CHECKING:
 class DiskEncryptionMenu(AbstractSubMenu):
 	def __init__(
 		self,
-		mods: List[DeviceModification],
+		disk_config: DiskLayoutConfiguration,
 		data_store: Dict[str, Any],
 		preset: Optional[DiskEncryption] = None
 	):
@@ -35,7 +36,7 @@ class DiskEncryptionMenu(AbstractSubMenu):
 		else:
 			self._preset = DiskEncryption()
 
-		self._modifications = mods
+		self._disk_config = disk_config
 		super().__init__(data_store=data_store)
 
 	def setup_selection_menu_options(self):
@@ -59,7 +60,7 @@ class DiskEncryptionMenu(AbstractSubMenu):
 		self._menu_options['partitions'] = \
 			Selector(
 				_('Partitions'),
-				func=lambda preset: select_partitions_to_encrypt(self._modifications.device_modifications, preset),
+				func=lambda preset: select_partitions_to_encrypt(self._disk_config.device_modifications, preset),
 				display_func=lambda x: f'{len(x)} {_("Partitions")}' if x else None,
 				dependencies=['encryption_password'],
 				default=self._preset.partitions,