X-Git-Url: https://code.grnet.gr/git/snf-image-creator/blobdiff_plain/5886f568337cad55facff1132344374553bd8b1f..dcf9274b9589ba4380412e49c8cbc2dccabd4b83:/image_creator/os_type/__init__.py

diff --git a/image_creator/os_type/__init__.py b/image_creator/os_type/__init__.py
index 8a39846..c0fa505 100644
--- a/image_creator/os_type/__init__.py
+++ b/image_creator/os_type/__init__.py
@@ -31,22 +31,22 @@
 # interpreted as representing official policies, either expressed
 # or implied, of GRNET S.A.
 
-from image_creator.util import output, FatalError
+from image_creator.util import FatalError
 
 import textwrap
 import re
 
 
-def get_os_class(distro, osfamily):
+def os_cls(distro, osfamily):
     module = None
     classname = None
     try:
-        module = __import__("image_creator.os_type.%s"
-            % distro, fromlist=['image_creator.os_type'])
+        module = __import__("image_creator.os_type.%s" % distro,
+                            fromlist=['image_creator.os_type'])
         classname = distro.capitalize()
     except ImportError:
-        module = __import__("image_creator.os_type.%s"
-            % osfamily, fromlist=['image_creator.os_type'])
+        module = __import__("image_creator.os_type.%s" % osfamily,
+                            fromlist=['image_creator.os_type'])
         classname = osfamily.capitalize()
 
     return getattr(module, classname)
@@ -70,9 +70,10 @@ def sysprep(enabled=True):
 class OSBase(object):
     """Basic operating system class"""
 
-    def __init__(self, rootdev, ghandler):
+    def __init__(self, rootdev, ghandler, output):
         self.root = rootdev
         self.g = ghandler
+        self.out = output
 
         # Collect metadata about the OS
         self.meta = {}
@@ -86,18 +87,20 @@ class OSBase(object):
 
     def list_syspreps(self):
 
-        objs = [getattr(self, name) for name in dir(self) \
-            if not name.startswith('_')]
+        objs = [getattr(self, name) for name in dir(self)
+                if not name.startswith('_')]
 
-        enabled = [x for x in objs if self._is_sysprep(x) and x.enabled]
-        disabled = [x for x in objs if self._is_sysprep(x) and not x.enabled]
+        return [x for x in objs if self._is_sysprep(x)]
 
-        return enabled, disabled
+    def sysprep_info(self, obj):
+        assert self._is_sysprep(obj), "Object is not a sysprep"
 
-    def _sysprep_change_status(self, name, status):
+        return (obj.__name__.replace('_', '-'), textwrap.dedent(obj.__doc__))
 
+    def get_sysprep_by_name(self, name):
+        """Returns the sysprep object with the given name"""
         error_msg = "Syprep operation %s does not exist for %s" % \
-                (name, self.__class__.__name__)
+                    (name, self.__class__.__name__)
 
         method_name = name.replace('-', '_')
         method = None
@@ -109,43 +112,45 @@ class OSBase(object):
         if not self._is_sysprep(method):
             raise FatalError(error_msg)
 
-        setattr(method.im_func, 'enabled', status)
+        return method
 
-    def enable_sysprep(self, name):
-        """Enable a system preperation operation"""
-        self._sysprep_change_status(name, True)
+    def enable_sysprep(self, obj):
+        """Enable a system preparation operation"""
+        setattr(obj.im_func, 'enabled', True)
 
-    def disable_sysprep(self, name):
-        """Disable a system preperation operation"""
-        self._sysprep_change_status(name, False)
+    def disable_sysprep(self, obj):
+        """Disable a system preparation operation"""
+        setattr(obj.im_func, 'enabled', False)
 
     def print_syspreps(self):
-        """Print enabled and disabled system preperation operations."""
+        """Print enabled and disabled system preparation operations."""
 
-        enabled, disabled = self.list_syspreps()
+        syspreps = self.list_syspreps()
+        enabled = filter(lambda x: x.enabled, syspreps)
+        disabled = filter(lambda x: not x.enabled, syspreps)
 
         wrapper = textwrap.TextWrapper()
         wrapper.subsequent_indent = '\t'
         wrapper.initial_indent = '\t'
         wrapper.width = 72
 
-        output("Enabled system preperation operations:")
+        self.out.output("Enabled system preparation operations:")
         if len(enabled) == 0:
-            output("(none)")
+            self.out.output("(none)")
         else:
             for sysprep in enabled:
                 name = sysprep.__name__.replace('_', '-')
                 descr = wrapper.fill(textwrap.dedent(sysprep.__doc__))
-                output('    %s:\n%s\n' % (name, descr))
+                self.out.output('    %s:\n%s\n' % (name, descr))
 
-        output("Disabled system preperation operations:")
+        self.out.output("Disabled system preparation operations:")
         if len(disabled) == 0:
-            output("(none)")
+            self.out.output("(none)")
         else:
             for sysprep in disabled:
                 name = sysprep.__name__.replace('_', '-')
                 descr = wrapper.fill(textwrap.dedent(sysprep.__doc__))
-                output('    %s:\n%s\n' % (name, descr))
+                self.out.output('    %s:\n%s\n' % (name, descr))
 
     @add_prefix
     def ls(self, directory):
@@ -202,15 +207,17 @@ class OSBase(object):
     def do_sysprep(self):
         """Prepere system for image creation."""
 
-        output('Preparing system for image creation:')
+        self.out.output('Preparing system for image creation:')
 
-        tasks, _ = self.list_syspreps()
-        size = len(tasks)
+        tasks = self.list_syspreps()
+        enabled = filter(lambda x: x.enabled, tasks)
+
+        size = len(enabled)
         cnt = 0
-        for task in tasks:
+        for task in enabled:
             cnt += 1
-            output(('(%d/%d)' % (cnt, size)).ljust(7), False)
+            self.out.output(('(%d/%d)' % (cnt, size)).ljust(7), False)
             task()
-        output()
+        self.out.output()
 
 # vim: set sta sts=4 shiftwidth=4 sw=4 et ai :