diff --git a/py_gen/codegen.py b/py_gen/codegen.py
index 4eafa48..f135f18 100644
--- a/py_gen/codegen.py
+++ b/py_gen/codegen.py
@@ -25,7 +25,7 @@
 # EPL for the specific language governing permissions and limitations
 # under the EPL.
 
-from collections import namedtuple
+from collections import defaultdict
 import loxi_globals
 import struct
 import template_utils
@@ -34,74 +34,74 @@
 import oftype
 from loxi_ir import *
 
-ofclasses_by_version = {}
+modules_by_version = {}
 
-# Return the name for the generated Python class
-def generate_pyname(cls):
-    if utils.class_is_action(cls):
-        return cls[10:]
-    elif utils.class_is_oxm(cls):
-        return cls[7:]
-    elif utils.class_is_meter_band(cls):
-        return cls[14:]
-    elif utils.class_is_instruction(cls):
-        return cls[15:]
-    else:
-        return cls[3:]
+# Map from inheritance root to module name
+roots = {
+    'of_header': 'message',
+    'of_action': 'action',
+    'of_oxm': 'oxm',
+    'of_instruction': 'instruction',
+    'of_meter_band': 'meter_band',
+}
+
+# Return the module and class names for the generated Python class
+def generate_pyname(ofclass):
+    for root, module_name in roots.items():
+        if ofclass.name == root:
+            return module_name, module_name
+        elif ofclass.is_instanceof(root):
+            if root == 'of_header':
+                # The input files don't prefix message names
+                return module_name, ofclass.name[3:]
+            else:
+                return module_name, ofclass.name[len(root)+1:]
+    return 'common', ofclass.name[3:]
 
 # Create intermediate representation, extended from the LOXI IR
 def build_ofclasses(version):
-    ofclasses = []
+    modules = defaultdict(list)
     for ofclass in loxi_globals.ir[version].classes:
-        if ofclass.virtual:
-            continue
-
-        ofclass.pyname = generate_pyname(ofclass.name)
-        ofclasses.append(ofclass)
-
-    return ofclasses
+        module_name, ofclass.pyname = generate_pyname(ofclass)
+        modules[module_name].append(ofclass)
+    return modules
 
 def generate_init(out, name, version):
     util.render_template(out, 'init.py', version=version)
 
 def generate_action(out, name, version):
-    ofclasses = [x for x in ofclasses_by_version[version]
-                 if utils.class_is_action(x.name)]
-    util.render_template(out, 'action.py', ofclasses=ofclasses, version=version)
+    util.render_template(out, 'action.py',
+                         ofclasses=modules_by_version[version]['action'],
+                         version=version)
 
 def generate_oxm(out, name, version):
-    ofclasses = [x for x in ofclasses_by_version[version]
-                 if utils.class_is_oxm(x.name)]
-    util.render_template(out, 'oxm.py', ofclasses=ofclasses, version=version)
+    util.render_template(out, 'oxm.py',
+                         ofclasses=modules_by_version[version]['oxm'],
+                         version=version)
 
 def generate_common(out, name, version):
-    ofclasses = [x for x in ofclasses_by_version[version]
-                 if not utils.class_is_message(x.name)
-                    and not utils.class_is_action(x.name)
-                    and not utils.class_is_instruction(x.name)
-                    and not utils.class_is_meter_band(x.name)
-                    and not utils.class_is_oxm(x.name)
-                    and not utils.class_is_list(x.name)]
-    util.render_template(out, 'common.py', ofclasses=ofclasses, version=version)
+    util.render_template(out, 'common.py',
+                         ofclasses=modules_by_version[version]['common'],
+                         version=version)
 
 def generate_const(out, name, version):
     util.render_template(out, 'const.py', version=version,
                          enums=loxi_globals.ir[version].enums)
 
 def generate_instruction(out, name, version):
-    ofclasses = [x for x in ofclasses_by_version[version]
-                 if utils.class_is_instruction(x.name)]
-    util.render_template(out, 'instruction.py', ofclasses=ofclasses, version=version)
+    util.render_template(out, 'instruction.py',
+                         ofclasses=modules_by_version[version]['instruction'],
+                         version=version)
 
 def generate_message(out, name, version):
-    ofclasses = [x for x in ofclasses_by_version[version]
-                 if utils.class_is_message(x.name)]
-    util.render_template(out, 'message.py', ofclasses=ofclasses, version=version)
+    util.render_template(out, 'message.py',
+                         ofclasses=modules_by_version[version]['message'],
+                         version=version)
 
 def generate_meter_band(out, name, version):
-    ofclasses = [x for x in ofclasses_by_version[version]
-                 if utils.class_is_meter_band(x.name)]
-    util.render_template(out, 'meter_band.py', ofclasses=ofclasses, version=version)
+    util.render_template(out, 'meter_band.py',
+                         ofclasses=modules_by_version[version]['meter_band'],
+                         version=version)
 
 def generate_pp(out, name, version):
     util.render_template(out, 'pp.py')
@@ -111,4 +111,4 @@
 
 def init():
     for version in loxi_globals.OFVersions.target_versions:
-        ofclasses_by_version[version] = build_ofclasses(version)
+        modules_by_version[version] = build_ofclasses(version)