Add {enable, disable}_guestfs methods in image cls

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

diff --git a/image_creator/disk.py b/image_creator/disk.py
index edcc4db..e4b6bc8 100644 (file)
--- a/image_creator/disk.py
+++ b/image_creator/disk.py
@@ -1,3 +1,5 @@
+# -*- coding: utf-8 -*-
+#
 # Copyright 2012 GRNET S.A. All rights reserved.
 #
 # Redistribution and use in source and binary forms, with or
@@ -31,6 +33,8 @@
 # interpreted as representing official policies, either expressed
 # or implied, of GRNET S.A.
 
+"""Module hosting the Disk class."""
+
 from image_creator.util import get_command
 from image_creator.util import try_fail_repeat
 from image_creator.util import free_space
@@ -50,7 +54,26 @@ losetup = get_command('losetup')
 blockdev = get_command('blockdev')
 
 
-TMP_CANDIDATES = ['/var/tmp', os.path.expanduser('~'), '/mnt']
+def get_tmp_dir(default=None):
+    """Check tmp directory candidates and return the one with the most
+    available space.
+    """
+    if default is not None:
+        return default
+
+    TMP_CANDIDATES = ['/var/tmp', os.path.expanduser('~'), '/mnt']
+
+    space = map(free_space, TMP_CANDIDATES)
+
+    max_idx = 0
+    max_val = space[0]
+    for i, val in zip(range(len(space)), space):
+        if val > max_val:
+            max_val = val
+            max_idx = i
+
+    # Return the candidate path with more available space
+    return TMP_CANDIDATES[max_idx]
 
 
 class Disk(object):
@@ -71,29 +94,10 @@ class Disk(object):
         self.out = output
         self.meta = {}
         self.tmp = tempfile.mkdtemp(prefix='.snf_image_creator.',
-                                    dir=self._get_tmp_dir(tmp))
+                                    dir=get_tmp_dir(tmp))
 
         self._add_cleanup(shutil.rmtree, self.tmp)
 
-    def _get_tmp_dir(self, default=None):
-        """Check tmp directory candidates and return the one with the most
-        available space.
-        """
-        if default is not None:
-            return default
-
-        space = map(free_space, TMP_CANDIDATES)
-
-        max_idx = 0
-        max_val = space[0]
-        for i, val in zip(range(len(space)), space):
-            if val > max_val:
-                max_val = val
-                max_idx = i
-
-        # Return the candidate path with more available space
-        return TMP_CANDIDATES[max_idx]
-
     def _add_cleanup(self, job, *args):
         """Add a new job in the cleanup list"""
         self._cleanup_jobs.append((job, args))
@@ -158,7 +162,7 @@ class Disk(object):
             self.out.success('looks like a block device')
 
         # Take a snapshot and return it to the user
-        self.out.output("Snapshotting media source...", False)
+        self.out.output("Snapshotting media source ...", False)
         size = blockdev('--getsz', sourcedev)
         cowfd, cow = tempfile.mkstemp(dir=self.tmp)
         os.close(cowfd)
@@ -170,27 +174,29 @@ class Disk(object):
         snapshot = uuid.uuid4().hex
         tablefd, table = tempfile.mkstemp()
         try:
-            os.write(tablefd, "0 %d snapshot %s %s n 8" %
-                              (int(size), sourcedev, cowdev))
+            try:
+                os.write(tablefd, "0 %d snapshot %s %s n 8" %
+                                  (int(size), sourcedev, cowdev))
+            finally:
+                os.close(tablefd)
+
             dmsetup('create', snapshot, table)
             self._add_cleanup(try_fail_repeat, dmsetup, 'remove', snapshot)
-
         finally:
             os.unlink(table)
         self.out.success('done')
         return "/dev/mapper/%s" % snapshot
 
-    def get_image(self, media):
-        """Returns a newly created ImageCreator instance."""
+    def get_image(self, media, **kargs):
+        """Returns a newly created Image instance."""
 
-        image = Image(media, self.out)
+        image = Image(media, self.out, **kargs)
         self._images.append(image)
         image.enable()
         return image
 
     def destroy_image(self, image):
-        """Destroys an ImageCreator instance previously created by
-        get_image_creator method.
+        """Destroys an Image instance previously created by get_image method.
         """
         self._images.remove(image)
         image.destroy()