Fix 1862 (#1884)

* Fix 1862

* Update

---------

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

2 files changed, 100 insertions, 58 deletions

diff --git a/archinstall/lib/disk/device_model.py b/archinstall/lib/disk/device_model.py
index 36dd0c4f..8e72390c 100644
--- a/archinstall/lib/disk/device_model.py
+++ b/archinstall/lib/disk/device_model.py
@@ -137,6 +137,10 @@ class Unit(Enum):
 
 	Percent = '%' 	# size in percentile
 
+	@staticmethod
+	def get_all_units() -> List[str]:
+		return [u.name for u in Unit]
+
 
 @dataclass
 class Size:
@@ -214,16 +218,25 @@ class Size:
 				value = int(self._normalize() / target_unit.value)  # type: ignore
 				return Size(value, target_unit)
 
+	def as_text(self) -> str:
+		return self.format_size(
+			self.unit,
+			self.sector_size
+		)
+
 	def format_size(
 		self,
 		target_unit: Unit,
-		sector_size: Optional[Size] = None
+		sector_size: Optional[Size] = None,
+		include_unit: bool = True
 	) -> str:
 		if self.unit == Unit.Percent:
 			return f'{self.value}%'
 		else:
 			target_size = self.convert(target_unit, sector_size)
-			return f'{target_size.value} {target_unit.name}'
+			if include_unit:
+				return f'{target_size.value} {target_unit.name}'
+			return f'{target_size.value}'
 
 	def _normalize(self) -> int:
 		"""
@@ -280,7 +293,7 @@ class _PartitionInfo:
 	mountpoints: List[Path]
 	btrfs_subvol_infos: List[_BtrfsSubvolumeInfo] = field(default_factory=list)
 
-	def as_json(self) -> Dict[str, Any]:
+	def table_data(self) -> Dict[str, Any]:
 		part_info = {
 			'Name': self.name,
 			'Type': self.type.value,
@@ -343,7 +356,7 @@ class _DeviceInfo:
 	read_only: bool
 	dirty: bool
 
-	def as_json(self) -> Dict[str, Any]:
+	def table_data(self) -> Dict[str, Any]:
 		total_free_space = sum([region.get_length(unit=Unit.MiB) for region in self.free_space_regions])
 		return {
 			'Model': self.model,
@@ -440,7 +453,7 @@ class SubvolumeModification:
 			'nodatacow': self.nodatacow
 		}
 
-	def as_json(self) -> Dict[str, Any]:
+	def table_data(self) -> Dict[str, Any]:
 		return {
 			'name': str(self.name),
 			'mountpoint': str(self.mountpoint),
@@ -465,12 +478,20 @@ class DeviceGeometry:
 	def get_length(self, unit: Unit = Unit.sectors) -> int:
 		return self._geometry.getLength(unit.name)
 
-	def as_json(self) -> Dict[str, Any]:
+	def table_data(self) -> Dict[str, Any]:
+		start = Size(self._geometry.start, Unit.sectors, self._sector_size)
+		end = Size(self._geometry.end, Unit.sectors, self._sector_size)
+		length = Size(self._geometry.getLength(), Unit.sectors, self._sector_size)
+
+		start_str = f'{self._geometry.start} / {start.format_size(Unit.B, include_unit=False)}'
+		end_str = f'{self._geometry.end} / {end.format_size(Unit.B, include_unit=False)}'
+		length_str = f'{self._geometry.getLength()} / {length.format_size(Unit.B, include_unit=False)}'
+
 		return {
 			'Sector size': self._sector_size.value,
-			'Start sector': self._geometry.start,
-			'End sector': self._geometry.end,
-			'Length': self._geometry.getLength()
+			'Start (sector/B)': start_str,
+			'End (sector/B)': end_str,
+			'Length (sectors/B)': length_str
 		}
 
 
@@ -700,7 +721,7 @@ class PartitionModification:
 			'btrfs': [vol.__dump__() for vol in self.btrfs_subvols]
 		}
 
-	def as_json(self) -> Dict[str, Any]:
+	def table_data(self) -> Dict[str, Any]:
 		"""
 		Called for displaying data in table format
 		"""
diff --git a/archinstall/lib/disk/partitioning_menu.py b/archinstall/lib/disk/partitioning_menu.py
index 89cf6293..4acb4e85 100644
--- a/archinstall/lib/disk/partitioning_menu.py
+++ b/archinstall/lib/disk/partitioning_menu.py
@@ -1,10 +1,11 @@
 from __future__ import annotations
 
+import re
 from pathlib import Path
 from typing import Any, Dict, TYPE_CHECKING, List, Optional, Tuple
 
 from .device_model import PartitionModification, FilesystemType, BDevice, Size, Unit, PartitionType, PartitionFlag, \
-	ModificationStatus
+	ModificationStatus, DeviceGeometry
 from ..menu import Menu, ListManager, MenuSelection, TextInput
 from ..output import FormattedOutput, warn
 from .subvolume_menu import SubvolumeMenu
@@ -192,22 +193,51 @@ class PartitioningList(ListManager):
 		choice = Menu(prompt, options, sort=False, skip=False).run()
 		return options[choice.single_value]
 
-	def _validate_sector(self, start_sector: str, end_sector: Optional[str] = None) -> bool:
-		if not start_sector.isdigit():
-			return False
+	def _validate_value(
+		self,
+		sector_size: Size,
+		total_size: Size,
+		value: str
+	) -> Optional[Size]:
+		match = re.match(r'([0-9]+)([a-zA-Z|%]*)', value, re.I)
+
+		if match:
+			value, unit = match.groups()
+
+			if unit == '%':
+				unit = Unit.Percent.name
+
+			if unit and unit not in Unit.get_all_units():
+				return None
+
+			unit = Unit[unit] if unit else Unit.sectors
+			return Size(int(value), unit, sector_size, total_size)
+
+		return None
+
+	def _enter_size(
+		self,
+		sector_size: Size,
+		total_size: Size,
+		prompt: str,
+		default: Size
+	) -> Size:
+		while True:
+			value = TextInput(prompt).run().strip()
+
+			size: Optional[Size] = None
 
-		if end_sector:
-			if end_sector.endswith('%'):
-				if not end_sector[:-1].isdigit():
-					return False
-			elif not end_sector.isdigit():
-				return False
-			elif int(start_sector) > int(end_sector):
-				return False
+			if not value:
+				size = default
+			else:
+				size = self._validate_value(sector_size, total_size, value)
 
-		return True
+			if size:
+				return size
 
-	def _prompt_sectors(self) -> Tuple[Size, Size]:
+			warn(f'Invalid value: {value}')
+
+	def _prompt_size(self) -> Tuple[Size, Size]:
 		device_info = self._device.device_info
 
 		text = str(_('Current free sectors on device {}:')).format(device_info.path) + '\n\n'
@@ -215,54 +245,45 @@ class PartitioningList(ListManager):
 		prompt = text + free_space_table + '\n'
 
 		total_sectors = device_info.total_size.format_size(Unit.sectors, device_info.sector_size)
-		prompt += str(_('Total sectors: {}')).format(total_sectors) + '\n'
+		total_bytes = device_info.total_size.format_size(Unit.B)
+
+		prompt += str(_('Total: {} / {}')).format(total_sectors, total_bytes) + '\n\n'
+		prompt += str(_('All entered values can be suffixed with a unit: B, KB, KiB, MB, MiB...')) + '\n'
+		prompt += str(_('If no unit is provided, the value is interpreted as sectors')) + '\n'
 		print(prompt)
 
-		largest_free_area = max(device_info.free_space_regions, key=lambda r: r.get_length())
+		largest_free_area: DeviceGeometry = max(device_info.free_space_regions, key=lambda r: r.get_length())
 
 		# prompt until a valid start sector was entered
-		while True:
-			start_prompt = str(_('Enter the start sector (default: {}): ')).format(largest_free_area.start)
-			start_sector = TextInput(start_prompt).run().strip()
-
-			if not start_sector or self._validate_sector(start_sector):
-				break
-
-			warn(f'Invalid start sector entered: {start_sector}')
+		default_start = Size(largest_free_area.start, Unit.sectors, device_info.sector_size)
+		start_prompt = str(_('Enter start (default: sector {}): ')).format(largest_free_area.start)
+		start_size = self._enter_size(
+			device_info.sector_size,
+			device_info.total_size,
+			start_prompt,
+			default_start
+		)
 
-		if not start_sector:
-			start_sector = str(largest_free_area.start)
-			end_sector = str(largest_free_area.end)
+		if start_size.value == largest_free_area.start:
+			end_size = Size(largest_free_area.end, Unit.sectors, device_info.sector_size)
 		else:
-			end_sector = '100%'
+			end_size = Size(100, Unit.Percent, total_size=device_info.total_size)
 
 		# prompt until valid end sector was entered
-		while True:
-			end_prompt = str(_('Enter the end sector of the partition (percentage or block number, default: {}): ')).format(end_sector)
-			end_value = TextInput(end_prompt).run().strip()
-
-			if not end_value or self._validate_sector(start_sector, end_value):
-				break
-
-			warn(f'Invalid end sector entered: {start_sector}')
-
-		# override the default value with the user value
-		if end_value:
-			end_sector = end_value
-
-		start_size = Size(int(start_sector), Unit.sectors, device_info.sector_size)
-
-		if end_sector.endswith('%'):
-			end_size = Size(int(end_sector[:-1]), Unit.Percent, device_info.sector_size, device_info.total_size)
-		else:
-			end_size = Size(int(end_sector), Unit.sectors, device_info.sector_size)
+		end_prompt = str(_('Enter end (default: {}): ')).format(end_size.as_text())
+		end_size = self._enter_size(
+			device_info.sector_size,
+			device_info.total_size,
+			end_prompt,
+			end_size
+		)
 
 		return start_size, end_size
 
 	def _create_new_partition(self) -> PartitionModification:
 		fs_type = self._prompt_partition_fs_type()
 
-		start_size, end_size = self._prompt_sectors()
+		start_size, end_size = self._prompt_size()
 		length = end_size - start_size
 
 		# new line for the next prompt