Update partitioning.py

iridium_installer/ui/pages/partitioning.py

@@ -54,93 +54,88 @@ def get_total_memory() -> int:
                     return int(parts[1]) * 1024
     return 0
 
-
     def calculate_auto_partitions(disk_device):
-    """
+        disk_size = 0
-    Generates an automatic partition layout.
+        try:
-      - 2GB EFI (or 1GB if disk < required)
+            # Get disk size in bytes
-      - Root (Rest)
+            result = subprocess.run(
-      - RAM + 2GB Swap (at end) (or 0 if disk < required)
+                ["lsblk", "-J", "-b", "-o", "NAME,SIZE,TYPE"],
-    """
+                capture_output=True,
-    disk_size = 0
+                text=True,
-    try:
+            )
-        # Get disk size in bytes
+            if result.returncode == 0:
-        result = subprocess.run(
+                data = json.loads(result.stdout)
-            ["lsblk", "-J", "-b", "-o", "NAME,SIZE,TYPE"],
+                devices = data.get("blockdevices", [])
-            capture_output=True,
+                dev_name = disk_device.replace("/dev/", "")
-            text=True,
-        )
-        if result.returncode == 0:
-            data = json.loads(result.stdout)
-            devices = data.get("blockdevices", [])
-            dev_name = disk_device.replace("/dev/", "")
 
-            for dev in devices:
+                for dev in devices:
-                if dev.get("name") == dev_name:
+                    if dev.get("name") == dev_name:
-                    disk_size = int(dev.get("size", 0))
+                        disk_size = int(dev.get("size", 0))
-                    break
+                        break
-    except Exception as e:
+        except Exception as e:
-        print(f"Error getting disk size for auto partitioning: {e}")
+            print(f"Error getting disk size for auto partitioning: {e}")
-        return []
-
-    if disk_size == 0:
-        return []
-
-    ram_size = get_total_memory()
-    disk_mb = disk_size / (1024*1024)
-
-    # Defaults
-    efi_size = 2 * 1024 * 1024 * 1024
-    swap_size = ram_size + (2 * 1024 * 1024 * 1024)
-    min_root_size = 10 * 1024 * 1024 * 1024 # 10GB
-
-    total_required = efi_size + swap_size + min_root_size
-    
-    use_swap = True
-    
-    if disk_size < total_required:
-        efi_size = 1 * 1024 * 1024 * 1024
-        use_swap = False
-        swap_size = 0
-        if disk_size < (efi_size + min_root_size):
-            print("Disk too small for automatic partitioning scheme.")
             return []
 
-    root_size = disk_size - efi_size - swap_size
+        if disk_size == 0:
+            return []
 
-    partitions = [
+        ram_size = get_total_memory()
-        {
+        disk_mb = disk_size / (1024 * 1024)
-            "type": "partition",
-            "name": "EFI System",
-            "filesystem": "vfat",
-            "mount_point": "/boot/efi",
-            "size": f"{efi_size / (1024**3):.1f} GB",
-            "bytes": efi_size,
-            "style_class": "part-efi",
-        },
-        {
-            "type": "partition",
-            "name": "Root",
-            "filesystem": "ext4",
-            "mount_point": "/",
-            "size": f"{root_size / (1024**3):.1f} GB",
-            "bytes": root_size,
-            "style_class": "part-root",
-        }
-    ]
 
-    if use_swap:
+        # Defaults
-        partitions.append({
+        efi_size = 2 * 1024 * 1024 * 1024
-            "type": "partition",
+        swap_size = ram_size + (2 * 1024 * 1024 * 1024)
-            "name": "Swap",
+        min_root_size = 10 * 1024 * 1024 * 1024  # 10GB
-            "filesystem": "swap",
-            "mount_point": "[SWAP]",
-            "size": f"{swap_size / (1024**3):.1f} GB",
-            "bytes": swap_size,
-            "style_class": "part-swap",
-        })
 
-    return partitions
+        total_required = efi_size + swap_size + min_root_size
+
+        use_swap = True
+
+        if disk_size < total_required:
+            efi_size = 1 * 1024 * 1024 * 1024
+            use_swap = False
+            swap_size = 0
+            if disk_size < (efi_size + min_root_size):
+                print("Disk too small for automatic partitioning scheme.")
+                return []
+
+        root_size = disk_size - efi_size - swap_size
+
+        partitions = [
+            {
+                "type": "partition",
+                "name": "EFI System",
+                "filesystem": "vfat",
+                "mount_point": "/boot/efi",
+                "size": f"{efi_size / (1024**3):.1f} GB",
+                "bytes": efi_size,
+                "style_class": "part-efi",
+            },
+            {
+                "type": "partition",
+                "name": "Root",
+                "filesystem": "ext4",
+                "mount_point": "/",
+                "size": f"{root_size / (1024**3):.1f} GB",
+                "bytes": root_size,
+                "style_class": "part-root",
+            },
+        ]
+
+        if use_swap:
+            partitions.append(
+                {
+                    "type": "partition",
+                    "name": "Swap",
+                    "filesystem": "swap",
+                    "mount_point": "[SWAP]",
+                    "size": f"{swap_size / (1024**3):.1f} GB",
+                    "bytes": swap_size,
+                    "style_class": "part-swap",
+                }
+            )
+
+        return partitions
 
 
 class PartitionSegment(Gtk.Button):
@@ -439,7 +434,7 @@ class PartitioningPage(Adw.Bin):
 
     def show_context_menu(self, widget, x, y):
         data = widget.part_data
-        self.selected_disk_path = self.current_disk_path # Assumes we store this
+        self.selected_disk_path = self.current_disk_path  # Assumes we store this
 
         popover = Gtk.Popover()
         popover.set_parent(widget)
@@ -464,22 +459,41 @@ class PartitioningPage(Adw.Bin):
             menu_box.append(btn)
 
         if data.get("type") == "partition":
-            add_menu_item("Select Mount Point", "folder-open-symbolic", 
+            add_menu_item(
-                         lambda: self.select_mount_point(data))
+                "Select Mount Point",
+                "folder-open-symbolic",
+                lambda: self.select_mount_point(data),
+            )
             separator = Gtk.Separator()
             menu_box.append(separator)
-            add_menu_item("Delete", "user-trash-symbolic", 
+            add_menu_item(
-                         lambda: self.delete_part(data), destructive=True)
+                "Delete",
+                "user-trash-symbolic",
+                lambda: self.delete_part(data),
+                destructive=True,
+            )
 
         elif data.get("type") == "empty":
-            add_menu_item("Create Partition", "list-add-symbolic", 
+            add_menu_item(
-                         lambda: self.create_part_dialog(data))
+                "Create Partition",
+                "list-add-symbolic",
+                lambda: self.create_part_dialog(data),
+            )
 
         popover.popup()
 
     def select_mount_point(self, data):
-        win = Adw.Window(title="Select Mount Point", modal=True, transient_for=self.get_root())
+        win = Adw.Window(
-        box = Gtk.Box(orientation=Gtk.Orientation.VERTICAL, spacing=12, margin_top=24, margin_bottom=24, margin_start=24, margin_end=24)
+            title="Select Mount Point", modal=True, transient_for=self.get_root()
+        )
+        box = Gtk.Box(
+            orientation=Gtk.Orientation.VERTICAL,
+            spacing=12,
+            margin_top=24,
+            margin_bottom=24,
+            margin_start=24,
+            margin_end=24,
+        )
         win.set_content(box)
 
         options = ["/", "/boot/efi", "[SWAP]", "None"]
@@ -493,7 +507,7 @@ class PartitioningPage(Adw.Bin):
         if current_mp in options:
             dropdown.set_selected(options.index(current_mp))
         else:
-             dropdown.set_selected(options.index("None"))
+            dropdown.set_selected(options.index("None"))
 
         box.append(Gtk.Label(label=f"Mount point for {data['name']}:"))
         box.append(dropdown)
@@ -511,9 +525,11 @@ class PartitioningPage(Adw.Bin):
         win.present()
 
     def delete_part(self, data):
-        from ...backend.disk import delete_partition
         # Extract partition number from name (e.g., nvme0n1p3 -> 3)
         import re
+
+        from ...backend.disk import delete_partition
+
         match = re.search(r"(\d+)$", data["name"])
         if match:
             part_num = int(match.group(1))
@@ -521,8 +537,17 @@ class PartitioningPage(Adw.Bin):
             self.load_partitions(self.current_disk_path)
 
     def create_part_dialog(self, data):
-        win = Adw.Window(title="Create Partition", modal=True, transient_for=self.get_root())
+        win = Adw.Window(
-        box = Gtk.Box(orientation=Gtk.Orientation.VERTICAL, spacing=12, margin_top=24, margin_bottom=24, margin_start=24, margin_end=24)
+            title="Create Partition", modal=True, transient_for=self.get_root()
+        )
+        box = Gtk.Box(
+            orientation=Gtk.Orientation.VERTICAL,
+            spacing=12,
+            margin_top=24,
+            margin_bottom=24,
+            margin_start=24,
+            margin_end=24,
+        )
         win.set_content(box)
 
         # Type dropdown
@@ -534,7 +559,7 @@ class PartitioningPage(Adw.Bin):
 
         # Size entry
         size_box = Gtk.Box(orientation=Gtk.Orientation.HORIZONTAL, spacing=12)
-        size_entry = Gtk.Entry(text=str(int(data["bytes"] / (1024*1024))))
+        size_entry = Gtk.Entry(text=str(int(data["bytes"] / (1024 * 1024))))
         size_entry.set_width_chars(15)
         size_box.append(size_entry)
         size_box.append(Gtk.Label(label="MB"))
@@ -547,6 +572,7 @@ class PartitioningPage(Adw.Bin):
 
         def on_create(b):
             from ...backend.disk import create_partition
+
             selected_type = type_names[type_dropdown.get_selected()]
 
             # Default name and fstype based on type
@@ -563,7 +589,7 @@ class PartitioningPage(Adw.Bin):
             size_text = size_entry.get_text()
             try:
                 size_mb = int(size_text)
-                max_mb = int(data["bytes"] / (1024*1024))
+                max_mb = int(data["bytes"] / (1024 * 1024))
                 if size_mb >= max_mb:
                     size_mb = 0
             except ValueError:
@@ -573,7 +599,9 @@ class PartitioningPage(Adw.Bin):
             type_code = types[selected_type]
 
             try:
-                create_partition(self.current_disk_path, size_mb, type_code, name, fstype)
+                create_partition(
+                    self.current_disk_path, size_mb, type_code, name, fstype
+                )
                 win.close()
                 self.load_partitions(self.current_disk_path)
             except Exception as e: