Rename Partition Entry

gpt_image/partition.py

@@ -1,5 +1,4 @@
 import uuid
-from dataclasses import dataclass
 from math import ceil
 
 from gpt_image.disk import Geometry
@@ -13,15 +12,9 @@ class Partition:
     from a table's partition list.
     """
 
-    @dataclass
-    class Type:
-        """GPT Partition Types
-
-        https://en.wikipedia.org/wiki/GUID_Partition_Table#Partition_entries
-        """
-
-        LinuxFileSystem = "0FC63DAF-8483-4772-8E79-3D69D8477DE4"
-        EFISystemPartition = "C12A7328-F81F-11D2-BA4B-00A0C93EC93B"
+    # https://en.wikipedia.org/wiki/GUID_Partition_Table#Partition_entries
+    LINUX_FILE_SYSTEM = "0FC63DAF-8483-4772-8E79-3D69D8477DE4"
+    EFI_SYSTEM_PARTITION = "C12A7328-F81F-11D2-BA4B-00A0C93EC93B"
 
     def __init__(
         self,
@@ -49,7 +42,7 @@ def __init__(
 
         # if name is set, this isn't an empty partition. Set relevant fields
         if name:
-            self.type_guid.data = uuid.UUID(Partition.Type.LinuxFileSystem).bytes_le
+            self.type_guid.data = uuid.UUID(Partition.LINUX_FILE_SYSTEM).bytes_le
             if not partition_guid:
                 self.partition_guid.data = uuid.uuid4().bytes_le
             else:
@@ -76,14 +69,14 @@ def as_bytes(self) -> bytes:
         return b"".join(byte_list)
 
 
-class PartitionEntry:
+class PartitionEntryArray:
     """Stores the Partition objects for a Table"""
 
     def __init__(self, geometry: Geometry, entry_size: int = 128):
         self.entries = [Partition()] * entry_size
         self._geometry = geometry
 
-    def add(self, partition: Partition):
+    def add(self, partition: Partition) -> None:
         """Add a partition to the entries
 
         Appends the Partition to the next available entry. Calculates the
@@ -136,12 +129,11 @@ def _get_last_lba(self, partition: Partition) -> int:
 
     def _get_next_partition(self) -> int:
         """Return the index of the next unused partition"""
-
-        # @TODO: handle error if partition not found
         for idx, part in enumerate(self.entries):
             # return the first partition index that has no name
             if int.from_bytes(part.partition_name.data, byteorder="little") == 0:
                 return idx
+        raise AttributeError("No partition entries available")
 
     def as_bytes(self) -> bytes:
         """Represent as bytes

gpt_image/table.py

@@ -7,7 +7,7 @@
 
 from gpt_image.disk import Disk, Geometry
 from gpt_image.entry import Entry
-from gpt_image.partition import Partition, PartitionEntry
+from gpt_image.partition import Partition, PartitionEntryArray
 
 
 class ProtectiveMBR:
@@ -39,7 +39,7 @@ def __init__(self, geometry: Geometry):
             self.partition_size,
         ]
 
-    def as_bytes(self):
+    def as_bytes(self) -> bytes:
         """Get the Protective MBR as bytes
 
         Does not include the signature
@@ -132,16 +132,16 @@ class Table:
     directly used.
     """
 
-    def __init__(self, disk: Disk, sector_size: int = 512) -> None:
+    def __init__(self, disk: Disk, sector_size: int = 512):
         self.disk = disk
         self.geometry = disk.geometry
         self.protective_mbr = ProtectiveMBR(self.geometry)
         self.primary_header = Header(self.geometry)
         self.secondary_header = Header(self.geometry, is_backup=True)
 
-        self.partitions = PartitionEntry(self.geometry)
+        self.partitions = PartitionEntryArray(self.geometry)
 
-    def write(self):
+    def write(self) -> None:
         """Write the table to disk"""
         # calculate partition checksum and write to header
         self.checksum_partitions(self.primary_header)
@@ -176,7 +176,7 @@ def write(self):
 
     def create_partition(
         self, name: str, size: int, guid: uuid.UUID, alignment: int = 8
-    ):
+    ) -> None:
         part = Partition(
             name,
             size,
@@ -185,14 +185,14 @@ def create_partition(
         )
         self.partitions.add(part)
 
-    def checksum_partitions(self, header: Header):
+    def checksum_partitions(self, header: Header) -> None:
         """Checksum the partition entries"""
         part_entry_bytes = self.partitions.as_bytes()
         header.partition_array_crc.data = binascii.crc32(part_entry_bytes).to_bytes(
             4, "little"
         )
 
-    def checksum_header(self, header: Header):
+    def checksum_header(self, header: Header) -> None:
         """Checksum the table header
 
         This CRC includes the partition checksum, and must be calculated