Flexible menu 2 (#916)

* Correct definition of btrfs standard layout

* New version of the FlexibleMenu
* Added new functionality to Selector
* Created a GeneralMenu class
* GlobalMenu is made a child of GeneralMenu

* Some refining in GeneralMenu
secret is now a general function

* log is invoked in GeneralMenu directly

* Correction at GeneralMenu

* Materialize again _setup_selection_menu_options. Gives more room to play

* Callbacks converted as methods
Synch() (data area and menu) decoupled from enable()
and made general before any run

* script swiss added to the patch set

* Only_hd gets a new implementation of the menu
flake8 corrections

* swiss.py description added

* New version of the FlexibleMenu
* Added new functionality to Selector
* Created a GeneralMenu class
* GlobalMenu is made a child of GeneralMenu

* changes from the rebase left dangling

* Modify order of execution between exec_menu and post_processing.
Added selector_name as parameter for exec_menu

* minor corrections to the scripts

* Adapt to PR #874

* Solve issue #936

* make ask_for_a_timezone as synonym to ask_timezone

* Adapted to nationalization framework (PR 893).
String still NOT adapted

* flake8 complains

* Use of archinstall.output_config instead of local copy at swiss.py

* Problems with the last merge

* more flake8 complains. caused by reverted changes re. ask*timezone

* git complains

Co-authored-by: Anton Hvornum <anton@hvornum.se>

4 files changed, 301 insertions, 37 deletions

diff --git a/archinstall/__init__.py b/archinstall/__init__.py
index 4a7339a9..bf09b188 100644
--- a/archinstall/__init__.py
+++ b/archinstall/__init__.py
@@ -35,7 +35,11 @@ from .lib.storage import *
 from .lib.systemd import *
 from .lib.user_interaction import *
 from .lib.menu import Menu
-from .lib.menu.selection_menu import GlobalMenu
+from .lib.menu.selection_menu import (
+	GlobalMenu,
+	Selector,
+	GeneralMenu
+)
 from .lib.translation import Translation, DeferredTranslation
 from .lib.plugins import plugins, load_plugin # This initiates the plugin loading ceremony
 from .lib.configuration import *
diff --git a/archinstall/lib/general.py b/archinstall/lib/general.py
index c9ebb921..ad7b8ad4 100644
--- a/archinstall/lib/general.py
+++ b/archinstall/lib/general.py
@@ -550,3 +550,7 @@ def json_stream_to_structure(id : str, stream :str, target :dict) -> bool :
 			log(f" {id} is neither a file nor is a JSON string:",level=logging.ERROR)
 			return False
 	return True
+
+def secret(x :str):
+	""" return * with len equal to to the input string """
+	return '*' * len(x)
diff --git a/archinstall/lib/menu/selection_menu.py b/archinstall/lib/menu/selection_menu.py
index f094c6dc..93299502 100644
--- a/archinstall/lib/menu/selection_menu.py
+++ b/archinstall/lib/menu/selection_menu.py
@@ -1,8 +1,10 @@
+from __future__ import annotations
 import sys
-from typing import Dict
+import logging
+from typing import Callable, Any, List, Iterator, Dict
 
 from .menu import Menu
-from ..general import SysCommand
+from ..general import SysCommand, secret
 from ..storage import storage
 from ..output import log
 from ..profiles import is_desktop_profile
@@ -34,13 +36,16 @@ from ..translation import Translation
 class Selector:
 	def __init__(
 		self,
-		description,
-		func=None,
-		display_func=None,
-		default=None,
-		enabled=False,
-		dependencies=[],
-		dependencies_not=[]
+		description :str,
+		func :Callable = None,
+		display_func :Callable = None,
+		default :Any = None,
+		enabled :bool = False,
+		dependencies :List = [],
+		dependencies_not :List = [],
+		exec_func :Callable = None,
+		preview_func :Callable = None,
+		mandatory :bool = False
 	):
 		"""
 		Create a new menu selection entry
@@ -70,6 +75,17 @@ class Selector:
 		:param dependencies_not: These are the exclusive options; the menu item will only be
 		displayed if non of the entries in the list have been specified
 		:type dependencies_not: list
+
+		:param exec_func: A function with the name and the result of the selection as input parameter and which returns boolean.
+		Can be used for any action deemed necessary after selection. If it returns True, exits the menu loop, if False,
+		menu returns to the selection screen. If not specified it is assumed the return is False
+		:type exec_func: Callable
+
+		:param preview_func: A callable which invokws a preview screen (not implemented)
+		:type preview_func: Callable
+
+		:param mandatory: A boolean which determines that the field is mandatory, i.e. menu can not be exited if it is not set
+		:type mandatory: bool
 		"""
 
 		self._description = description
@@ -79,26 +95,29 @@ class Selector:
 		self.enabled = enabled
 		self._dependencies = dependencies
 		self._dependencies_not = dependencies_not
+		self.exec_func = exec_func
+		self.preview_func = preview_func
+		self.mandatory = mandatory
 
 	@property
-	def dependencies(self):
+	def dependencies(self) -> dict:
 		return self._dependencies
 
 	@property
-	def dependencies_not(self):
+	def dependencies_not(self) -> dict:
 		return self._dependencies_not
 
 	@property
 	def current_selection(self):
 		return self._current_selection
 
-	def set_enabled(self):
-		self.enabled = True
+	def set_enabled(self, status :bool = True):
+		self.enabled = status
 
-	def update_description(self, description):
+	def update_description(self, description :str):
 		self._description = description
 
-	def menu_text(self):
+	def menu_text(self) -> str:
 		current = ''
 
 		if self._display_func:
@@ -113,29 +132,271 @@ class Selector:
 
 		return f'{self._description} {current}'
 
-	def set_current_selection(self, current):
+	def set_current_selection(self, current :str):
 		self._current_selection = current
 
-	def has_selection(self):
+	def has_selection(self) -> bool:
 		if self._current_selection is None:
 			return False
 		return True
 
-	def is_empty(self):
+	def get_selection(self) -> Any:
+		return self._current_selection
+
+	def is_empty(self) -> bool:
 		if self._current_selection is None:
 			return True
 		elif isinstance(self._current_selection, (str, list, dict)) and len(self._current_selection) == 0:
 			return True
-
 		return False
 
+	def is_enabled(self) -> bool:
+		return self.enabled
+
+	def is_mandatory(self) -> bool:
+		return self.mandatory
+
+	def set_mandatory(self, status :bool = True):
+		self.mandatory = status
+		if status and not self.is_enabled():
+			self.set_enabled(True)
 
-class GlobalMenu:
-	def __init__(self):
+
+class GeneralMenu:
+	def __init__(self, data_store :dict = None):
+		"""
+		Create a new selection menu.
+
+		:param data_store:  Area (Dict) where the resulting data will be held. At least an entry for each option. Default area is self._data_store (not preset in the call, due to circular references
+		:type  data_store:  Dict
+
+		"""
 		self._translation = Translation.load_nationalization()
+		self.is_context_mgr = False
+		self._data_store = data_store if data_store is not None else {}
 		self._menu_options = {}
 		self._setup_selection_menu_options()
 
+	def __enter__(self, *args :Any, **kwargs :Any) -> GeneralMenu:
+		self.is_context_mgr = True
+		return self
+
+	def __exit__(self, *args :Any, **kwargs :Any) -> None:
+		# TODO: https://stackoverflow.com/questions/28157929/how-to-safely-handle-an-exception-inside-a-context-manager
+		# TODO: skip processing when it comes from a planified exit
+		if len(args) >= 2 and args[1]:
+			log(args[1], level=logging.ERROR, fg='red')
+			print("    Please submit this issue (and file) to https://github.com/archlinux/archinstall/issues")
+			raise args[1]
+
+		for key in self._menu_options:
+			sel = self._menu_options[key]
+			if key and key not in self._data_store:
+				self._data_store[key] = sel._current_selection
+
+		self.exit_callback()
+
+	def _setup_selection_menu_options(self):
+		""" Define the menu options.
+			Menu options can be defined here in a subclass or done per progam calling self.set_option()
+		"""
+		return
+
+	def pre_callback(self, selector_name):
+		""" will be called before each action in the menu """
+		return
+
+	def post_callback(self, selector_name :str, value :Any):
+		""" will be called after each action in the menu """
+		return True
+
+	def exit_callback(self):
+		""" will be called at the end of the processing of the menu """
+		return
+
+	def synch(self, selector_name :str, omit_if_set :bool = False,omit_if_disabled :bool = False):
+		""" loads menu options with data_store value """
+		arg = self._data_store.get(selector_name, None)
+		# don't display the menu option if it was defined already
+		if arg is not None and omit_if_set:
+			return
+
+		if not self.option(selector_name).is_enabled() and omit_if_disabled:
+			return
+
+		if arg is not None:
+			self._menu_options[selector_name].set_current_selection(arg)
+
+	def enable(self, selector_name :str, omit_if_set :bool = False , mandatory :bool = False):
+		""" activates menu options """
+		if self._menu_options.get(selector_name, None):
+			self._menu_options[selector_name].set_enabled(True)
+			if mandatory:
+				self._menu_options[selector_name].set_mandatory(True)
+			self.synch(selector_name,omit_if_set)
+		else:
+			print(f'No selector found: {selector_name}')
+			sys.exit(1)
+
+	def run(self):
+		""" Calls the Menu framework"""
+		# we synch all the options just in case
+		for item in self.list_options():
+			self.synch(item)
+		while True:
+			# Before continuing, set the preferred keyboard layout/language in the current terminal.
+			# 	This will just help the user with the next following questions.
+			self._set_kb_language()
+			enabled_menus = self._menus_to_enable()
+			menu_text = [m.text for m in enabled_menus.values()]
+			selection = Menu('Set/Modify the below options', menu_text, sort=False).run()
+			if selection:
+				selection = selection.strip()
+			if selection:
+				# if this calls returns false, we exit the menu. We allow for an callback for special processing on realeasing control
+				if not self._process_selection(selection):
+					break
+		if not self.is_context_mgr:
+			self.__exit__()
+
+	def _process_selection(self, selection :str) -> bool:
+		"""  determines and executes the selection y
+			Can / Should be extended to handle specific selection issues
+			Returns true if the menu shall continue, False if it has ended
+		"""
+		# find the selected option in our option list
+		option = [[k, v] for k, v in self._menu_options.items() if v.text.strip() == selection]
+		if len(option) != 1:
+			raise ValueError(f'Selection not found: {selection}')
+		selector_name = option[0][0]
+		selector = option[0][1]
+
+		return self.exec_option(selector_name,selector)
+
+	def exec_option(self,selector_name :str, p_selector :Selector = None) -> bool:
+		""" processes the exection of a given menu entry
+		- pre process callback
+		- selection function
+		- post process callback
+		- exec action
+		returns True if the loop has to continue, false if the loop can be closed
+		"""
+		if not p_selector:
+			selector = self.option(selector_name)
+		else:
+			selector = p_selector
+
+		self.pre_callback(selector_name)
+
+		result = None
+		if selector.func:
+			result = selector.func()
+			self._menu_options[selector_name].set_current_selection(result)
+			self._data_store[selector_name] = result
+		exec_ret_val = selector.exec_func(selector_name,result) if selector.exec_func else False
+		self.post_callback(selector_name,result)
+		if exec_ret_val and self._check_mandatory_status():
+			return False
+		return True
+		""" old behaviour
+		# we allow for a callback after we get the result
+		self.post_callback(selector_name,result)
+		# we have a callback, by option, to determine if we can exit the menu. Only if ALL mandatory fields are written
+		if selector.exec_func:
+			if selector.exec_func(result) and self._check_mandatory_status():
+				return False
+		"""
+		return True
+
+	def _set_kb_language(self):
+		""" general for ArchInstall"""
+		# Before continuing, set the preferred keyboard layout/language in the current terminal.
+		# This will just help the user with the next following questions.
+		if self._data_store.get('keyboard-layout', None) and len(self._data_store['keyboard-layout']):
+			set_keyboard_language(self._data_store['keyboard-layout'])
+
+	def _verify_selection_enabled(self, selection_name :str) -> bool:
+		""" general """
+		if selection := self._menu_options.get(selection_name, None):
+			if not selection.enabled:
+				return False
+
+			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():
+						return False
+
+			if len(selection.dependencies_not) > 0:
+				for d in selection.dependencies_not:
+					if not self._menu_options.get(d).is_empty():
+						return False
+			return True
+
+		raise ValueError(f'No selection found: {selection_name}')
+
+	def _menus_to_enable(self) -> dict:
+		""" general """
+		enabled_menus = {}
+
+		for name, selection in self._menu_options.items():
+			if self._verify_selection_enabled(name):
+				enabled_menus[name] = selection
+
+		return enabled_menus
+
+	def option(self,name :str) -> Selector:
+		# TODO check inexistent name
+		return self._menu_options[name]
+
+	def list_options(self) -> Iterator:
+		""" Iterator to retrieve the enabled menu option names
+		"""
+		for item in self._menu_options:
+			yield item
+
+	def list_enabled_options(self) -> Iterator:
+		""" Iterator to retrieve the enabled menu options at a given time.
+		The results are dynamic (if between calls to the iterator some elements -still not retrieved- are (de)activated
+		"""
+		for item in self._menu_options:
+			if item in self._menus_to_enable():
+				yield item
+
+	def set_option(self, name :str, selector :Selector):
+		self._menu_options[name] = selector
+		self.synch(name)
+
+	def _check_mandatory_status(self) -> bool:
+		for field in self._menu_options:
+			option = self._menu_options[field]
+			if option.is_mandatory() and not option.has_selection():
+				return False
+		return True
+
+	def set_mandatory(self, field :str, status :bool):
+		self.option(field).set_mandatory(status)
+
+	def mandatory_overview(self) -> [int, int]:
+		mandatory_fields = 0
+		mandatory_waiting = 0
+		for field in self._menu_options:
+			option = self._menu_options[field]
+			if option.is_mandatory():
+				mandatory_fields += 1
+				if not option.has_selection():
+					mandatory_waiting += 1
+		return mandatory_fields, mandatory_waiting
+
+	def _select_archinstall_language(self, default_lang):
+		language = select_archinstall_language(default_lang)
+		self._translation.activate(language)
+		return language
+
+
+class GlobalMenu(GeneralMenu):
+	def __init__(self,data_store):
+		super().__init__(data_store=data_store)
+
 	def _setup_selection_menu_options(self):
 		self._menu_options['archinstall-language'] = \
 			Selector(
@@ -170,8 +431,7 @@ class GlobalMenu:
 		self._menu_options['!encryption-password'] = \
 			Selector(
 				_('Set encryption password'),
-				lambda: get_password(prompt=str(_('Enter disk encryption password (leave blank for no encryption): '))),
-				display_func=lambda x: self._secret(x) if x else 'None',
+				display_func=lambda x: secret(x) if x else 'None',
 				dependencies=['harddrives'])
 		self._menu_options['swap'] = \
 			Selector(
@@ -188,7 +448,7 @@ class GlobalMenu:
 			Selector(
 				_('Set root password'),
 				lambda: self._set_root_password(),
-				display_func=lambda x: self._secret(x) if x else 'None')
+				display_func=lambda x: secret(x) if x else 'None')
 		self._menu_options['!superusers'] = \
 			Selector(
 				_('Specify superuser account'),
@@ -236,7 +496,9 @@ class GlobalMenu:
 		self._menu_options['install'] = \
 			Selector(
 				self._install_text(),
+				exec_func=lambda n,v: True if self._missing_configs() == 0 else False,
 				enabled=True)
+
 		self._menu_options['abort'] = Selector(_('Abort'), enabled=True)
 
 	def enable(self, selector_name, omit_if_set=False):
@@ -300,8 +562,11 @@ class GlobalMenu:
 		text = self._install_text()
 		self._menu_options.get('install').update_description(text)
 
-	def _post_processing(self):
-		if storage['arguments'].get('harddrives', None) and storage['arguments'].get('!encryption-password', None):
+	def post_callback(self,name :str = None ,result :Any = None):
+		self._update_install(name,result)
+
+	def exit_callback(self):
+		if self._data_store.get('harddrives', None) and self._data_store.get('!encryption-password', None):
 			# If no partitions was marked as encrypted, but a password was supplied and we have some disks to format..
 			# Then we need to identify which partitions to encrypt. This will default to / (root).
 			if len(list(encrypted_partitions(storage['arguments'].get('disk_layouts', [])))) == 0:
@@ -337,11 +602,6 @@ class GlobalMenu:
 
 		return missing
 
-	def _select_archinstall_language(self, default_lang):
-		language = select_archinstall_language(default_lang)
-		self._translation.activate(language)
-		return language
-
 	def _set_root_password(self):
 		prompt = str(_('Enter root password (leave blank to disable root): '))
 		password = get_password(prompt=prompt)
@@ -386,9 +646,6 @@ class GlobalMenu:
 
 		return harddrives
 
-	def _secret(self, x):
-		return '*' * len(x)
-
 	def _select_profile(self):
 		profile = select_profile()
 
@@ -444,4 +701,4 @@ class GlobalMenu:
 			if self._verify_selection_enabled(name):
 				enabled_menus[name] = selection
 
-		return enabled_menus
+		return enabled_menus
\ No newline at end of file
diff --git a/archinstall/lib/user_interaction.py b/archinstall/lib/user_interaction.py
index badf4cb4..7ed143f1 100644
--- a/archinstall/lib/user_interaction.py
+++ b/archinstall/lib/user_interaction.py
@@ -351,7 +351,6 @@ def ask_for_a_timezone() -> str:
 
 	return selected_tz
 
-
 def ask_for_bootloader(advanced_options :bool = False) -> str:
 	bootloader = "systemd-bootctl" if has_uefi() else "grub-install"
 	if has_uefi():