--- a/mamba/habitats/editor.py	Sun Sep 11 15:42:57 2011 +0200
+++ b/mamba/habitats/editor.py	Sun Sep 11 15:47:47 2011 +0200
@@ -1,8 +1,12 @@
 """Habitat for editing levels."""
 
+import pygame.display
+from pygame.locals import SWSURFACE
+
 from mamba.engine import Habitat
 from mamba.widgets.level import LevelWidget
 from mamba.level import Level
+from mamba.constants import SCREEN, EDIT_SCREEN, NAME
 
 
 class EditorHabitat(Habitat):
@@ -10,3 +14,19 @@
         super(EditorHabitat, self).__init__()
         self.level = Level(level_name)
         self.container.add(LevelWidget(self.level))
+
+    def on_enter(self):
+        # We need to juggle the display to the correct size
+        # This is a horrible hack
+        pygame.display.quit()
+        pygame.display.init()
+        pygame.display.set_mode(EDIT_SCREEN, SWSURFACE)
+        pygame.display.set_caption('%s Level editor' % NAME)
+
+    def on_exit(self):
+        # We need to juggle the display to the correct size
+        # This is a horrible hack
+        pygame.display.quit()
+        pygame.display.init()
+        pygame.display.set_mode(SCREEN, SWSURFACE)
+        pygame.display.set_caption(NAME)