Update the subvolume menu - fix for #1278 (#1297)

* Update subvolume

* Add mypy compliance

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

2 files changed, 75 insertions, 138 deletions

diff --git a/archinstall/lib/user_interaction/partitioning_conf.py b/archinstall/lib/user_interaction/partitioning_conf.py
index caf5f5df..63b7c7df 100644
--- a/archinstall/lib/user_interaction/partitioning_conf.py
+++ b/archinstall/lib/user_interaction/partitioning_conf.py
@@ -351,18 +351,16 @@ def manage_new_and_existing_partitions(block_device: 'BlockDevice') -> Dict[str,
 				if partition is not None:
 					if not block_device_struct["partitions"][partition].get('btrfs', {}):
 						block_device_struct["partitions"][partition]['btrfs'] = {}
-					if not block_device_struct["partitions"][partition]['btrfs'].get('subvolumes', {}):
-						block_device_struct["partitions"][partition]['btrfs']['subvolumes'] = {}
+					if not block_device_struct["partitions"][partition]['btrfs'].get('subvolumes', []):
+						block_device_struct["partitions"][partition]['btrfs']['subvolumes'] = []
 
 					prev = block_device_struct["partitions"][partition]['btrfs']['subvolumes']
-					result = SubvolumeList(_("Manage btrfs subvolumes for current partition"),prev).run()
-					if result:
-						block_device_struct["partitions"][partition]['btrfs']['subvolumes'] = result
-					else:
-						del block_device_struct["partitions"][partition]['btrfs']
+					result = SubvolumeList(_("Manage btrfs subvolumes for current partition"), prev).run()
+					block_device_struct["partitions"][partition]['btrfs']['subvolumes'] = result
 
 	return block_device_struct
 
+
 def select_encrypted_partitions(
 	title :str,
 	partitions :List[Partition],
diff --git a/archinstall/lib/user_interaction/subvolume_config.py b/archinstall/lib/user_interaction/subvolume_config.py
index 94e6f5d7..a54ec891 100644
--- a/archinstall/lib/user_interaction/subvolume_config.py
+++ b/archinstall/lib/user_interaction/subvolume_config.py
@@ -1,155 +1,94 @@
-from typing import Dict, List
+from typing import Dict, List, Optional, Any, TYPE_CHECKING
 
 from ..menu.list_manager import ListManager
 from ..menu.menu import MenuSelectionType
-from ..menu.selection_menu import Selector, GeneralMenu
 from ..menu.text_input import TextInput
 from ..menu import Menu
+from ..models.subvolume import Subvolume
+
+if TYPE_CHECKING:
+	_: Any
 
-"""
-UI classes
-"""
 
 class SubvolumeList(ListManager):
-	def __init__(self,prompt,list):
-		self.ObjectNullAction = None # str(_('Add'))
-		self.ObjectDefaultAction = str(_('Add'))
-		super().__init__(prompt,list,None,self.ObjectNullAction,self.ObjectDefaultAction)
-
-	def reformat(self, data: Dict) -> Dict:
-		def presentation(key :str, value :Dict):
-			text = _(" Subvolume :{:16}").format(key)
-			if isinstance(value,str):
-				text += _(" mounted at {:16}").format(value)
-			else:
-				if value.get('mountpoint'):
-					text += _(" mounted at {:16}").format(value['mountpoint'])
-				else:
-					text += (' ' * 28)
-
-				if value.get('options',[]):
-					text += _(" with option {}").format(', '.join(value['options']))
-			return text
-
-		formatted = {presentation(k, v): k for k, v in data.items()}
-		return {k: v for k, v in sorted(formatted.items(), key=lambda e: e[0])}
+	def __init__(self, prompt: str, current_volumes: List[Subvolume]):
+		self._actions = [
+			str(_('Add subvolume')),
+			str(_('Edit subvolume')),
+			str(_('Delete subvolume'))
+		]
+		super().__init__(prompt, current_volumes, self._actions, self._actions[0])
 
-	def action_list(self):
-		return super().action_list()
+	def reformat(self, data: List[Subvolume]) -> Dict[str, Subvolume]:
+		return {e.display(): e for e in data}
 
-	def exec_action(self, data: Dict):
-		if self.target:
-			origkey, origval = list(self.target.items())[0]
-		else:
-			origkey = None
+	def action_list(self):
+		active_user = self.target if self.target else None
 
-		if self.action == str(_('Delete')):
-			del data[origkey]
+		if active_user is None:
+			return [self._actions[0]]
 		else:
-			if self.action == str(_('Add')):
-				self.target = {}
-				print(_('\n Fill the desired values for a new subvolume \n'))
-				with SubvolumeMenu(self.target,self.action) as add_menu:
-					for elem in ['name','mountpoint','options']:
-						add_menu.exec_option(elem)
-			else:
-				SubvolumeMenu(self.target,self.action).run()
+			return self._actions[1:]
 
-			data.update(self.target)
+	def _prompt_options(self, editing: Optional[Subvolume] = None) -> List[str]:
+		preset_options = []
+		if editing:
+			preset_options = editing.options
 
-		return data
-
-
-class SubvolumeMenu(GeneralMenu):
-	def __init__(self,parameters,action=None):
-		self.data = parameters
-		self.action = action
-		self.ds = {}
-		self.ds['name'] = None
-		self.ds['mountpoint'] = None
-		self.ds['options'] = None
-		if self.data:
-			origkey,origval = list(self.data.items())[0]
-			self.ds['name'] = origkey
-			if isinstance(origval,str):
-				self.ds['mountpoint'] = origval
-			else:
-				self.ds['mountpoint'] = self.data[origkey].get('mountpoint')
-				self.ds['options'] = self.data[origkey].get('options')
-
-		super().__init__(data_store=self.ds)
-
-	def _setup_selection_menu_options(self):
-		self._menu_options['name'] = Selector(
-			str(_('Subvolume name ')),
-			self._select_subvolume_name if not self.action or self.action in (str(_('Add')), str(_('Copy'))) else None,
-			mandatory=True,
-			enabled=True)
-
-		self._menu_options['mountpoint'] = Selector(
-			str(_('Subvolume mountpoint')),
-			self._select_subvolume_mount_point if not self.action or self.action in (str(_('Add')),str(_('Edit'))) else None,
-			enabled=True)
-
-		self._menu_options['options'] = Selector(
-			str(_('Subvolume options')),
-			self._select_subvolume_options if not self.action or self.action in (str(_('Add')),str(_('Edit'))) else None,
-			enabled=True)
-
-		self._menu_options['save'] = Selector(
-			str(_('Save')),
-			exec_func=lambda n,v:True,
-			enabled=True)
-
-		self._menu_options['cancel'] = Selector(
-			str(_('Cancel')),
-			# func = lambda pre:True,
-			exec_func=lambda n,v:self.fast_exit(n),
-			enabled=True)
-
-		self.cancel_action = 'cancel'
-		self.save_action = 'save'
-		self.bottom_list = [self.save_action,self.cancel_action]
-
-	def fast_exit(self,accion):
-		if self.option(accion).get_selection():
-			for item in self.list_options():
-				if self.option(item).is_mandatory():
-					self.option(item).set_mandatory(False)
-		return True
-
-	def exit_callback(self):
-		# we exit without moving data
-		if self.option(self.cancel_action).get_selection():
-			return
-		if not self.ds['name']:
-			return
-		else:
-			key = self.ds['name']
-			value = {}
-			if self.ds['mountpoint']:
-				value['mountpoint'] = self.ds['mountpoint']
-			if self.ds['options']:
-				value['options'] = self.ds['options']
-			self.data.update({key : value})
-
-	def _select_subvolume_name(self,value):
-		return TextInput(str(_("Subvolume name :")),value).run()
-
-	def _select_subvolume_mount_point(self,value):
-		return TextInput(str(_("Select a mount point :")),value).run()
-
-	def _select_subvolume_options(self,value) -> List[str]:
-		# def __init__(self, title, p_options, skip=True, multi=False, default_option=None, sort=True):
 		choice = Menu(
 			str(_("Select the desired subvolume options ")),
 			['nodatacow','compress'],
 			skip=True,
-			preset_values=value,
+			preset_values=preset_options,
 			multi=True
 		).run()
 
 		if choice.type_ == MenuSelectionType.Selection:
-			return choice.value
+			return choice.value  # type: ignore
 
 		return []
+
+	def _add_subvolume(self, editing: Optional[Subvolume] = None) -> Optional[Subvolume]:
+		name = TextInput(f'\n\n{_("Subvolume name")}: ', editing.name if editing else '').run()
+
+		if not name:
+			return None
+
+		mountpoint = TextInput(f'\n{_("Subvolume mountpoint")}: ', editing.mountpoint if editing else '').run()
+
+		if not mountpoint:
+			return None
+
+		options = self._prompt_options(editing)
+
+		subvolume = Subvolume(name, mountpoint)
+		subvolume.compress = 'compress' in options
+		subvolume.nodatacow = 'nodatacow' in options
+
+		return subvolume
+
+	def exec_action(self, data: List[Subvolume]) -> List[Subvolume]:
+		if self.target:
+			active_subvolume = self.target
+		else:
+			active_subvolume = None
+
+		if self.action == self._actions[0]:  # add
+			new_subvolume = self._add_subvolume()
+
+			if new_subvolume is not None:
+				# in case a user with the same username as an existing user
+				# was created we'll replace the existing one
+				data = [d for d in data if d.name != new_subvolume.name]
+				data += [new_subvolume]
+		elif self.action == self._actions[1]:  # edit subvolume
+			new_subvolume = self._add_subvolume(active_subvolume)
+
+			if new_subvolume is not None:
+				# we'll remove the original subvolume and add the modified version
+				data = [d for d in data if d.name != active_subvolume.name and d.name != new_subvolume.name]
+				data += [new_subvolume]
+		elif self.action == self._actions[2]:  # delete
+			data = [d for d in data if d != active_subvolume]
+
+		return data