#10619: refactor wic codebase. Part 3

index a9144e2..fefe88e 100644
@@ -36,25 +36,8 @@ from wic.plugin import pluginmgr
 from wic.pluginbase import ImagerPlugin
 from wic.utils.errors import CreatorError, ImageError
 from wic.utils.misc import get_bitbake_var, exec_cmd, exec_native_cmd
-from wic.utils.partitionedfs import Image
+from wic.utils.partitionedfs import PartitionedImage
 
-class DiskImage():
-    """
-    A Disk backed by a file.
-    """
-    def __init__(self, device, size):
-        self.size = size
-        self.device = device
-        self.created = False
-
-    def create(self):
-        if self.created:
-            return
-        # create sparse disk image
-        with open(self.device, 'w') as sparse:
-            os.ftruncate(sparse.fileno(), self.size)
-
-        self.created = True
 
 class DirectPlugin(ImagerPlugin):
     """

@@ -189,9 +172,10 @@ class DirectPlugin(ImagerPlugin):
         filesystems from the artifacts directly and combine them into
         a partitioned image.
         """
-        self._image = Image(self.native_sysroot)
+        image_path = self._full_path(self.workdir, self.parts[0].disk, "direct")
+        self._image = PartitionedImage(image_path, self.ptable_format,
+                                       self.native_sysroot)
 
-        disk_ids = {}
         for num, part in enumerate(self.parts, 1):
             # as a convenience, set source to the boot partition source
             # instead of forcing it to be set via bootloader --source

@@ -203,10 +187,8 @@ class DirectPlugin(ImagerPlugin):
                 if self.ptable_format == 'gpt':
                     part.uuid = str(uuid.uuid4())
                 else: # msdos partition table
-                    if part.disk not in disk_ids:
-                        disk_ids[part.disk] = int.from_bytes(os.urandom(4), 'little')
-                    disk_id = disk_ids[part.disk]
-                    part.uuid = '%0x-%02d' % (disk_id, self._get_part_num(num, self.parts))
+                    part.uuid = '%0x-%02d' % (self._image.identifier,
+                                              self._get_part_num(num, self.parts))
 
         fstab_path = self._write_fstab(self.rootfs_dir.get("ROOTFS_DIR"))
 

@@ -225,11 +207,6 @@ class DirectPlugin(ImagerPlugin):
                         part.size = int(round(float(rsize_bb)))
             # need to create the filesystems in order to get their
             # sizes before we can add them and do the layout.
-            # Image.create() actually calls __format_disks() to create
-            # the disk images and carve out the partitions, then
-            # self.assemble() calls Image.assemble() which calls
-            # __write_partitition() for each partition to dd the fs
-            # into the partitions.
             part.prepare(self, self.workdir, self.oe_builddir, self.rootfs_dir,
                          self.bootimg_dir, self.kernel_dir, self.native_sysroot)
 

@@ -238,26 +215,14 @@ class DirectPlugin(ImagerPlugin):
         if fstab_path:
             shutil.move(fstab_path + ".orig", fstab_path)
 
-        self._image.layout_partitions(self.ptable_format)
-
-        for disk_name, disk in self._image.disks.items():
-            full_path = self._full_path(self.workdir, disk_name, "direct")
-            msger.debug("Adding disk %s as %s with size %s bytes" \
-                        % (disk_name, full_path, disk['min_size']))
-            disk_obj = DiskImage(full_path, disk['min_size'])
-            self._image.add_disk(disk_name, disk_obj, disk_ids.get(disk_name))
-
+        self._image.layout_partitions()
         self._image.create()
 
     def assemble(self):
         """
-        Assemble partitions into disk image(s)
+        Assemble partitions into disk image
         """
-        for disk_name, disk in self._image.disks.items():
-            full_path = self._full_path(self.workdir, disk_name, "direct")
-            msger.debug("Assembling disk %s as %s with size %s bytes" \
-                        % (disk_name, full_path, disk['min_size']))
-            self._image.assemble(full_path)
+        self._image.assemble()
 
     def finalize(self):
         """

@@ -267,26 +232,25 @@ class DirectPlugin(ImagerPlugin):
         creating and installing a bootloader configuration.
         """
         source_plugin = self.ks.bootloader.source
+        disk_name = self.parts[0].disk
         if source_plugin:
             name = "do_install_disk"
             methods = pluginmgr.get_source_plugin_methods(source_plugin,
                                                           {name: None})
-            for disk_name, disk in self._image.disks.items():
-                methods["do_install_disk"](disk, disk_name, self, self.workdir,
-                                           self.oe_builddir, self.bootimg_dir,
-                                           self.kernel_dir, self.native_sysroot)
-
-        for disk_name, disk in self._image.disks.items():
-            full_path = self._full_path(self.workdir, disk_name, "direct")
-            # Generate .bmap
-            if self.bmap:
-                msger.debug("Generating bmap file for %s" % disk_name)
-                exec_native_cmd("bmaptool create %s -o %s.bmap" % (full_path, full_path),
-                                self.native_sysroot)
-            # Compress the image
-            if self.compressor:
-                msger.debug("Compressing disk %s with %s" % (disk_name, self.compressor))
-                exec_cmd("%s %s" % (self.compressor, full_path))
+            methods["do_install_disk"](self._image, disk_name, self, self.workdir,
+                                       self.oe_builddir, self.bootimg_dir,
+                                       self.kernel_dir, self.native_sysroot)
+
+        full_path = self._image.path
+        # Generate .bmap
+        if self.bmap:
+            msger.debug("Generating bmap file for %s" % disk_name)
+            exec_native_cmd("bmaptool create %s -o %s.bmap" % (full_path, full_path),
+                            self.native_sysroot)
+        # Compress the image
+        if self.compressor:
+            msger.debug("Compressing disk %s with %s" % (disk_name, self.compressor))
+            exec_cmd("%s %s" % (self.compressor, full_path))
 
     def print_info(self):
         """

@@ -294,13 +258,12 @@ class DirectPlugin(ImagerPlugin):
         """
         msg = "The new image(s) can be found here:\n"
 
-        for disk_name in self._image.disks:
-            extension = "direct" + {"gzip": ".gz",
-                                    "bzip2": ".bz2",
-                                    "xz": ".xz",
-                                    None: ""}.get(self.compressor)
-            full_path = self._full_path(self.outdir, disk_name, extension)
-            msg += '  %s\n\n' % full_path
+        extension = "direct" + {"gzip": ".gz",
+                                "bzip2": ".bz2",
+                                "xz": ".xz",
+                                None: ""}.get(self.compressor)
+        full_path = self._full_path(self.outdir, self.parts[0].disk, extension)
+        msg += '  %s\n\n' % full_path
 
         msg += 'The following build artifacts were used to create the image(s):\n'
         for part in self.parts:

index 4b9b265..0be2b78 100644
@@ -63,7 +63,7 @@ class BootimgPcbiosPlugin(SourcePlugin):
 
         full_path = creator._full_path(workdir, disk_name, "direct")
         msger.debug("Installing MBR on disk %s as %s with size %s bytes" \
-                    % (disk_name, full_path, disk['min_size']))
+                    % (disk_name, full_path, disk.min_size))
 
         rcode = runner.show(['dd', 'if=%s' % mbrfile,
                              'of=%s' % full_path, 'conv=notrunc'])

index ca28bc0..fb34235 100644
@@ -473,13 +473,12 @@ class IsoImagePlugin(SourcePlugin):
         utility for booting via BIOS from disk storage devices.
         """
 
+        iso_img = "%s.p1" % disk.path
         full_path = creator._full_path(workdir, disk_name, "direct")
-        iso_img = "%s.p1" % full_path
         full_path_iso = creator._full_path(workdir, disk_name, "iso")
 
         isohybrid_cmd = "isohybrid -u %s" % iso_img
-        msger.debug("running command: %s" % \
-                    isohybrid_cmd)
+        msger.debug("running command: %s" % isohybrid_cmd)
         exec_native_cmd(isohybrid_cmd, native_sysroot)
 
         # Replace the image created by direct plugin with the one created by

@@ -487,9 +486,6 @@ class IsoImagePlugin(SourcePlugin):
         # mkisofs has a very specific MBR is system area of the ISO image, and
         # direct plugin adds and configures an another MBR.
         msger.debug("Replaceing the image created by direct plugin\n")
-        os.remove(full_path)
+        os.remove(disk.path)
         shutil.copy2(iso_img, full_path_iso)
         shutil.copy2(full_path_iso, full_path)
-
-        # Remove temporary ISO file
-        os.remove(iso_img)

index cb1da93..9e79a13 100644
@@ -204,9 +204,9 @@ class RootfsPlugin(SourcePlugin):
         if not os.path.exists(mbrfile):
             msger.error("Couldn't find %s. Has syslinux-native been baked?" % mbrfile)
 
-        full_path = disk['disk'].device
+        full_path = disk.path
         msger.debug("Installing MBR on disk %s as %s with size %s bytes" \
-                    % (disk_name, full_path, disk['min_size']))
+                    % (disk_name, full_path, disk.min_size))
 
         ret_code = runner.show(['dd', 'if=%s' % mbrfile, 'of=%s' % full_path, 'conv=notrunc'])
         if ret_code != 0:

350 lines not on this page · raw message

index fefe88e..1204431 100644
@@ -31,13 +31,12 @@ import tempfile
 from time import strftime
 
 from wic import msger
+from wic.filemap import sparse_copy
 from wic.ksparser import KickStart, KickStartError
 from wic.plugin import pluginmgr
 from wic.pluginbase import ImagerPlugin
 from wic.utils.errors import CreatorError, ImageError
 from wic.utils.misc import get_bitbake_var, exec_cmd, exec_native_cmd
-from wic.utils.partitionedfs import PartitionedImage
-
 
 class DirectPlugin(ImagerPlugin):
     """

@@ -316,3 +315,278 @@ class DirectPlugin(ImagerPlugin):
 
         # remove work directory
         shutil.rmtree(self.workdir, ignore_errors=True)
+
+# Overhead of the MBR partitioning scheme (just one sector)
+MBR_OVERHEAD = 1
+
+# Overhead of the GPT partitioning scheme
+GPT_OVERHEAD = 34
+
+# Size of a sector in bytes
+SECTOR_SIZE = 512
+
+class PartitionedImage():
+    """
+    Partitioned image in a file.
+    """
+
+    def __init__(self, path, ptable_format, native_sysroot=None):
+        self.path = path  # Path to the image file
+        self.numpart = 0  # Number of allocated partitions
+        self.realpart = 0 # Number of partitions in the partition table
+        self.offset = 0   # Offset of next partition (in sectors)
+        self.min_size = 0 # Minimum required disk size to fit
+                          # all partitions (in bytes)
+        self.ptable_format = ptable_format  # Partition table format
+        # Disk system identifier
+        self.identifier = int.from_bytes(os.urandom(4), 'little')
+
+        self.partitions = []
+        self.partimages = []
+        # Size of a sector used in calculations
+        self.sector_size = SECTOR_SIZE
+        self.native_sysroot = native_sysroot
+
+    def add_partition(self, part):
+        """
+        Add the next partition. Partitions have to be added in the
+        first-to-last order.
+        """
+        part.ks_pnum = len(self.partitions)
+
+        # Converting kB to sectors for parted
+        part.size_sec = part.disk_size * 1024 // self.sector_size
+
+        self.partitions.append(part)
+
+    def layout_partitions(self):
+        """ Layout the partitions, meaning calculate the position of every
+        partition on the disk. The 'ptable_format' parameter defines the
+        partition table format and may be "msdos". """
+
+        msger.debug("Assigning %s partitions to disks" % self.ptable_format)
+
+        # Go through partitions in the order they are added in .ks file
+        for num in range(len(self.partitions)):
+            part = self.partitions[num]
+
+            if self.ptable_format == 'msdos' and part.part_type:
+                # The --part-type can also be implemented for MBR partitions,
+                # in which case it would map to the 1-byte "partition type"
+                # filed at offset 3 of the partition entry.
+                raise ImageError("setting custom partition type is not " \
+                                 "implemented for msdos partitions")
+
+            # Get the disk where the partition is located
+            self.numpart += 1
+            if not part.no_table:
+                self.realpart += 1
+
+            if self.numpart == 1:
+                if self.ptable_format == "msdos":
+                    overhead = MBR_OVERHEAD
+                elif self.ptable_format == "gpt":
+                    overhead = GPT_OVERHEAD
+
+                # Skip one sector required for the partitioning scheme overhead
+                self.offset += overhead
+
+            if self.realpart > 3:
+                # Reserve a sector for EBR for every logical partition
+                # before alignment is performed.
+                if self.ptable_format == "msdos":
+                    self.offset += 1
+
+            if part.align:
+                # If not first partition and we do have alignment set we need
+                # to align the partition.
+                # FIXME: This leaves a empty spaces to the disk. To fill the
+                # gaps we could enlargea the previous partition?
+
+                # Calc how much the alignment is off.
+                align_sectors = self.offset % (part.align * 1024 // self.sector_size)
+
+                if align_sectors:
+                    # If partition is not aligned as required, we need
+                    # to move forward to the next alignment point
+                    align_sectors = (part.align * 1024 // self.sector_size) - align_sectors
+
+                    msger.debug("Realignment for %s%s with %s sectors, original"
+                                " offset %s, target alignment is %sK." %
+                                (part.disk, self.numpart, align_sectors,
+                                 self.offset, part.align))
+
+                    # increase the offset so we actually start the partition on right alignment
+                    self.offset += align_sectors
+
+            part.start = self.offset
+            self.offset += part.size_sec
+
+            part.type = 'primary'
+            if not part.no_table:
+                part.num = self.realpart
+            else:
+                part.num = 0
+
+            if self.ptable_format == "msdos":
+                # only count the partitions that are in partition table
+                if len([p for p in self.partitions if not p.no_table]) > 4:
+                    if self.realpart > 3:
+                        part.type = 'logical'
+                        part.num = self.realpart + 1
+
+            msger.debug("Assigned %s to %s%d, sectors range %d-%d size %d "
+                        "sectors (%d bytes)." \
+                            % (part.mountpoint, part.disk, part.num,
+                               part.start, self.offset - 1,
+                               part.size_sec, part.size_sec * self.sector_size))
+
+        # Once all the partitions have been layed out, we can calculate the
+        # minumim disk size
+        self.min_size = self.offset
+        if self.ptable_format == "gpt":
+            self.min_size += GPT_OVERHEAD
+
+        self.min_size *= self.sector_size
+
+    def _create_partition(self, device, parttype, fstype, start, size):
+        """ Create a partition on an image described by the 'device' object. """
+
+        # Start is included to the size so we need to substract one from the end.
+        end = start + size - 1
+        msger.debug("Added '%s' partition, sectors %d-%d, size %d sectors" %
+                    (parttype, start, end, size))
+
+        cmd = "parted -s %s unit s mkpart %s" % (device, parttype)
+        if fstype:
+            cmd += " %s" % fstype
+        cmd += " %d %d" % (start, end)
+
+        return exec_native_cmd(cmd, self.native_sysroot)
+
+    def create(self):
+        msger.debug("Creating sparse file %s" % self.path)
+        with open(self.path, 'w') as sparse:
+            os.ftruncate(sparse.fileno(), self.min_size)
+
+        msger.debug("Initializing partition table for %s" % self.path)
+        exec_native_cmd("parted -s %s mklabel %s" %
+                        (self.path, self.ptable_format), self.native_sysroot)
+
+        msger.debug("Set disk identifier %x" % self.identifier)
+        with open(self.path, 'r+b') as img:
+            img.seek(0x1B8)
+            img.write(self.identifier.to_bytes(4, 'little'))
+
+        msger.debug("Creating partitions")
+
+        for part in self.partitions:
+            if part.num == 0:
+                continue
+
+            if self.ptable_format == "msdos" and part.num == 5:
+                # Create an extended partition (note: extended
+                # partition is described in MBR and contains all
+                # logical partitions). The logical partitions save a
+                # sector for an EBR just before the start of a
+                # partition. The extended partition must start one
+                # sector before the start of the first logical
+                # partition. This way the first EBR is inside of the
+                # extended partition. Since the extended partitions
+                # starts a sector before the first logical partition,
+                # add a sector at the back, so that there is enough
+                # room for all logical partitions.
+                self._create_partition(self.path, "extended",
+                                       None, part.start - 1,
+                                       self.offset - part.start + 1)
+
+            if part.fstype == "swap":
+                parted_fs_type = "linux-swap"
+            elif part.fstype == "vfat":
+                parted_fs_type = "fat32"
+            elif part.fstype == "msdos":
+                parted_fs_type = "fat16"
+            elif part.fstype == "ontrackdm6aux3":
+                parted_fs_type = "ontrackdm6aux3"
+            else:
+                # Type for ext2/ext3/ext4/btrfs
+                parted_fs_type = "ext2"
+
+            # Boot ROM of OMAP boards require vfat boot partition to have an
+            # even number of sectors.
+            if part.mountpoint == "/boot" and part.fstype in ["vfat", "msdos"] \
+               and part.size_sec % 2:
+                msger.debug("Subtracting one sector from '%s' partition to " \
+                            "get even number of sectors for the partition" % \
+                            part.mountpoint)
+                part.size_sec -= 1
+
+            self._create_partition(self.path, part.type,
+                                   parted_fs_type, part.start, part.size_sec)
+
+            if part.part_type:
+                msger.debug("partition %d: set type UID to %s" % \
+                            (part.num, part.part_type))
+                exec_native_cmd("sgdisk --typecode=%d:%s %s" % \
+                                         (part.num, part.part_type,
+                                          self.path), self.native_sysroot)
+
+            if part.uuid and self.ptable_format == "gpt":
+                msger.debug("partition %d: set UUID to %s" % \
+                            (part.num, part.uuid))
+                exec_native_cmd("sgdisk --partition-guid=%d:%s %s" % \
+                                (part.num, part.uuid, self.path),
+                                self.native_sysroot)
+
+            if part.label and self.ptable_format == "gpt":
+                msger.debug("partition %d: set name to %s" % \
+                            (part.num, part.label))
+                exec_native_cmd("parted -s %s name %d %s" % \
+                                (self.path, part.num, part.label),
+                                self.native_sysroot)
+
+            if part.active:
+                flag_name = "legacy_boot" if self.ptable_format == 'gpt' else "boot"
+                msger.debug("Set '%s' flag for partition '%s' on disk '%s'" % \
+                            (flag_name, part.num, self.path))
+                exec_native_cmd("parted -s %s set %d %s on" % \
+                                (self.path, part.num, flag_name),
+                                self.native_sysroot)
+            if part.system_id:
+                exec_native_cmd("sfdisk --part-type %s %s %s" % \
+                                (self.path, part.num, part.system_id),
+                                self.native_sysroot)
+
+            # Parted defaults to enabling the lba flag for fat16 partitions,
+            # which causes compatibility issues with some firmware (and really
+            # isn't necessary).
+            if parted_fs_type == "fat16":
+                if self.ptable_format == 'msdos':
+                    msger.debug("Disable 'lba' flag for partition '%s' on disk '%s'" % \
+                                (part.num, self.path))
+                    exec_native_cmd("parted -s %s set %d lba off" % \
+                                    (self.path, part.num),
+                                    self.native_sysroot)
+
+    def cleanup(self):
+        # remove partition images
+        for image in self.partimages:
+            os.remove(image)
+
+    def assemble(self):
+        msger.debug("Installing partitions")
+
+        for part in self.partitions:
+            source = part.source_file
+            if source:
+                # install source_file contents into a partition
+                sparse_copy(source, self.path, part.start * self.sector_size)
+
+                msger.debug("Installed %s in partition %d, sectors %d-%d, "
+                            "size %d sectors" % \
+                            (source, part.num, part.start,
+                             part.start + part.size_sec - 1, part.size_sec))
+
+                partimage = self.path + '.p%d' % part.num
+                os.rename(source, partimage)
+                self.partimages.append(partimage)

301 lines not on this page · raw message

index 1204431..a2bc4c4 100644
@@ -149,15 +149,6 @@ class DirectPlugin(ImagerPlugin):
 
         return updated
 
-    def set_bootimg_dir(self, bootimg_dir):
-        """
-        Accessor for bootimg_dir, the actual location used for the source
-        of the bootimg.  Should be set by source plugins (only if they
-        change the default bootimg source) so the correct info gets
-        displayed for print_outimage_info().
-        """
-        self.bootimg_dir = bootimg_dir
-
     def _full_path(self, path, name, extention):
         """ Construct full file path to a file we generate. """
         return os.path.join(path, "%s-%s.%s" % (self.name, name, extention))

index dd6c920..95316c8 100644
@@ -187,7 +187,7 @@ class BootimgEFIPlugin(SourcePlugin):
             if not bootimg_dir:
                 msger.error("Couldn't find DEPLOY_DIR_IMAGE, exiting\n")
             # just so the result notes display it
-            creator.set_bootimg_dir(bootimg_dir)
+            creator.bootimg_dir = bootimg_dir
 
         staging_kernel_dir = kernel_dir
 

index 0be2b78..e5f6a32 100644
@@ -151,7 +151,7 @@ class BootimgPcbiosPlugin(SourcePlugin):
             if not _has_syslinux(bootimg_dir):
                 msger.error("Please build syslinux first\n")
             # just so the result notes display it
-            creator.set_bootimg_dir(bootimg_dir)
+            creator.bootimg_dir = bootimg_dir
 
         staging_kernel_dir = kernel_dir
 

deleted file mode 100644
index 0a81575..0000000
@@ -1,22 +0,0 @@
-#
-# OpenEmbedded wic utils library
-#
-# Copyright (c) 2013, Intel Corporation.
-# All rights reserved.
-#
-# This program is free software; you can redistribute it and/or modify
-# it under the terms of the GNU General Public License version 2 as
-# published by the Free Software Foundation.
-#
-# This program is distributed in the hope that it will be useful,
-# but WITHOUT ANY WARRANTY; without even the implied warranty of
-# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
-# GNU General Public License for more details.
-#
-# You should have received a copy of the GNU General Public License along
-# with this program; if not, write to the Free Software Foundation, Inc.,
-# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
-#
-# AUTHORS
-# Tom Zanussi <tom.zanussi (at] linux.intel.com>
-#

index 9bcccb3..d0a1cad 100644
@@ -35,7 +35,7 @@ from wic.filemap import sparse_copy
 from wic.ksparser import KickStart, KickStartError
 from wic.plugin import pluginmgr
 from wic.pluginbase import ImagerPlugin
-from wic.utils.errors import CreatorError, ImageError
+from wic.utils.errors import ImageError
 from wic.utils.misc import get_bitbake_var, exec_cmd, exec_native_cmd
 
 class DirectPlugin(ImagerPlugin):

@@ -83,8 +83,6 @@ class DirectPlugin(ImagerPlugin):
             self.assemble()
             self.finalize()
             self.print_info()
-        except CreatorError:
-            raise
         finally:
             self.cleanup()
 

index 41d3cc6..d9fa805 100644
@@ -135,7 +135,7 @@ class KickStart():
         part.add_argument('--fstype')
         part.add_argument('--label')
         part.add_argument('--no-table', action='store_true')
-        part.add_argument('--ondisk', '--ondrive', dest='disk')
+        part.add_argument('--ondisk', '--ondrive', dest='disk', default='sda')
         part.add_argument("--overhead-factor", type=overheadtype)
         part.add_argument('--part-type')
         part.add_argument('--rootfs-dir')

index a2bc4c4..9bcccb3 100644
@@ -290,10 +290,7 @@ class DirectPlugin(ImagerPlugin):
 
     def cleanup(self):
         if self._image:
-            try:
-                self._image.cleanup()
-            except ImageError as err:
-                msger.warning("%s" % err)
+            self._image.cleanup()
 
         # Move results to the output dir
         if not os.path.exists(self.outdir):

index d0a1cad..ffe6c84 100644
@@ -74,6 +74,19 @@ class DirectPlugin(ImagerPlugin):
         self.ptable_format = self.ks.bootloader.ptable
         self.parts = self.ks.partitions
 
+        # calculate the real partition number, accounting for partitions not
+        # in the partition table and logical partitions
+        realnum = 0
+        for part in self.parts:
+            if part.no_table:
+                part.realnum = 0
+            else:
+                realnum += 1
+                if self.ptable_format == 'msdos' and realnum > 3:
+                    part.realnum = realnum + 1
+                    continue
+                part.realnum = realnum
+
     def do_create(self):
         """
         Plugin entry point.

@@ -86,22 +99,6 @@ class DirectPlugin(ImagerPlugin):
         finally:
             self.cleanup()
 
-    def _get_part_num(self, num, parts):
-        """calculate the real partition number, accounting for partitions not
-        in the partition table and logical partitions
-        """
-        realnum = 0
-        for pnum, part in enumerate(parts, 1):
-            if not part.no_table:
-                realnum += 1
-            if pnum == num:
-                if  part.no_table:
-                    return 0
-                if self.ptable_format == 'msdos' and realnum > 3:
-                    # account for logical partition numbering, ex. sda5..
-                    return realnum + 1
-                return realnum
-
     def _write_fstab(self, image_rootfs):
         """overriden to generate fstab (temporarily) in rootfs. This is called
         from _create, make sure it doesn't get called from

@@ -128,15 +125,14 @@ class DirectPlugin(ImagerPlugin):
     def _update_fstab(self, fstab_lines, parts):
         """Assume partition order same as in wks"""
         updated = False
-        for num, part in enumerate(parts, 1):
-            pnum = self._get_part_num(num, parts)
-            if not pnum or not part.mountpoint \
+        for part in parts:
+            if not part.realnum or not part.mountpoint \
                or part.mountpoint in ("/", "/boot"):
                 continue
 
             # mmc device partitions are named mmcblk0p1, mmcblk0p2..
             prefix = 'p' if  part.disk.startswith('mmcblk') else ''
-            device_name = "/dev/%s%s%d" % (part.disk, prefix, pnum)
+            device_name = "/dev/%s%s%d" % (part.disk, prefix, part.realnum)
 
             opts = part.fsopts if part.fsopts else "defaults"
             line = "\t".join([device_name, part.mountpoint, part.fstype,

@@ -164,7 +160,7 @@ class DirectPlugin(ImagerPlugin):
         self._image = PartitionedImage(image_path, self.ptable_format,
                                        self.native_sysroot)
 
-        for num, part in enumerate(self.parts, 1):
+        for part in self.parts:
             # as a convenience, set source to the boot partition source
             # instead of forcing it to be set via bootloader --source
             if not self.ks.bootloader.source and part.mountpoint == "/boot":

@@ -175,8 +171,7 @@ class DirectPlugin(ImagerPlugin):
                 if self.ptable_format == 'gpt':
                     part.uuid = str(uuid.uuid4())
                 else: # msdos partition table
-                    part.uuid = '%0x-%02d' % (self._image.identifier,
-                                              self._get_part_num(num, self.parts))
+                    part.uuid = '%0x-%02d' % (self._image.identifier, part.realnum)
 
         fstab_path = self._write_fstab(self.rootfs_dir.get("ROOTFS_DIR"))
 

@@ -277,14 +272,13 @@ class DirectPlugin(ImagerPlugin):
 
         Assume partition order same as in wks
         """
-        for num, part in enumerate(self.parts, 1):
+        for part in self.parts:
             if part.mountpoint == "/":
                 if part.uuid:
                     return "PARTUUID=%s" % part.uuid
                 else:
                     suffix = 'p' if part.disk.startswith('mmcblk') else ''
-                    pnum = self._get_part_num(num, self.parts)
-                    return "/dev/%s%s%-d" % (part.disk, suffix, pnum)
+                    return "/dev/%s%s%-d" % (part.disk, suffix, part.realnum)
 
     def cleanup(self):
         if self._image:

index ffe6c84..18cd277 100644
@@ -158,7 +158,7 @@ class DirectPlugin(ImagerPlugin):
         """
         image_path = self._full_path(self.workdir, self.parts[0].disk, "direct")
         self._image = PartitionedImage(image_path, self.ptable_format,
-                                       self.native_sysroot)
+                                       self.parts, self.native_sysroot)
 
         for part in self.parts:
             # as a convenience, set source to the boot partition source

@@ -166,13 +166,6 @@ class DirectPlugin(ImagerPlugin):
             if not self.ks.bootloader.source and part.mountpoint == "/boot":
                 self.ks.bootloader.source = part.source
 
-            # generate parition UUIDs
-            if not part.uuid and part.use_uuid:
-                if self.ptable_format == 'gpt':
-                    part.uuid = str(uuid.uuid4())
-                else: # msdos partition table
-                    part.uuid = '%0x-%02d' % (self._image.identifier, part.realnum)
-
         fstab_path = self._write_fstab(self.rootfs_dir.get("ROOTFS_DIR"))
 
         for part in self.parts:

@@ -310,7 +303,7 @@ class PartitionedImage():
     Partitioned image in a file.
     """
 
-    def __init__(self, path, ptable_format, native_sysroot=None):
+    def __init__(self, path, ptable_format, partitions, native_sysroot=None):
         self.path = path  # Path to the image file
         self.numpart = 0  # Number of allocated partitions
         self.realpart = 0 # Number of partitions in the partition table

@@ -321,12 +314,20 @@ class PartitionedImage():
         # Disk system identifier
         self.identifier = int.from_bytes(os.urandom(4), 'little')
 
-        self.partitions = []
+        self.partitions = partitions
         self.partimages = []
         # Size of a sector used in calculations
         self.sector_size = SECTOR_SIZE
         self.native_sysroot = native_sysroot
 
+        # generate parition UUIDs
+        for part in self.partitions:
+            if not part.uuid and part.use_uuid:
+                if self.ptable_format == 'gpt':
+                    part.uuid = str(uuid.uuid4())
+                else: # msdos partition table
+                    part.uuid = '%0x-%02d' % (self.identifier, part.realnum)
+
     def add_partition(self, part):
         """
         Add the next partition. Partitions have to be added in the

@@ -337,8 +338,6 @@ class PartitionedImage():
         # Converting kB to sectors for parted
         part.size_sec = part.disk_size * 1024 // self.sector_size
 
-        self.partitions.append(part)
-
     def layout_partitions(self):
         """ Layout the partitions, meaning calculate the position of every
         partition on the disk. The 'ptable_format' parameter defines the

index 18cd277..64a3368 100644
@@ -87,6 +87,13 @@ class DirectPlugin(ImagerPlugin):
                     continue
                 part.realnum = realnum
 
+        # as a convenience, set source to the boot partition source
+        # instead of forcing it to be set via bootloader --source
+        for part in self.parts:
+            if not self.ks.bootloader.source and part.mountpoint == "/boot":
+                self.ks.bootloader.source = part.source
+                break
+
     def do_create(self):
         """
         Plugin entry point.

@@ -160,12 +167,6 @@ class DirectPlugin(ImagerPlugin):
         self._image = PartitionedImage(image_path, self.ptable_format,
                                        self.parts, self.native_sysroot)
 
-        for part in self.parts:
-            # as a convenience, set source to the boot partition source
-            # instead of forcing it to be set via bootloader --source
-            if not self.ks.bootloader.source and part.mountpoint == "/boot":
-                self.ks.bootloader.source = part.source
-
         fstab_path = self._write_fstab(self.rootfs_dir.get("ROOTFS_DIR"))
 
         for part in self.parts:

index 64a3368..9c4109e 100644
@@ -182,12 +182,8 @@ class DirectPlugin(ImagerPlugin):
                     rsize_bb = get_bitbake_var('ROOTFS_SIZE', image_name)
                     if rsize_bb:
                         part.size = int(round(float(rsize_bb)))
-            # need to create the filesystems in order to get their
-            # sizes before we can add them and do the layout.
-            part.prepare(self, self.workdir, self.oe_builddir, self.rootfs_dir,
-                         self.bootimg_dir, self.kernel_dir, self.native_sysroot)
 
-            self._image.add_partition(part)
+        self._image.prepare(self)
 
         if fstab_path:
             shutil.move(fstab_path + ".orig", fstab_path)

@@ -329,15 +325,17 @@ class PartitionedImage():
                 else: # msdos partition table
                     part.uuid = '%0x-%02d' % (self.identifier, part.realnum)
 
-    def add_partition(self, part):
-        """
-        Add the next partition. Partitions have to be added in the
-        first-to-last order.
-        """
-        part.ks_pnum = len(self.partitions)
+    def prepare(self, imager):
+        """Prepare an image. Call prepare method of all image partitions."""
+        for part in self.partitions:
+            # need to create the filesystems in order to get their
+            # sizes before we can add them and do the layout.
+            part.prepare(imager, imager.workdir, imager.oe_builddir,
+                         imager.rootfs_dir, imager.bootimg_dir,
+                         imager.kernel_dir, imager.native_sysroot)
 
-        # Converting kB to sectors for parted
-        part.size_sec = part.disk_size * 1024 // self.sector_size
+            # Converting kB to sectors for parted
+            part.size_sec = part.disk_size * 1024 // self.sector_size
 
     def layout_partitions(self):
         """ Layout the partitions, meaning calculate the position of every

index 9c4109e..c517904 100644
@@ -74,19 +74,6 @@ class DirectPlugin(ImagerPlugin):
         self.ptable_format = self.ks.bootloader.ptable
         self.parts = self.ks.partitions
 
-        # calculate the real partition number, accounting for partitions not
-        # in the partition table and logical partitions
-        realnum = 0
-        for part in self.parts:
-            if part.no_table:
-                part.realnum = 0
-            else:
-                realnum += 1
-                if self.ptable_format == 'msdos' and realnum > 3:
-                    part.realnum = realnum + 1
-                    continue
-                part.realnum = realnum
-
         # as a convenience, set source to the boot partition source
         # instead of forcing it to be set via bootloader --source
         for part in self.parts:

@@ -317,6 +304,19 @@ class PartitionedImage():
         self.sector_size = SECTOR_SIZE
         self.native_sysroot = native_sysroot
 
+        # calculate the real partition number, accounting for partitions not
+        # in the partition table and logical partitions
+        realnum = 0
+        for part in self.partitions:
+            if part.no_table:
+                part.realnum = 0
+            else:
+                realnum += 1
+                if self.ptable_format == 'msdos' and realnum > 3:
+                    part.realnum = realnum + 1
+                    continue
+                part.realnum = realnum
+
         # generate parition UUIDs
         for part in self.partitions:
             if not part.uuid and part.use_uuid:

index c517904..481d24d 100644
@@ -81,6 +81,10 @@ class DirectPlugin(ImagerPlugin):
                 self.ks.bootloader.source = part.source
                 break
 
+        image_path = self._full_path(self.workdir, self.parts[0].disk, "direct")
+        self._image = PartitionedImage(image_path, self.ptable_format,
+                                       self.parts, self.native_sysroot)
+
     def do_create(self):
         """
         Plugin entry point.

@@ -150,10 +154,6 @@ class DirectPlugin(ImagerPlugin):
         filesystems from the artifacts directly and combine them into
         a partitioned image.
         """
-        image_path = self._full_path(self.workdir, self.parts[0].disk, "direct")
-        self._image = PartitionedImage(image_path, self.ptable_format,
-                                       self.parts, self.native_sysroot)
-
         fstab_path = self._write_fstab(self.rootfs_dir.get("ROOTFS_DIR"))
 
         for part in self.parts: