Extend the mypy checks (#2120)

* Extend the mypy checks

* Update

* Update

* Update

---------

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

11 files changed, 54 insertions, 30 deletions

diff --git a/archinstall/__init__.py b/archinstall/__init__.py
index 07b85f96..11b47c48 100644
--- a/archinstall/__init__.py
+++ b/archinstall/__init__.py
@@ -225,8 +225,8 @@ def load_config():
 	if arguments.get('servers', None) is not None:
 		storage['_selected_servers'] = arguments.get('servers', None)
 
-	if arguments.get('network_config', None) is not None:
-		config = NetworkConfiguration.parse_arg(arguments.get('network_config'))
+	if (net_config := arguments.get('network_config', None)) is not None:
+		config = NetworkConfiguration.parse_arg(net_config)
 		arguments['network_config'] = config
 
 	if arguments.get('!users', None) is not None or arguments.get('!superusers', None) is not None:
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,
diff --git a/archinstall/lib/global_menu.py b/archinstall/lib/global_menu.py
index b38dac0b..deda2ef6 100644
--- a/archinstall/lib/global_menu.py
+++ b/archinstall/lib/global_menu.py
@@ -176,8 +176,11 @@ class GlobalMenu(AbstractMenu):
 		self._menu_options['abort'] = Selector(_('Abort'), exec_func=lambda n,v:exit(1))
 
 	def _missing_configs(self) -> List[str]:
-		def check(s):
-			return self._menu_options.get(s).has_selection()
+		def check(s) -> bool:
+			obj = self._menu_options.get(s)
+			if obj and obj.has_selection():
+				return True
+			return False
 
 		def has_superuser() -> bool:
 			sel = self._menu_options['!users']
@@ -228,7 +231,7 @@ class GlobalMenu(AbstractMenu):
 		return config.type.display_msg()
 
 	def _disk_encryption(self, preset: Optional[disk.DiskEncryption]) -> Optional[disk.DiskEncryption]:
-		mods: Optional[List[disk.DeviceModification]] = self._menu_options['disk_config'].current_selection
+		mods: Optional[disk.DiskLayoutConfiguration] = self._menu_options['disk_config'].current_selection
 
 		if not mods:
 			# this should not happen as the encryption menu has the disk_config as dependency
@@ -263,7 +266,7 @@ class GlobalMenu(AbstractMenu):
 
 	def _prev_additional_pkgs(self):
 		selector = self._menu_options['packages']
-		if selector.has_selection():
+		if selector.current_selection:
 			packages: List[str] = selector.current_selection
 			return format_cols(packages, None)
 		return None
diff --git a/archinstall/lib/installer.py b/archinstall/lib/installer.py
index ad98d9a8..8e716d3d 100644
--- a/archinstall/lib/installer.py
+++ b/archinstall/lib/installer.py
@@ -131,7 +131,11 @@ class Installer:
 		We need to wait for it before we continue since we opted in to use a custom mirror/region.
 		"""
 		info('Waiting for time sync (systemd-timesyncd.service) to complete.')
-		while SysCommand('timedatectl show --property=NTPSynchronized --value').decode() != 'yes':
+
+		while True:
+			time_val = SysCommand('timedatectl show --property=NTPSynchronized --value').decode()
+			if time_val and time_val.strip() == 'yes':
+				break
 			time.sleep(1)
 
 		info('Waiting for automatic mirror selection (reflector) to complete.')
@@ -237,7 +241,7 @@ class Installer:
 			gen_enc_file = self._disk_encryption.should_generate_encryption_file(part_mod)
 
 			luks_handler = Luks2(
-				part_mod.dev_path,
+				part_mod.safe_dev_path,
 				mapper_name=part_mod.mapper_name,
 				password=self._disk_encryption.encryption_password
 			)
@@ -281,8 +285,10 @@ class Installer:
 		self._fstab_entries.append(f'{file} none swap defaults 0 0')
 
 		if enable_resume:
-			resume_uuid = SysCommand(f'findmnt -no UUID -T {self.target}{file}').decode('UTF-8').strip()
-			resume_offset = SysCommand(f'/usr/bin/filefrag -v {self.target}{file}').decode().split('0:', 1)[1].split(":", 1)[1].split("..", 1)[0].strip()
+			resume_uuid = SysCommand(f'findmnt -no UUID -T {self.target}{file}').decode()
+			resume_offset = SysCommand(
+				f'/usr/bin/filefrag -v {self.target}{file}'
+			).decode().split('0:', 1)[1].split(":", 1)[1].split("..", 1)[0].strip()
 
 			self._hooks.append('resume')
 			self._kernel_params.append(f'resume=UUID={resume_uuid}')
diff --git a/archinstall/lib/menu/abstract_menu.py b/archinstall/lib/menu/abstract_menu.py
index 306c500a..053f3c30 100644
--- a/archinstall/lib/menu/abstract_menu.py
+++ b/archinstall/lib/menu/abstract_menu.py
@@ -14,7 +14,7 @@ class Selector:
 	def __init__(
 		self,
 		description: str,
-		func: Optional[Callable[[str], Any]] = None,
+		func: Optional[Callable[[Any], Any]] = None,
 		display_func: Optional[Callable] = None,
 		default: Optional[Any] = None,
 		enabled: bool = False,
diff --git a/archinstall/lib/models/gen.py b/archinstall/lib/models/gen.py
index cc8d7605..fb7e5751 100644
--- a/archinstall/lib/models/gen.py
+++ b/archinstall/lib/models/gen.py
@@ -1,5 +1,5 @@
 from dataclasses import dataclass
-from typing import Optional, List
+from typing import Optional, List, Dict, Any
 
 
 @dataclass
@@ -87,6 +87,10 @@ class PackageSearchResult:
 	makedepends: List[str]
 	checkdepends: List[str]
 
+	@staticmethod
+	def from_json(data: Dict[str, Any]) -> 'PackageSearchResult':
+		return PackageSearchResult(**data)
+
 	@property
 	def pkg_version(self) -> str:
 		return self.pkgver
@@ -107,8 +111,18 @@ class PackageSearch:
 	page: int
 	results: List[PackageSearchResult]
 
-	def __post_init__(self):
-		self.results = [PackageSearchResult(**x) for x in self.results]
+	@staticmethod
+	def from_json(data: Dict[str, Any]) -> 'PackageSearch':
+		results = [PackageSearchResult.from_json(r) for r in data['results']]
+
+		return PackageSearch(
+			version=data['version'],
+			limit=data['limit'],
+			valid=data['valid'],
+			num_pages=data['num_pages'],
+			page=data['page'],
+			results=results
+		)
 
 
 @dataclass
diff --git a/archinstall/lib/packages/packages.py b/archinstall/lib/packages/packages.py
index 7491df07..e495b03f 100644
--- a/archinstall/lib/packages/packages.py
+++ b/archinstall/lib/packages/packages.py
@@ -55,8 +55,8 @@ def package_search(package :str) -> PackageSearch:
 		raise PackageError(f"Could not locate package: [{response.code}] {response}")
 
 	data = response.read().decode('UTF-8')
-
-	return PackageSearch(**json.loads(data))
+	json_data = json.loads(data)
+	return PackageSearch.from_json(json_data)
 
 
 def find_package(package :str) -> List[PackageSearchResult]:
diff --git a/archinstall/lib/profile/profiles_handler.py b/archinstall/lib/profile/profiles_handler.py
index 03039321..515cdfe9 100644
--- a/archinstall/lib/profile/profiles_handler.py
+++ b/archinstall/lib/profile/profiles_handler.py
@@ -138,16 +138,16 @@ class ProfileHandler:
 			profiles = [profiles]
 
 		for profile in profiles:
-			self._profiles.append(profile)
+			self.profiles.append(profile)
 
-		self._verify_unique_profile_names(self._profiles)
+		self._verify_unique_profile_names(self.profiles)
 
 	def remove_custom_profiles(self, profiles: Union[TProfile, List[TProfile]]):
 		if not isinstance(profiles, list):
 			profiles = [profiles]
 
 		remove_names = [p.name for p in profiles]
-		self._profiles = [p for p in self._profiles if p.name not in remove_names]
+		self._profiles = [p for p in self.profiles if p.name not in remove_names]
 
 	def get_profile_by_name(self, name: str) -> Optional[Profile]:
 		return next(filter(lambda x: x.name == name, self.profiles), None)  # type: ignore
diff --git a/archinstall/lib/translationhandler.py b/archinstall/lib/translationhandler.py
index 33230562..3ea4c70e 100644
--- a/archinstall/lib/translationhandler.py
+++ b/archinstall/lib/translationhandler.py
@@ -206,4 +206,4 @@ class DeferredTranslation:
 	@classmethod
 	def install(cls):
 		import builtins
-		builtins._ = cls
+		builtins._ = cls  # type: ignore
diff --git a/pyproject.toml b/pyproject.toml
index 36ee0492..445aa277 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -61,7 +61,7 @@ packages = ["archinstall"]
 python_version = "3.11"
 files = "archinstall/"
 exclude = "tests"
-#check_untyped_defs=true
+check_untyped_defs=true
 
 [tool.bandit]
 targets = ["archinstall"]