diff --git a/iridium_installer/ui/pages/partitioning.py b/iridium_installer/ui/pages/partitioning.py
index 078979f..de2f860 100644
--- a/iridium_installer/ui/pages/partitioning.py
+++ b/iridium_installer/ui/pages/partitioning.py
@@ -43,7 +43,6 @@ CSS = """
 
 
 def get_total_memory() -> int:
-    """Returns total system memory in bytes."""
     try:
         return os.sysconf("SC_PAGE_SIZE") * os.sysconf("SC_PHYS_PAGES")
     except (ValueError, OSError):
@@ -54,88 +53,89 @@ def get_total_memory() -> int:
                     return int(parts[1]) * 1024
     return 0
 
-    def calculate_auto_partitions(disk_device):
-        disk_size = 0
-        try:
-            # Get disk size in bytes
-            result = subprocess.run(
-                ["lsblk", "-J", "-b", "-o", "NAME,SIZE,TYPE"],
-                capture_output=True,
-                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:
-                    if dev.get("name") == dev_name:
-                        disk_size = int(dev.get("size", 0))
-                        break
-        except Exception as e:
-            print(f"Error getting disk size for auto partitioning: {e}")
+def calculate_auto_partitions(disk_device):
+    disk_size = 0
+    try:
+        # Get disk size in bytes
+        result = subprocess.run(
+            ["lsblk", "-J", "-b", "-o", "NAME,SIZE,TYPE"],
+            capture_output=True,
+            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:
+                if dev.get("name") == dev_name:
+                    disk_size = int(dev.get("size", 0))
+                    break
+    except Exception as 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 []
 
-        if disk_size == 0:
-            return []
+    root_size = disk_size - efi_size - swap_size
 
-        ram_size = get_total_memory()
-        disk_mb = disk_size / (1024 * 1024)
+    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",
+        },
+    ]
 
-        # 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
-
-        partitions = [
+    if use_swap:
+        partitions.append(
             {
                 "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",
-            },
-        ]
+                "name": "Swap",
+                "filesystem": "swap",
+                "mount_point": "[SWAP]",
+                "size": f"{swap_size / (1024**3):.1f} GB",
+                "bytes": swap_size,
+                "style_class": "part-swap",
+            }
+        )
 
-        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
+    return partitions
 
 
 class PartitionSegment(Gtk.Button):