--- 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):