Implement a WizardExit exception in dialog_wizard

[snf-image-creator] / image_creator / dialog_wizard.py

diff --git a/image_creator/dialog_wizard.py b/image_creator/dialog_wizard.py
index 8e122a8..636292e 100644 (file)
--- a/image_creator/dialog_wizard.py
+++ b/image_creator/dialog_wizard.py
@@ -44,6 +44,10 @@ from image_creator.output.cli import OutputWthProgress
 PAGE_WIDTH = 70
 
 
+class WizardExit(Exception):
+    pass
+
+
 class Wizard:
     def __init__(self, session):
         self.session = session
@@ -56,7 +60,10 @@ class Wizard:
     def run(self):
         idx = 0
         while True:
-            idx += self.pages[idx].run(self.session, idx, len(self.pages))
+            try:
+                idx += self.pages[idx].run(self.session, idx, len(self.pages))
+            except WizardExit:
+                return False
 
             if idx >= len(self.pages):
                 break
@@ -69,7 +76,6 @@ class Wizard:
 class WizardPage:
     NEXT = 1
     PREV = -1
-    EXIT = -255
 
     def run(self, session, index, total):
         raise NotImplementedError
@@ -124,7 +130,7 @@ class WizardYesNoPage(WizardPage):
             if ret == d.DIALOG_CANCEL:
                 return self.PREV
             elif ret == d.DIALOG_EXTRA:
-                return self.EXIT
+                raise WizardExit
             elif ret == d.DIALOG_OK:
                 return self.NEXT
 
@@ -177,7 +183,7 @@ def extract_image(session):
     device.out = out
     image_os.out = out
 
-    out.output()
+    out.clear()
 
     #Sysprep
     device.mount(False)
@@ -187,7 +193,9 @@ def extract_image(session):
 
     #Shrink
     size = device.shrink()
+
     metadata.update(device.meta)
+    metadata['DESCRIPTION'] = wizard['ImageDescription']
 
     #MD5
     md5 = MD5(out)