diff --git a/py_gen/codegen.py b/py_gen/codegen.py
index 9c32e7f..bbf9c5d 100644
--- a/py_gen/codegen.py
+++ b/py_gen/codegen.py
@@ -26,15 +26,20 @@
 # 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 wire version to directory name
+versions = {
+    1: "of10",
+    2: "of11",
+    3: "of12",
+    4: "of13",
+}
 
 # Map from inheritance root to module name
 roots = {
@@ -69,63 +74,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 wire_version, subdir in versions.items():
+        version = loxi_globals.OFVersions.from_wire(wire_version)
+        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