From bc58b5c515e0338ee1fb221bf323584b4dcadd5b Mon Sep 17 00:00:00 2001
From: Ian Wienand <iwienand@redhat.com>
Date: Thu, 18 May 2017 14:50:02 +1000
Subject: [PATCH] Move parts of Partition creation into object

Move Partition() object creation into the actual Partition object,
rather than having the logic within the Partitioning() object

Change-Id: I833ed419a0fca38181a9e2db28e5af87500d8ba4
---
 .../block_device/level1/partition.py          | 45 ++++++++++++++-----
 .../block_device/level1/partitioning.py       | 35 ++-------------
 2 files changed, 38 insertions(+), 42 deletions(-)

diff --git a/diskimage_builder/block_device/level1/partition.py b/diskimage_builder/block_device/level1/partition.py
index 3f314a91d..a6954b717 100644
--- a/diskimage_builder/block_device/level1/partition.py
+++ b/diskimage_builder/block_device/level1/partition.py
@@ -12,6 +12,8 @@
 
 import logging
 
+from diskimage_builder.block_device.exception import \
+    BlockDeviceSetupException
 from diskimage_builder.block_device.tree_config import TreeConfig
 from diskimage_builder.graph.digraph import Digraph
 
@@ -30,7 +32,10 @@ class PartitionTreeConfig(object):
 
         for partition in config_value:
             name = partition['name']
-            nconfig = {'name': name}
+            nconfig = {
+                'name': name,
+                'base': base_name
+            }
             for k, v in partition.items():
                 if k not in plugin_manager:
                     nconfig[k] = v
@@ -48,17 +53,37 @@ class Partition(Digraph.Node):
     type_string = "partitions"
     tree_config = TreeConfig("partitions")
 
-    def __init__(self, name, flags, size, ptype, base, partitioning,
-                 prev_partition):
-        Digraph.Node.__init__(self, name)
-        self.name = name
-        self.flags = flags
-        self.size = size
-        self.ptype = ptype
-        self.base = base
-        self.partitioning = partitioning
+    flag_boot = 1
+    flag_primary = 2
+
+    def __init__(self, config, parent, prev_partition):
+        if 'name' not in config:
+            raise BlockDeviceSetupException(
+                "Missing 'name' in partition config: %s" % config)
+        self.name = config['name']
+
+        Digraph.Node.__init__(self, self.name)
+
+        self.base = config['base']
+        self.partitioning = parent
         self.prev_partition = prev_partition
 
+        self.flags = set()
+        if 'flags' in config:
+            for f in config['flags']:
+                if f == 'boot':
+                    self.flags.add(self.flag_boot)
+                elif f == 'primary':
+                    self.flags.add(self.flag_primary)
+                else:
+                    raise BlockDeviceSetupException("Unknown flag: %s" % f)
+
+        if 'size' not in config:
+            raise BlockDeviceSetupException("No size in partition" % self.name)
+        self.size = config['size']
+
+        self.ptype = int(config['type'], 16) if 'type' in config else 0x83
+
     def __repr__(self):
         return "<Partition [%s] on [%s] size [%s] prev [%s]>" \
             % (self.name, self.base, self.size,
diff --git a/diskimage_builder/block_device/level1/partitioning.py b/diskimage_builder/block_device/level1/partitioning.py
index 27d27f0b8..6a8c8880d 100644
--- a/diskimage_builder/block_device/level1/partitioning.py
+++ b/diskimage_builder/block_device/level1/partitioning.py
@@ -61,9 +61,6 @@ class Partitioning(Digraph.Node):
 
     tree_config = PartitioningTreeConfig()
 
-    flag_boot = 1
-    flag_primary = 2
-
     def __init__(self, config, default_config):
         logger.debug("Creating Partitioning object; config [%s]" % config)
         # Because using multiple partitions of one base is done
@@ -106,35 +103,9 @@ class Partitioning(Digraph.Node):
         prev_partition = None
 
         for part_cfg in config['partitions']:
-            if 'name' not in part_cfg:
-                raise BlockDeviceSetupException(
-                    "Missing 'name' in partition config")
-            part_name = part_cfg['name']
-
-            flags = set()
-            if 'flags' in part_cfg:
-                for f in part_cfg['flags']:
-                    if f == 'boot':
-                        flags.add(Partitioning.flag_boot)
-                    elif f == 'primary':
-                        flags.add(Partitioning.flag_primary)
-                    else:
-                        raise BlockDeviceSetupException(
-                            "Unknown flag [%s] in partitioning for [%s]"
-                            % (f, part_name))
-
-            if 'size' not in part_cfg:
-                raise BlockDeviceSetupException("No 'size' in partition [%s]"
-                                                % part_name)
-            size = part_cfg['size']
-
-            ptype = int(part_cfg['type'], 16) if 'type' in part_cfg else 0x83
-
-            np = Partition(part_name, flags, size, ptype, self.base, self,
-                           prev_partition)
+            np = Partition(part_cfg, self, prev_partition)
             self.partitions.append(np)
             prev_partition = np
-            logger.debug(part_cfg)
 
     def _size_of_block_dev(self, dev):
         with open(dev, "r") as fd:
@@ -201,9 +172,9 @@ class Partitioning(Digraph.Node):
         with MBR(image_path, disk_size, self.align) as part_impl:
             for part_cfg in self.partitions:
                 part_name = part_cfg.get_name()
-                part_bootflag = Partitioning.flag_boot \
+                part_bootflag = Partition.flag_boot \
                                 in part_cfg.get_flags()
-                part_primary = Partitioning.flag_primary \
+                part_primary = Partition.flag_primary \
                                in part_cfg.get_flags()
                 part_size = part_cfg.get_size()
                 part_free = part_impl.free()