--- a/mamba/widgets/base.py	Wed Sep 14 20:31:16 2011 +0200
+++ b/mamba/widgets/base.py	Wed Sep 14 20:49:48 2011 +0200
@@ -16,6 +16,7 @@
         self.rect = rect
         self.focussable = False
         self.focussed = False
+        self.modal = False
         self.parent = None
         self.callbacks = collections.defaultdict(list)
 
@@ -42,9 +43,37 @@
         pass
 
     def grab_focus(self):
-        if self.focussable and self.parent:
-            self.parent.defocus()
+        if self.focussable:
+            # Find the root and current focus
+            root = self
+            this_modal_base = None
+            while root.parent is not None:
+                if root.modal and this_modal_base is None:
+                    this_modal_base = root
+                root = root.parent
+            focus = root
+            focus_modal_base = None
+            while (isinstance(focus, Container)
+                    and focus.focussed_child is not None):
+                if focus.modal:
+                    focus_modal_base = focus
+                focus = focus.children[focus.focussed_child]
+
+            # Don't leave a modal Widget
+            if focus_modal_base and focus_modal_base != this_modal_base:
+                return False
+
+            root.defocus()
+            widget = self
+            while widget.parent is not None:
+                parent = widget.parent
+                if isinstance(parent, Container):
+                    idx = parent.children.index(widget)
+                    parent.focussed_child = idx
+                widget = parent
             self.focussed = True
+            return True
+        return False
 
 
 class Button(Widget):
@@ -75,25 +104,13 @@
             for child in self.children:
                 if child.rect.collidepoint(ev.pos):
                     if ev.type == MOUSEBUTTONDOWN and child.focussable:
-                        root = self
-                        while root.parent is not None:
-                            root = root.parent
-                        root.defocus()
-                        widget = child
-                        while widget.parent is not None:
-                            parent = widget.parent
-                            if isinstance(parent, Container):
-                                idx = parent.children.index(widget)
-                                parent.focussed_child = idx
-                            widget = parent
-                        child.focussed = True
-
+                        child.grab_focus()
                     if child.event(ev):
                         return True
 
         elif ev.type == KEYDOWN:
-            for child in self.children:
-                if child.focussed:
+            for i, child in enumerate(self.children):
+                if child.focussed or i == self.focussed_child:
                     if child.event(ev):
                         return True
         else:
@@ -128,8 +145,11 @@
         if self.focussed_child is not None:
             child = self.children[self.focussed_child]
             if isinstance(child, Container):
-                child.defocus()
+                if not child.defocus():
+                    return False
             child.focussed = False
+            self.focussed_child = None
+            return True
 
     def adjust_focus(self, direction):
         """Try and adjust focus in direction (integer)
@@ -139,6 +159,10 @@
             if isinstance(child, Container):
                 if child.adjust_focus(direction):
                     return True
+                elif child.modal:
+                    # We're modal, go back
+                    if child.adjust_focus(-direction):
+                        return True
             else:
                 child.focussed = False
 

--- a/mamba/widgets/game.py	Wed Sep 14 20:31:16 2011 +0200
+++ b/mamba/widgets/game.py	Wed Sep 14 20:49:48 2011 +0200
@@ -16,6 +16,8 @@
         rect = Rect(offset, world.get_size())
         super(GameWidget, self).__init__(rect)
         self.focussable = True
+        self.add_callback(KEYDOWN, self.action_callback)
+        self.add_callback(SnakeDiedEvent, self.snake_died)
 
     def create_action_map(self):
         actions = {}
@@ -25,24 +27,22 @@
         actions[K_UP] = (self.world.snake.set_orientation, (UP,))
         return actions
 
-    def event(self, ev):
-        if ev.type == KEYDOWN and ev.key in self.actions:
+    def action_callback(self, ev, widget):
+        if ev.key in self.actions:
             func, args = self.actions[ev.key]
             func(*args)
             return True
-        elif SnakeDiedEvent.matches(ev):
-            self.snake_died()
-            return True
 
     def draw(self, surface):
         self.world.update()
         self.world.draw(surface)
 
-    def snake_died(self):
+    def snake_died(self, ev, widget):
         self.world.pause()
         messagebox = MessageBox((300, 200), 'You died!', self.restart)
         self.parent.add(messagebox)
         messagebox.grab_focus()
+        return True
 
     def restart(self):
         self.world.restart()

--- a/mamba/widgets/messagebox.py	Wed Sep 14 20:31:16 2011 +0200
+++ b/mamba/widgets/messagebox.py	Wed Sep 14 20:49:48 2011 +0200
@@ -1,5 +1,5 @@
 import pygame
-from pygame.constants import SRCALPHA, K_RETURN, KEYDOWN
+from pygame.constants import SRCALPHA
 
 from mamba.widgets.base import Container
 from mamba.widgets.text import TextWidget, TextButton
@@ -14,7 +14,7 @@
         self.post_callback = post_callback
         self.color = color
         self.prepare()
-        self.focussable = True
+        self.modal = True
 
     def prepare(self):
         cont = TextWidget((0, 0), "Press [OK] or Enter to continue")
@@ -34,7 +34,7 @@
         cont.rect.topleft = cont_pos
         self.add(message)
         self.add(cont)
-        ok_button = TextButton((0, 0), 'OK')
+        self.ok_button = ok_button = TextButton((0, 0), 'OK')
         ok_pos = (self.rect.left + 5 + width / 2 - ok_button.rect.width / 2,
                 cont_pos[1] + 5 + cont.rect.height)
         ok_button.rect.topleft = ok_pos
@@ -47,13 +47,10 @@
         surface.blit(background, self.rect)
         super(MessageBox, self).draw(surface)
 
-    def event(self, ev):
-        if ev.type == KEYDOWN and ev.key == K_RETURN:
-            self.close(ev, self)
-            return True
-        return super(MessageBox, self).event(ev)
-
     def close(self, ev, widget):
         self.parent.remove(self)
         if self.post_callback:
             self.post_callback()
+
+    def grab_focus(self):
+        return self.ok_button.grab_focus()