From 4253cab7733a234e3e15b4f8a2c7ea3ebf25b977 Mon Sep 17 00:00:00 2001
From: Ian Wienand <iwienand@redhat.com>
Date: Wed, 31 May 2017 11:24:55 +1000
Subject: [PATCH] Make BlockDeviceState implement a dict

While plugins treat the state as just a dictionary, it's nice for the
driver functions to keep state related functions encapsulated in the
state object singleton.  Wrap the internal state dictionary so we can
pass the BlockDeviceState directly without dereferencing.

Change-Id: Ic0193c64d645ed1312f898cbfef87841f460799c
---
 diskimage_builder/block_device/blockdevice.py | 31 ++++++++++++-------
 1 file changed, 20 insertions(+), 11 deletions(-)

diff --git a/diskimage_builder/block_device/blockdevice.py b/diskimage_builder/block_device/blockdevice.py
index c112d671b..ed2c9587e 100644
--- a/diskimage_builder/block_device/blockdevice.py
+++ b/diskimage_builder/block_device/blockdevice.py
@@ -13,6 +13,7 @@
 # under the License.
 
 import codecs
+import collections
 import json
 import logging
 import os
@@ -39,12 +40,12 @@ def _load_json(file_name):
     return None
 
 
-class BlockDeviceState(object):
+class BlockDeviceState(collections.MutableMapping):
     """The global state singleton
 
     An reference to an instance of this object is passed between nodes
-    as a global repository.  It contains a single dictionary "state"
-    and a range of helper functions.
+    as a global repository.  It wraps a single dictionary "state"
+    and provides a few helper functions.
 
     This is used in two contexts:
 
@@ -56,8 +57,6 @@ class BlockDeviceState(object):
       from disk and are thus passed the full state.
     """
     # XXX:
-    #  - might it make sense to make state implement MutableMapping,
-    #    so that callers treat it as a dictionary?
     #  - we could implement getters/setters such that if loaded from
     #    disk, the state is read-only? or make it append-only
     #    (i.e. you can't overwrite existing keys)
@@ -78,6 +77,21 @@ class BlockDeviceState(object):
         else:
             self.state = {}
 
+    def __delitem__(self, key):
+        del self.state[key]
+
+    def __getitem__(self, key):
+        return self.state[key]
+
+    def __setitem__(self, key, value):
+        self.state[key] = value
+
+    def __iter__(self):
+        return iter(self.state)
+
+    def __len__(self):
+        return len(self.state)
+
     def save_state(self, filename):
         """Persist the state to disk
 
@@ -289,7 +303,6 @@ class BlockDevice(object):
         # dictionary.  They can only be accessed after the state has
         # been dumped; i.e. after cmd_create() called.
         state = BlockDeviceState(self.state_json_file_name)
-        state = state.state
 
         if symbol == 'image-block-partition':
             # If there is no partition needed, pass back directly the
@@ -313,7 +326,6 @@ class BlockDevice(object):
         # State should have been created by prior calls; we only need
         # the dict
         state = BlockDeviceState(self.state_json_file_name)
-        state = state.state
 
         tmp_fstab = os.path.join(self.state_dir, "fstab")
         with open(tmp_fstab, "wt") as fstab_fd:
@@ -360,7 +372,7 @@ class BlockDevice(object):
         try:
             dg, call_order = create_graph(self.config, self.params)
             for node in call_order:
-                node.create(state.state, rollback)
+                node.create(state, rollback)
         except Exception:
             logger.exception("Create failed; rollback initiated")
             for rollback_cb in reversed(rollback):
@@ -380,7 +392,6 @@ class BlockDevice(object):
         # (? more details?)
         try:
             state = BlockDeviceState(self.state_json_file_name)
-            state = state.state
         except BlockDeviceSetupException:
             logger.info("State already cleaned - no way to do anything here")
             return 0
@@ -401,7 +412,6 @@ class BlockDevice(object):
         # State should have been created by prior calls; we only need
         # the dict
         state = BlockDeviceState(self.state_json_file_name)
-        state = state.state
 
         # Deleting must be done in reverse order
         dg, call_order = create_graph(self.config, self.params)
@@ -420,7 +430,6 @@ class BlockDevice(object):
         # State should have been created by prior calls; we only need
         # the dict
         state = BlockDeviceState(self.state_json_file_name)
-        state = state.state
 
         # Deleting must be done in reverse order
         dg, call_order = create_graph(self.config, self.params)