PyQTDiskPartitioningDemo/partition_demo.py

import sys, parted
from PyQt5.QtCore import *
from PyQt5.QtGui import *
from PyQt5.QtWidgets import *


####################################################################
def main():
    app = QApplication(sys.argv)
    app.setApplicationName('Disk Partitioning Demo')
    w = PartioningDemo()
    w.show()
    sys.exit(app.exec_())


####################################################################
class PartioningDemo(QWidget):
    def __init__(self, *args):
        QWidget.__init__(self, *args)

        self.device_list = parted.getAllDevices()
        device_list_widget = QWidget()
        device_list_layout = QHBoxLayout()
        self.device_selection_combobox = QComboBox()
        self.refresh_button = QPushButton("Refresh")
        self.refresh_button.pressed.connect(self.refresh_device_list)
        for disk in self.device_list:
            try:
                if parted.Disk(disk).type == "msdos":
                    self.device_selection_combobox.addItem(
                        "{} {} GB ({})".format(disk.model, format(disk.getSize(unit="GB"), '.2f'), disk.path),
                        userData=disk.path)
            except parted.DiskLabelException:
                disk = parted.freshDisk(disk, 'msdos')
                # For discarding CD-ROM like devices
                try:
                    disk.commit()
                except parted.IOException:
                    pass
                else:
                    disk = disk.device
                    self.device_selection_combobox.addItem(
                        "{} {} GB ({})".format(disk.model, format(disk.getSize(unit="GB"), '.2f'), disk.path),
                        userData=disk.path)

        self.device_selection_combobox.currentIndexChanged.connect(self.selected_device_changed)

        if self.device_selection_combobox.currentData():
            self.current_device = parted.getDevice(self.device_selection_combobox.currentData())

            self.current_disk = parted.Disk(self.current_device)

            self.partition_list_widget = QListWidget()

            for part in self.current_disk.partitions:

                part_info = self.get_part_info(part, "GB")

                part_item = QListWidgetItem(
                    "{}\t\t{} GB\t{}\t{}".format(part_info["path"], part_info["size"], part_info["file_system_type"],
                                                 part_info["flags"]))
                part_item.setData(Qt.UserRole, part_info["order"])
                if part_info["type"] == parted.PARTITION_NORMAL:
                    part_item.setIcon(QIcon("images/primary.xpm"))
                elif part_info["type"] == parted.PARTITION_EXTENDED:
                    part_item.setIcon(QIcon("images/extended.xpm"))
                elif part_info["type"] == parted.PARTITION_LOGICAL:
                    part_item.setIcon(QIcon("images/logical.xpm"))
                self.partition_list_widget.addItem(part_item)

            for free_space in self.current_disk.getFreeSpacePartitions():
                total_size = 0
                part_info = self.get_part_info(free_space, "GB")
                if float(part_info["size"]) > 0:
                    if part_info["type"] == 5:
                        extended_unallocated_item = QListWidgetItem(
                            "{}\t{} GB".format("Unallocated", part_info["size"]))
                        extended_unallocated_item.setIcon(QIcon("images/blank.xpm"))
                        extended_unallocated_item.setData(Qt.UserRole, "unallocated")
                        self.partition_list_widget.addItem(extended_unallocated_item)
                    if part_info["type"] == parted.PARTITION_FREESPACE:
                        total_size = total_size + float(part_info["size"])
                    unallocated_item = QListWidgetItem("{}\t{} GB".format("Unallocated", total_size))
                    unallocated_item.setIcon(QIcon("images/blank.xpm"))
                    unallocated_item.setData(Qt.UserRole, "unallocated")
                    self.partition_list_widget.addItem(unallocated_item)

        device_list_layout.addWidget(self.device_selection_combobox)
        device_list_layout.addWidget(self.refresh_button)
        device_list_widget.setLayout(device_list_layout)
        layout = QVBoxLayout()
        layout.addWidget(device_list_widget)
        layout.addWidget(self.partition_list_widget)
        partition_type_index_widget = QLabel()
        partition_type_index_widget.setPixmap(QPixmap("images/lejant.png"))
        partition_type_index_widget.setAlignment(Qt.AlignCenter)
        layout.addWidget(partition_type_index_widget)
        self.partition_list_widget.itemClicked.connect(self.on_part_select)
        self.partition_list_widget.itemDoubleClicked.connect(self.change_partition_format)

        op_widget = QWidget()
        op_buttons_layout = QHBoxLayout()
        self.new_part_btn = QPushButton("Add new partition")
        self.new_part_btn.pressed.connect(self.add_new_part)
        self.delete_part_btn = QPushButton("Delete partition")
        self.delete_part_btn.pressed.connect(self.delete_part)
        self.save_btn = QPushButton("Save")
        self.save_btn.pressed.connect(self.save)
        op_buttons_layout.addWidget(self.new_part_btn)
        op_buttons_layout.addWidget(self.delete_part_btn)
        op_buttons_layout.addWidget(self.save_btn)
        op_widget.setLayout(op_buttons_layout)
        layout.addWidget(op_widget)

        self.delete_part_btn.setEnabled(False)
        self.setLayout(layout)

    def save(self):
        self.current_disk.commit()

    def change_partition_format(self, selected_part):
        if selected_part.data(Qt.UserRole) != "unallocated":
            for part in self.current_disk.partitions:
                if part.number == selected_part.data(Qt.UserRole):
                    print(part.path)

    def get_part_info(self, part, size_unit):
        part_info = {}
        part_info["path"] = part.path
        if size_unit == "GB":
            part_info["size"] = format(part.getSize(unit=size_unit), '.2f')
        else:
            part_info["size"] = part.getSize(unit=size_unit)

        part_info["file_system_type"] = "Unknown"

        if part.fileSystem:
            if part.fileSystem.type.startswith('linux-swap'):
                part_info["file_system_type"] = "swap"
            else:
                part_info["file_system_type"] = part.fileSystem.type

        try:
            part_info["flags"] = part.getFlagsAsString()
        except:
            pass
        part_info["order"] = part.number
        part_info["type"] = part.type

        return part_info

    def refresh_device_list(self):
        self.device_selection_combobox.clear()
        self.device_list = parted.getAllDevices()
        for disk in self.device_list:
            try:
                if parted.Disk(disk).type == "msdos":
                    self.device_selection_combobox.addItem(
                        "{} {} GB ({})".format(disk.model, format(disk.getSize(unit="GB"), '.2f'), disk.path),
                        userData=disk.path)
            except parted.DiskLabelException:
                disk = parted.freshDisk(disk, 'msdos')
                # For discarding CD-ROM like devices
                try:
                    disk.commit()
                except parted.IOException:
                    pass
                else:
                    disk = disk.device
                    self.device_selection_combobox.addItem(
                        "{} {} GB ({})".format(disk.model, format(disk.getSize(unit="GB"), '.2f'), disk.path),
                        userData=disk.path)

    def selected_device_changed(self):
        if self.device_selection_combobox.currentData():
            self.current_device = parted.getDevice(self.device_selection_combobox.currentData())
            self.current_disk = parted.Disk(self.current_device)
            self.refresh_partition_list()

    def refresh_partition_list(self):
        self.partition_list_widget.clear()
        for part in self.current_disk.partitions:

            part_info = self.get_part_info(part, "GB")

            item = QListWidgetItem(
                "{}\t\t{} GB\t{}\t{}".format(part_info["path"], part_info["size"], part_info["file_system_type"],
                                             part_info["flags"]))
            item.setData(Qt.UserRole, part_info["order"])
            if part_info["type"] == parted.PARTITION_NORMAL:
                item.setIcon(QIcon("images/primary.xpm"))
            elif part_info["type"] == parted.PARTITION_EXTENDED:
                item.setIcon(QIcon("images/extended.xpm"))
            elif part_info["type"] == parted.PARTITION_LOGICAL:
                item.setIcon(QIcon("images/logical.xpm"))
            self.partition_list_widget.addItem(item)

        for free_part in self.current_disk.getFreeSpacePartitions():
            total_size = 0
            part_info = self.get_part_info(free_part, "GB")
            if float(part_info["size"]) > 0:
                if part_info["type"] == 5:
                    extended_unallocated_item = QListWidgetItem("{}\t{} GB".format("Unallocated", part_info["size"]))
                    extended_unallocated_item.setIcon(QIcon("images/blank.xpm"))
                    extended_unallocated_item.setData(Qt.UserRole, "unallocated")
                    self.partition_list_widget.addItem(extended_unallocated_item)
                if part_info["type"] == parted.PARTITION_FREESPACE:
                    total_size = total_size + float(part_info["size"])
                unallocated_item = QListWidgetItem("{}\t{} GB".format("Unallocated", total_size))
                unallocated_item.setIcon(QIcon("images/blank.xpm"))
                unallocated_item.setData(Qt.UserRole, "unallocated")
                self.partition_list_widget.addItem(unallocated_item)

    def on_part_select(self, selected_part):

        if selected_part.data(Qt.UserRole) != "unallocated":
            self.delete_part_btn.setEnabled(True)
        else:
            self.delete_part_btn.setEnabled(False)

    def delete_part(self):
        part_order = self.partition_list_widget.currentItem().data(Qt.UserRole)
        for part in self.current_disk.partitions:
            if part.number == part_order:
                try:
                    self.current_disk.deletePartition(part)
                    self.refresh_partition_list()
                except parted.PartitionException:
                    QMessageBox.warning(self, "Warning",
                                        "Extended partitions cannot be deleted before logical partitions.")
        self.partition_list_widget.setCurrentRow(self.partition_list_widget.count() - 2)

    def add_new_part(self):
        if self.get_max_unallocated_region():
            region = self.get_max_unallocated_region()
            number_of_primary_parts = len(self.current_disk.getPrimaryPartitions())
            number_of_extended_parts = ext_count = 1 if self.current_disk.getExtendedPartition() else 0
            parts_avail = self.current_disk.maxPrimaryPartitionCount - (
                        number_of_primary_parts + number_of_extended_parts)
            if not parts_avail and not ext_count:
                QMessageBox.warning(self,
                                    "Warning",
                                    """You cannot create more than 4 primary partition.
                                     If you want more partitions, 
                                     delete one of primary partition and create extended one. """
                                    )
            else:
                if parts_avail:
                    if not number_of_extended_parts and parts_avail > 1:
                        self.create_part(region, parted.PARTITION_NORMAL)
                        self.refresh_partition_list()
                    elif parts_avail == 1:
                        self.create_part(region, parted.PARTITION_EXTENDED)
                        self.refresh_partition_list()

                if number_of_extended_parts:
                    ext_part = self.current_disk.getExtendedPartition()
                    try:
                        region = ext_part.geometry.intersect(region)
                    except ArithmeticError:
                        QMessageBox.critical(self, "Error",
                                             "There is no free space for new partition!"
                                             " Try increase size for extended partitions.")
                    else:
                        self.create_part(region, parted.PARTITION_LOGICAL)
                        self.refresh_partition_list()


        else:
            QMessageBox.critical(self, "There is no free space for new partition!")

    def get_max_unallocated_region(self):
        max_size = -1
        region = None
        alignment = self.current_device.optimumAlignment

        for _region in self.current_disk.getFreeSpaceRegions():
            if _region.length > max_size and _region.length > alignment.grainSize:
                region = _region
                max_size = _region.length

        return region

    def create_part(self, region, part_type):

        if part_type == parted.PARTITION_NORMAL or part_type == parted.PARTITION_EXTENDED:
            for free_part in self.current_disk.getFreeSpacePartitions():
                part_info = self.get_part_info(free_part, "GB")
                if part_info["type"] == parted.PARTITION_FREESPACE:
                    max_size = float(part_info["size"])
        elif part_type == part_type == parted.PARTITION_LOGICAL:
            for free_part in self.current_disk.getFreeSpacePartitions():
                part_info = self.get_part_info(free_part, "GB")
                if part_info["type"] == 5:
                    max_size = float(part_info["size"])

        alignment = self.current_device.optimalAlignedConstraint
        constraint = self.current_device.getConstraint()
        data = {
            'start': constraint.startAlign.alignUp(region, region.start),
            'end': constraint.endAlign.alignDown(region, region.end),
        }

        selected_size, ok = QInputDialog().getDouble(self, 'Create partition', 'Size in GB:', min=0.001, value=1,
                                                     max=max_size, decimals=3)

        if ok:
            data["end"] = int(data["start"]) + int(
                parted.sizeToSectors(float(selected_size), "GiB", self.current_device.sectorSize))
            try:
                geometry = parted.Geometry(device=self.current_device, start=int(data["start"]), end=int(data["end"]))
                partition = parted.Partition(
                    disk=self.current_disk,
                    type=part_type,
                    geometry=geometry,
                )

                self.current_disk.addPartition(partition=partition, constraint=constraint)
            except (parted.PartitionException, parted.GeometryException, parted.CreateException) as e:
                # GeometryException raised when incorrect start/end values applied (e.g. start < end),
                # CreateException is raised when the partition doesn't fit on the disk, etc.
                # PartedException is raised on generic errors (e.g. start/end values out of range)
                raise RuntimeError(e.message)


if __name__ == "__main__":
    main()