4 files changed, 40 insertions, 27 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
diff --git a/gen_exps.py b/gen_exps.py
index 8688ae9..6488cdc 100755
--- a/gen_exps.py
+++ b/gen_exps.py
@@ -1,18 +1,14 @@
 #!/usr/bin/env python
 from __future__ import print_function
 
+import gen.generator as gen
 import os
 import re
 import shutil as sh
+import sys
 
-from gen.edf_generators import GedfGenerator,PedfGenerator,CedfGenerator
 from optparse import OptionParser
 
-# There has to be a better way to do this...
-GENERATORS = {'C-EDF':CedfGenerator,
-              'P-EDF':PedfGenerator,
-              'G-EDF':GedfGenerator}
-
 def parse_args():
     parser = OptionParser("usage: %prog [options] [files...] "
                           "[generators...] [param=val[,val]...]")
@@ -51,13 +47,14 @@ def main():
 
     # Print generator information on the command line
     if opts.list_gens:
-        print(", ".join(GENERATORS.keys()))
+        print(", ".join(gen.get_generators()))
     if opts.described != None:
         for generator in opts.described.split(','):
-            if generator not in GENERATORS:
+            if generator not in gen.get_generators():
                 print("No generator '%s'" % generator)
             else:
-                GENERATORS[generator]().print_help()
+                sys.stdout.write("Generator '%s', " % generator)
+                gen.get_generators()[generator]().print_help()
     if opts.list_gens or opts.described:
         return 0
 
@@ -65,10 +62,9 @@ def main():
 
     # Ensure some generator is loaded
     args = list(set(args) - set(params))
-    #TODO: get every loaded plugin, try and use that generator
-    args = args or ['C-EDF', 'G-EDF', 'P-EDF']
+    args = args or gen.get_generators().keys()
 
-    # Split into files to load, named generators
+    # Split into files to load and named generators
     files = filter(os.path.exists, args)
     gen_list = list(set(args) - set(files))
 
@@ -86,13 +82,15 @@ def main():
         os.mkdir(opts.out_dir)
 
     for gen_name, gen_params in exp_sets:
-        if gen_name not in GENERATORS:
-            raise ValueError("Invalid generator name: %s" % gen_name)
+        if gen_name not in gen.get_generators():
+            raise ValueError("Invalid generator '%s'" % gen_name)
 
         print("Creating experiments using %s generator..." % gen_name)
 
         params = dict(gen_params.items() + global_params.items())
-        generator = GENERATORS[gen_name](params=params)
+        clazz  = gen.get_generators()[gen_name]
+
+        generator = clazz(params=params)
 
         generator.create_exps(opts.out_dir, opts.force, opts.trials)