From 5878db3e8138a735de724fd877ad5edcbc9fb14b Mon Sep 17 00:00:00 2001
From: Martin Gracik <mgracik@redhat.com>
Date: Tue, 15 Mar 2011 11:37:44 +0100
Subject: [PATCH] Use factory to get the image classes

---
 src/pylorax/__init__.py |  8 ++------
 src/pylorax/images.py   | 18 ++++++++++++++++++
 2 files changed, 20 insertions(+), 6 deletions(-)

diff --git a/src/pylorax/__init__.py b/src/pylorax/__init__.py
index 97979cba..d1b58903 100644
--- a/src/pylorax/__init__.py
+++ b/src/pylorax/__init__.py
@@ -367,12 +367,8 @@ class Lorax(BaseLoraxClass):
                             self.variant, self.basearch)
 
         # get the image class
-        if self.basearch == "ppc":
-            imgclass = images.PPC
-        elif self.basearch in ("i386", "x86_64"):
-            imgclass = images.X86
-        else:
-            raise Exception("not supported arch '{0}'".format(self.basearch))
+        factory = images.Factory()
+        imgclass = factory.get_class(self.basearch)
 
         i = imgclass(kernellist=self.outputtree.kernels,
                      installtree=self.installtree,
diff --git a/src/pylorax/images.py b/src/pylorax/images.py
index f47fd610..954290be 100644
--- a/src/pylorax/images.py
+++ b/src/pylorax/images.py
@@ -667,3 +667,21 @@ class S390(object):
 
     def create_boot(self, efiboot=None):
         pass
+
+
+class Factory(object):
+
+    DISPATCH_MAP = {"ppc": PPC,
+                    "i386": X86,
+                    "x86_64": X86,
+                    "s390": S390,
+                    "s390x": S390}
+
+    def __init__(self):
+        pass
+
+    def get_class(self, arch):
+        if arch in self.DISPATCH_MAP:
+            return self.DISPATCH_MAP[arch]
+        else:
+            raise Exception("no support for {0}".format(arch))