Create generator map in gen/__init__.py.

3 files changed, 27 insertions, 12 deletions

diff --git a/gen/__init__.py b/gen/__init__.py
index e69de29..8c60b46 100644
--- a/gen/__init__.py
+++ b/gen/__init__.py
@@ -0,0 +1,6 @@
+import generator as gen
+import edf_generators as edf
+
+gen.register_generator("G-EDF", edf.GedfGenerator)
+gen.register_generator("P-EDF", edf.PedfGenerator)
+gen.register_generator("C-EDF", edf.CedfGenerator)
diff --git a/gen/edf_generators.py b/gen/edf_generators.py
index 78b0ba1..3f05b76 100644
--- a/gen/edf_generators.py
+++ b/gen/edf_generators.py
@@ -9,8 +9,8 @@ TP_PART_TASK = TP_TBASE.format("-p $t.cpu")
 
 class EdfGenerator(gen.Generator):
     '''Creates sporadic task sets with the most common Litmus options.'''
-    def __init__(self, name, templates, options, params):
-        super(EdfGenerator, self).__init__(name, templates,
+    def __init__(self, scheduler, templates, options, params):
+        super(EdfGenerator, self).__init__(scheduler, templates,
                                            self.__make_options() + options,
                                            params)
 
@@ -45,8 +45,8 @@ class EdfGenerator(gen.Generator):
 
 
 class PartitionedGenerator(EdfGenerator):
-    def __init__(self, name, templates, options, params):
-        super(PartitionedGenerator, self).__init__(name,
+    def __init__(self, scheduler, templates, options, params):
+        super(PartitionedGenerator, self).__init__(scheduler,
             templates + [TP_PART_TASK], options, params)
 
     def _customize(self, taskset, exp_params):
diff --git a/gen/generator.py b/gen/generator.py
index 56563b1..f35a22b 100644
--- a/gen/generator.py
+++ b/gen/generator.py
@@ -51,14 +51,14 @@ class Generator(object):
     This class also performs checks of parameter values and prints out help.
     All subclasses must implement _create_exp.
     '''
-    def __init__(self, name, templates, options, params):
+    def __init__(self, scheduler, templates, options, params):
         self.options = self.__make_options(params) + options
 
         self.__setup_params(params)
 
-        self.params   = params
-        self.template = "\n".join([TP_RM] + templates)
-        self.name     = name
+        self.params    = params
+        self.template  = "\n".join([TP_RM] + templates)
+        self.scheduler = scheduler
 
     def __make_options(self, params):
         '''Return generic Litmus options.'''
@@ -137,7 +137,7 @@ class Generator(object):
 
         exp_params_file = self.out_dir + "/" + DEFAULTS['params_file']
         with open(exp_params_file, 'wa') as f:
-            params['scheduler'] = self.name
+            params['scheduler'] = self.scheduler
             pprint.pprint(params, f)
 
         if tasks:
@@ -210,7 +210,7 @@ class Generator(object):
         for dp in DesignPointGenerator(self.params):
             for trial in xrange(trials):
                 # Create directory name from relevant parameters
-                dir_leaf  = "sched=%s_%s" % (self.name, col_map.encode(dp))
+                dir_leaf  = "sched=%s_%s" % (self.scheduler, col_map.encode(dp))
                 dir_leaf  = dir_leaf.strip('_') # If there are none
                 dir_leaf += ("_trial=%s" % trial) if trials > 1 else ""
 
@@ -239,12 +239,12 @@ class Generator(object):
 
     def print_help(self):
         display_options = [o for o in self.options if not o.hidden]
-        s = str(Template("""Generator $name:
+        s = str(Template("""scheduler $scheduler:
         #for $o in $options
         $o.name -- $o.help
         \tDefault: $o.default
         \tAllowed: $o.types
-        #end for""", searchList={'name':self.name, 'options':display_options}))
+        #end for""", searchList={'scheduler':self.scheduler, 'options':display_options}))
 
         # Has to be an easier way to print this out...
         for line in s.split("\n"):
@@ -258,3 +258,12 @@ class Generator(object):
                     res = [" "*Generator.HELP_INDENT +res[-1]]
                     i = Generator.HELP_INDENT + len(word)
             print(", ".join(res))
+
+
+generators = {}
+
+def register_generator(name, clazz):
+    generators[name] = clazz
+
+def get_generators():
+    return generators