--- a/mamba/engine.py	Sun Sep 11 16:54:09 2011 +0200
+++ b/mamba/engine.py	Sun Sep 11 17:23:42 2011 +0200
@@ -24,7 +24,6 @@
         """Game loop."""
         get_events = pygame.event.get
         flip = pygame.display.flip
-        surface = pygame.display.get_surface()
         while True:
             events = get_events()
             for ev in events:
@@ -34,6 +33,7 @@
                     self.set_habitat(ev.habitat)
                 else:
                     self._habitat.dispatch(ev)
+            surface = pygame.display.get_surface()
             self._habitat.draw(surface)
             flip()
 
@@ -41,24 +41,27 @@
 class Habitat(object):
 
     def __init__(self):
-        self.surface = pygame.Surface(SCREEN)
+        self.surface = None
         self.container = Container(pygame.Rect((0, 0), SCREEN))
 
     def on_enter(self):
         """Called when this becomes the current habitat."""
-        pass
+        # Create the surface here as flipping between editor and
+        # other things kills pygame.display
+        self.surface = pygame.Surface(SCREEN)
 
     def on_exit(self):
         """Called when this stops being the current habitat."""
-        pass
+        self.surface = None
 
     def dispatch(self, ev):
         self.container.event(ev)
 
     def draw(self, surface):
-        self.surface.fill(pygame.Color('black'))
-        self.container.draw(self.surface)
-        surface.blit(self.surface, self.surface.get_rect())
+        if self.surface:
+            self.surface.fill(pygame.Color('black'))
+            self.container.draw(self.surface)
+            surface.blit(self.surface, self.surface.get_rect())
 
 
 class UserEvent(object):

--- a/mamba/habitats/editor.py	Sun Sep 11 16:54:09 2011 +0200
+++ b/mamba/habitats/editor.py	Sun Sep 11 17:23:42 2011 +0200
@@ -1,12 +1,12 @@
 """Habitat for editing levels."""
 
 import pygame.display
-from pygame.locals import SWSURFACE
+from pygame.locals import SWSURFACE, KEYDOWN
 
-from mamba.engine import Habitat
+from mamba.engine import Habitat, NewHabitatEvent
 from mamba.widgets.level import LevelWidget
 from mamba.level import Level
-from mamba.constants import SCREEN, EDIT_SCREEN, NAME
+from mamba.constants import SCREEN, EDIT_SCREEN, NAME, ESCAPE_KEYS
 
 
 class EditorHabitat(Habitat):
@@ -14,6 +14,7 @@
         super(EditorHabitat, self).__init__()
         self.level = Level(level_name)
         self.container.add(LevelWidget(self.level))
+        self.container.add_callback(KEYDOWN, self.keydown_event)
 
     def on_enter(self):
         # We need to juggle the display to the correct size
@@ -22,11 +23,18 @@
         pygame.display.init()
         pygame.display.set_mode(EDIT_SCREEN, SWSURFACE)
         pygame.display.set_caption('%s Level editor' % NAME)
+        super(EditorHabitat, self).on_enter()
 
     def on_exit(self):
         # We need to juggle the display to the correct size
         # This is a horrible hack
+        super(EditorHabitat, self).on_exit()
         pygame.display.quit()
         pygame.display.init()
         pygame.display.set_mode(SCREEN, SWSURFACE)
         pygame.display.set_caption(NAME)
+
+    def keydown_event(self, ev, widget):
+        if ev.key in ESCAPE_KEYS:
+            from mamba.habitats.mainmenu import MainMenu
+            NewHabitatEvent.post(MainMenu())