--- a/mamba/habitats/editor.py	Sat Sep 17 22:10:01 2011 +0200
+++ b/mamba/habitats/editor.py	Sat Sep 17 22:12:16 2011 +0200
@@ -28,6 +28,7 @@
 class EditorHabitat(Habitat):
     def __init__(self, level):
         super(EditorHabitat, self).__init__(EDIT_SCREEN)
+        self.toolbar = {}
         self.level = level
         self.container.paused = False
         self.edit_widget = EditLevelWidget(self.level)
@@ -130,38 +131,43 @@
             self.edit_widget.set_tool('.')
             button_height += (self.floor_button.surface.get_height()
                     + button_padding)
-        if self.mode == 'Tiles':
-            tile_map = TILE_MAP
-        elif self.mode == 'Things':
-            tile_map = THING_MAP
+        if self.mode in self.toolbar:
+            # FIXME: This needs to be recreated on tileset changes
+            self.tool_widget = self.toolbar[self.mode]
         else:
-            tile_map = []
-        tool_list = []
-        for tile_char in sorted(tile_map):
-            try:
-                tile = self.level.tileset[tile_char]
-            except pygame.error:
-                # Ignore stuff we can't load for now
-                continue
-            if tile is None:
-                continue
-            if tile.name:
-                text = tile.name
+            if self.mode == 'Tiles':
+                tile_map = TILE_MAP
+            elif self.mode == 'Things':
+                tile_map = THING_MAP
             else:
-                text = 'Tile'
-            tile_button = ImageButtonWidget((0, 0), tile.image, text,
-                    color='white')
-            tile_button.add_callback('clicked', self.change_tool,
-                    tile_char, text)
-            tool_list.append(tile_button)
-        if self.mode == "Sprites":
-            for name in ['Add', 'Edit', 'Delete']:
-                tile_button = TextButton((0, 0), '%s Sprite' % name)
-                tile_button.add_callback('clicked', self.sprite_tool,
-                        name)
+                tile_map = []
+            tool_list = []
+            for tile_char in sorted(tile_map):
+                try:
+                    tile = self.level.tileset[tile_char]
+                except pygame.error:
+                    # Ignore stuff we can't load for now
+                    continue
+                if tile is None:
+                    continue
+                if tile.name:
+                    text = tile.name
+                else:
+                    text = 'Tile'
+                tile_button = ImageButtonWidget((0, 0), tile.image, text,
+                        color='white')
+                tile_button.add_callback('clicked', self.change_tool,
+                        tile_char, text)
                 tool_list.append(tile_button)
-        self.tool_widget = ToolListWidget((button_left, button_height),
-                tool_list, MAX_TOOLS, start_key=K_2)
+            if self.mode == "Sprites":
+                for name in ['Add', 'Edit', 'Delete']:
+                    tile_button = TextButton((0, 0), '%s Sprite' % name)
+                    tile_button.add_callback('clicked', self.sprite_tool,
+                            name)
+                    tool_list.append(tile_button)
+            self.tool_widget = ToolListWidget((button_left, button_height),
+                    tool_list, MAX_TOOLS, start_key=K_2)
+            self.toolbar[self.mode] = self.tool_widget
         self.container.add(self.tool_widget)
         button_height = self.container.rect.top + MODE_HEIGHT
         if self.mode == 'Tiles':