Fix some mypy things (#1023)

Co-authored-by: Daniel Girtler <girtler.daniel@gmail.com>
Co-authored-by: Anton Hvornum <anton@hvornum.se>

1 files changed, 14 insertions, 12 deletions

diff --git a/archinstall/lib/menu/selection_menu.py b/archinstall/lib/menu/selection_menu.py
index a6c408b7..79f582f6 100644
--- a/archinstall/lib/menu/selection_menu.py
+++ b/archinstall/lib/menu/selection_menu.py
@@ -2,13 +2,16 @@ from __future__ import annotations
 import sys
 import logging
 
-from typing import Callable, Any, List, Iterator, Tuple, Optional
+from typing import Callable, Any, List, Iterator, Tuple, Optional, Dict, TYPE_CHECKING
 
 from .menu import Menu
 from ..locale_helpers import set_keyboard_language
 from ..output import log
 from ..translation import Translation
 
+if TYPE_CHECKING:
+	_: Any
+
 def select_archinstall_language(default='English'):
 	"""
 	copied from user_interaction/general_conf.py as a temporary measure
@@ -131,7 +134,7 @@ class Selector:
 	def text(self):
 		return self.menu_text()
 
-	def set_current_selection(self, current :str):
+	def set_current_selection(self, current :Optional[str]):
 		self._current_selection = current
 
 	def has_selection(self) -> bool:
@@ -179,7 +182,7 @@ class GeneralMenu:
 		self.is_context_mgr = False
 		self._data_store = data_store if data_store is not None else {}
 		self.auto_cursor = auto_cursor
-		self._menu_options = {}
+		self._menu_options: Dict[str, Selector] = {}
 		self._setup_selection_menu_options()
 		self.preview_size = preview_size
 
@@ -251,7 +254,7 @@ class GeneralMenu:
 		return None
 
 	def _find_selection(self, selection_name: str) -> Tuple[str, Selector]:
-		option = [[k, v] for k, v in self._menu_options.items() if v.text.strip() == selection_name.strip()]
+		option = [(k, v) for k, v in self._menu_options.items() if v.text.strip() == selection_name.strip()]
 		if len(option) != 1:
 			raise ValueError(f'Selection not found: {selection_name}')
 		config_name = option[0][0]
@@ -346,12 +349,12 @@ class GeneralMenu:
 
 			if len(selection.dependencies) > 0:
 				for d in selection.dependencies:
-					if not self._verify_selection_enabled(d) or self._menu_options.get(d).is_empty():
+					if not self._verify_selection_enabled(d) or self._menu_options[d].is_empty():
 						return False
 
 			if len(selection.dependencies_not) > 0:
 				for d in selection.dependencies_not:
-					if not self._menu_options.get(d).is_empty():
+					if not self._menu_options[d].is_empty():
 						return False
 			return True
 
@@ -399,7 +402,7 @@ class GeneralMenu:
 	def set_mandatory(self, field :str, status :bool):
 		self.option(field).set_mandatory(status)
 
-	def mandatory_overview(self) -> [int, int]:
+	def mandatory_overview(self) -> Tuple[int, int]:
 		mandatory_fields = 0
 		mandatory_waiting = 0
 		for field in self._menu_options:
@@ -553,7 +556,7 @@ class GlobalMenu(GeneralMenu):
 
 	def _update_install_text(self, name :str = None, result :Any = None):
 		text = self._install_text()
-		self._menu_options.get('install').update_description(text)
+		self._menu_options['install'].update_description(text)
 
 	def post_callback(self,name :str = None ,result :Any = None):
 		self._update_install_text(name, result)
@@ -596,7 +599,7 @@ class GlobalMenu(GeneralMenu):
 		if not check('harddrives'):
 			missing += ['Hard drives']
 		if check('harddrives'):
-			if not self._menu_options.get('harddrives').is_empty() and not check('disk_layouts'):
+			if not self._menu_options['harddrives'].is_empty() and not check('disk_layouts'):
 				missing += ['Disk layout']
 
 		return missing
@@ -621,12 +624,11 @@ class GlobalMenu(GeneralMenu):
 		return ntp
 
 	def _select_harddrives(self, old_harddrives : list) -> list:
-		# old_haddrives = storage['arguments'].get('harddrives', [])
 		harddrives = select_harddrives(old_harddrives)
 
 		# in case the harddrives got changed we have to reset the disk layout as well
 		if old_harddrives != harddrives:
-			self._menu_options.get('disk_layouts').set_current_selection(None)
+			self._menu_options['disk_layouts'].set_current_selection(None)
 			storage['arguments']['disk_layouts'] = {}
 
 		if not harddrives:
@@ -639,7 +641,7 @@ class GlobalMenu(GeneralMenu):
 			choice = Menu(prompt, ['yes', 'no'], default_option='yes').run()
 
 			if choice == 'no':
-				exit(1)
+				return self._select_harddrives(old_harddrives)
 
 		return harddrives