diff --git a/cppwg/writers/class_writer.py b/cppwg/writers/class_writer.py
index c35afdf..5810459 100644
--- a/cppwg/writers/class_writer.py
+++ b/cppwg/writers/class_writer.py
@@ -27,8 +27,8 @@ class CppClassWrapperWriter(CppBaseWrapperWriter):
         The class information
     wrapper_templates : Dict[str, str]
         String templates with placeholders for generating wrapper code
-    module_class_decls : List[pygccxml.declarations.class_t]
-        A list of decls for all classes in the module
+    module_classes : Dict[pygccxml.declarations.class_t, str]
+        A dictionary of decls and names for all classes in the module
     has_shared_ptr : bool
         Whether the class uses shared pointers
     hpp_string : str
@@ -41,19 +41,19 @@ def __init__(
         self,
         class_info: "CppClassInfo",  # noqa: F821
         wrapper_templates: Dict[str, str],
-        module_class_decls: List["class_t"],  # noqa: F821
+        module_classes: Dict["class_t", str],  # noqa: F821
     ) -> None:
         logger = logging.getLogger()
 
         super().__init__(wrapper_templates)
 
-        self.class_info: "CppClassInfo" = class_info  # noqa: F821
+        self.class_info = class_info
 
         if len(self.class_info.cpp_names) != len(self.class_info.py_names):
             logger.error("C++ and Python class name lists should be the same length")
             raise AssertionError()
 
-        self.module_class_decls: List["class_t"] = module_class_decls  # noqa: F821
+        self.module_classes = module_classes
 
         self.has_shared_ptr: bool = True
 
@@ -318,8 +318,8 @@ def write(self, work_dir: str) -> None:
                     continue
 
                 # Check if the base class is also wrapped in the module
-                if base.related_class in self.module_class_decls:
-                    bases += f", {base.related_class.name}"
+                if base.related_class in self.module_classes:
+                    bases += f", {self.module_classes[base]}"
 
             # Add the class registration
             class_definition_dict = {
diff --git a/cppwg/writers/module_writer.py b/cppwg/writers/module_writer.py
index 50fca41..269a9c8 100644
--- a/cppwg/writers/module_writer.py
+++ b/cppwg/writers/module_writer.py
@@ -2,7 +2,7 @@
 
 import logging
 import os
-from typing import Dict, List
+from typing import Dict
 
 from cppwg.utils.constants import CPPWG_EXT, CPPWG_HEADER_COLLECTION_FILENAME
 from cppwg.writers.class_writer import CppClassWrapperWriter
@@ -26,8 +26,9 @@ class CppModuleWrapperWriter:
         String templates with placeholders for generating wrapper code
     wrapper_root : str
         The output directory for the generated wrapper code
-    class_decls : List[pygccxml.declarations.class_t]
-        A list of declarations of all classes to be wrapped in the module
+
+    classes : Dict[pygccxml.declarations.class_t, str]
+        A dictionary of decls and names for all classes to be wrapped in the module
     """
 
     def __init__(
@@ -40,16 +41,17 @@ def __init__(
         self.wrapper_templates: Dict[str, str] = wrapper_templates
         self.wrapper_root: str = wrapper_root
 
-        # For convenience, store a list of declarations of all
+        # For convenience, store a dictionary of decl->name pairs for all
         # classes to be wrapped in the module
-        self.class_decls: List["class_t"] = []  # noqa: F821
+        self.classes: Dict["class_t", str] = {}  # noqa: F821
 
         for class_info in self.module_info.class_collection:
             # Skip excluded classes
             if class_info.excluded:
                 continue
 
-            self.class_decls.extend(class_info.decls)
+            for decl, cpp_name in zip(class_info.decls, class_info.cpp_names):
+                self.classes[decl] = cpp_name
 
     def write_module_wrapper(self) -> None:
         """
@@ -162,7 +164,7 @@ def write_class_wrappers(self) -> None:
             class_writer = CppClassWrapperWriter(
                 class_info,
                 self.wrapper_templates,
-                self.class_decls,
+                self.classes,
             )
 
             # Write the class wrappers into /path/to/wrapper_root/modulename/