diff --git a/lang_python.py b/lang_python.py
index 019b62d..18fca27 100644
--- a/lang_python.py
+++ b/lang_python.py
@@ -28,12 +28,6 @@
 """
 Python backend for LOXI
 
-This language specific file defines a dictionary 'targets' that
-defines the generated files and the functions used to generate them.
-
-For each generated file there is a generate_* function in py_gen.codegen
-and a Tenjin template under py_gen/templates.
-
 Target directory structure:
     pyloxi:
         loxi:
@@ -64,51 +58,9 @@
 """
 
 import os
-from loxi_globals import OFVersions
-import loxi_globals
-import loxi_utils.loxi_utils as loxi_utils
-import py_gen
-import py_gen.util
 import py_gen.codegen
-import template_utils
 
-versions = {
-    1: "of10",
-    2: "of11",
-    3: "of12",
-    4: "of13",
-}
-
-prefix = 'pyloxi/loxi'
-
-modules = {
-    1: ["action", "common", "const", "message", "util"],
-    2: ["action", "common", "const", "instruction", "message", "util"],
-    3: ["action", "common", "const", "instruction", "message", "oxm", "util"],
-    4: ["action", "action_id", "common", "const", "instruction", "instruction_id", "message", "meter_band", "oxm", "bsn_tlv", "util"],
-}
-
-def make_gen(name, version):
-    fn = getattr(py_gen.codegen, "generate_" + name)
-    return lambda out, name: fn(out, name, version)
-
-def static(template_name):
-    return lambda out, name: py_gen.util.render_template(out, template_name)
-
-targets = {
-    prefix+'/__init__.py': static('toplevel_init.py'),
-    prefix+'/pp.py': static('pp.py'),
-    prefix+'/generic_util.py': static('generic_util.py'),
-}
-
-for version, subdir in versions.items():
-    targets['%s/%s/__init__.py' % (prefix, subdir)] = make_gen('init', version)
-    for module in modules[version]:
-        filename = '%s/%s/%s.py' % (prefix, subdir, module)
-        targets[filename] = make_gen(module, OFVersions.from_wire(version))
+PREFIX = 'pyloxi/loxi'
 
 def generate(install_dir):
-    py_gen.codegen.init()
-    for (name, fn) in targets.items():
-        with template_utils.open_output(install_dir, name) as outfile:
-            fn(outfile, os.path.basename(name))
+    py_gen.codegen.codegen(os.path.join(install_dir, PREFIX))
diff --git a/py_gen/codegen.py b/py_gen/codegen.py
index 9c32e7f..3b0988e 100644
--- a/py_gen/codegen.py
+++ b/py_gen/codegen.py
@@ -26,16 +26,13 @@
 # under the EPL.
 
 from collections import defaultdict
+import os
 import loxi_globals
-import struct
 import template_utils
 import loxi_utils.loxi_utils as utils
 import util
-import oftype
 from loxi_ir import *
 
-modules_by_version = {}
-
 # Map from inheritance root to module name
 roots = {
     'of_header': 'message',
@@ -69,63 +66,36 @@
         modules[module_name].append(ofclass)
     return modules
 
-def generate_init(out, name, version):
-    util.render_template(out, 'init.py', version=version)
+def codegen(install_dir):
+    def render(name, template_name=None, **ctx):
+        if template_name is None:
+            template_name = os.path.basename(name)
+        with template_utils.open_output(install_dir, name) as out:
+            util.render_template(out, template_name, **ctx)
 
-def generate_action(out, name, version):
-    util.render_template(out, 'module.py',
-                         ofclasses=modules_by_version[version]['action'],
-                         version=version)
+    render('__init__.py', template_name='toplevel_init.py')
+    render('pp.py')
+    render('generic_util.py')
 
-def generate_action_id(out, name, version):
-    util.render_template(out, 'module.py',
-                         ofclasses=modules_by_version[version]['action_id'],
-                         version=version)
+    for version in loxi_globals.OFVersions.all_supported:
+        subdir = 'of' + version.version.replace('.', '')
+        modules = build_ofclasses(version)
 
-def generate_oxm(out, name, version):
-    util.render_template(out, 'module.py',
-                         ofclasses=modules_by_version[version]['oxm'],
-                         version=version)
+        render(os.path.join(subdir, '__init__.py'), template_name='init.py',
+               version=version, modules=modules.keys())
 
-def generate_common(out, name, version):
-    util.render_template(out, 'module.py',
-                         ofclasses=modules_by_version[version]['common'],
-                         version=version,
-                         extra_template='_common_extra.py')
+        render(os.path.join(subdir, 'util.py'), version=version)
 
-def generate_const(out, name, version):
-    util.render_template(out, 'const.py', version=version,
-                         enums=loxi_globals.ir[version].enums)
+        render(os.path.join(subdir, 'const.py'), version=version,
+               enums=loxi_globals.ir[version].enums)
 
-def generate_instruction(out, name, version):
-    util.render_template(out, 'module.py',
-                         ofclasses=modules_by_version[version]['instruction'],
-                         version=version)
+        args_by_module = {
+            'common': { 'extra_template' : '_common_extra.py' },
+            'message': { 'extra_template' : '_message_extra.py' },
+        }
 
-def generate_instruction_id(out, name, version):
-    util.render_template(out, 'module.py',
-                         ofclasses=modules_by_version[version]['instruction_id'],
-                         version=version)
-
-def generate_message(out, name, version):
-    util.render_template(out, 'module.py',
-                         ofclasses=modules_by_version[version]['message'],
-                         version=version,
-                         extra_template='_message_extra.py')
-
-def generate_meter_band(out, name, version):
-    util.render_template(out, 'module.py',
-                         ofclasses=modules_by_version[version]['meter_band'],
-                         version=version)
-
-def generate_util(out, name, version):
-    util.render_template(out, 'util.py', version=version)
-
-def generate_bsn_tlv(out, name, version):
-    util.render_template(out, 'module.py',
-                         ofclasses=modules_by_version[version]['bsn_tlv'],
-                         version=version)
-
-def init():
-    for version in loxi_globals.OFVersions.target_versions:
-        modules_by_version[version] = build_ofclasses(version)
+        for name, ofclasses in modules.items():
+            args = args_by_module.get(name, {})
+            render(os.path.join(subdir, name + '.py'), template_name='module.py',
+                   version=version, ofclasses=ofclasses, modules=modules.keys(),
+                   **args)
diff --git a/py_gen/templates/init.py b/py_gen/templates/init.py
index 3b73baa..aa968b6 100644
--- a/py_gen/templates/init.py
+++ b/py_gen/templates/init.py
@@ -29,17 +29,10 @@
 
 :: include('_autogen.py')
 
-import action, common, const, message
-:: if version >= 2:
-import instruction
-:: #endif
-:: if version >= 3:
-import oxm
-:: #endif
-:: if version >= 4:
-import meter_band
-import bsn_tlv
-:: #endif
+import const
+:: for module in modules:
+import ${module}
+:: #endfor
 from const import *
 from common import *
 from loxi import ProtocolError
diff --git a/py_gen/templates/module.py b/py_gen/templates/module.py
index dfe23e8..331044d 100644
--- a/py_gen/templates/module.py
+++ b/py_gen/templates/module.py
@@ -34,20 +34,9 @@
 import struct
 import loxi
 import const
-import common
-import action
-:: if version >= OFVersions.VERSION_1_1:
-import instruction
-:: #endif
-:: if version >= OFVersions.VERSION_1_2:
-import oxm
-:: #endif
-:: if version >= OFVersions.VERSION_1_3:
-import action_id
-import instruction_id
-import meter_band
-import bsn_tlv
-:: #endif
+:: for module in modules:
+import ${module}
+:: #endfor
 import util
 import loxi.generic_util