From 6ba3fef9bc7f05868c46135b5923e72f1b1b2d63 Mon Sep 17 00:00:00 2001
From: Oliver <480930+rivo@users.noreply.github.com>
Date: Tue, 26 Dec 2017 17:34:58 +0100
Subject: [PATCH] Added a method to Application to intercept key presses.

---
 application.go | 98 +++++++++++++++++++++++++++++++++++++++++++++-----
 demos/basic.go |  5 +++
 2 files changed, 94 insertions(+), 9 deletions(-)

diff --git a/application.go b/application.go
index d57cc44..54038e1 100644
--- a/application.go
+++ b/application.go
@@ -7,8 +7,12 @@ import (
 )
 
 // Application represents the top node of an application.
+//
+// It is not strictly required to use this class as none of the other classes
+// depend on it. However, it provides useful tools to set up an application and
+// plays nicely with all widgets.
 type Application struct {
-	sync.Mutex
+	sync.RWMutex
 
 	// The application's screen.
 	screen tcell.Screen
@@ -18,11 +22,57 @@ type Application struct {
 
 	// The root primitive to be seen on the screen.
 	root Primitive
+
+	// Key overrides.
+	keyOverrides map[tcell.Key]func(p Primitive) bool
+
+	// Rune overrides.
+	runeOverrides map[rune]func(p Primitive) bool
 }
 
 // NewApplication creates and returns a new application.
 func NewApplication() *Application {
-	return &Application{}
+	return &Application{
+		keyOverrides:  make(map[tcell.Key]func(p Primitive) bool),
+		runeOverrides: make(map[rune]func(p Primitive) bool),
+	}
+}
+
+// SetKeyCapture installs a global capture function for the given key. It
+// intercepts all events for the given key and routes them to the handler.
+// The handler receives the Primitive to which the key is originally redirected,
+// the one which has focus, or nil if it was not directed to a Primitive. The
+// handler also returns whether or not the key event is then forwarded to that
+// Primitive.
+//
+// Special keys (e.g. Escape, Enter, or Ctrl-A) are defined by the "key"
+// argument. The "ch" rune is ignored. Other keys (e.g. "a", "h", or "5") are
+// specified by their rune, with key set to tcell.KeyRune. See also
+// https://godoc.org/github.com/gdamore/tcell#EventKey for more information.
+//
+// To remove a handler again, provide a nil handler for the same key.
+//
+// The application itself will exit when Ctrl-C is pressed. You can intercept
+// this with this function as well.
+func (a *Application) SetKeyCapture(key tcell.Key, ch rune, handler func(p Primitive) bool) *Application {
+	if key == tcell.KeyRune {
+		if handler != nil {
+			a.runeOverrides[ch] = handler
+		} else {
+			if _, ok := a.runeOverrides[ch]; ok {
+				delete(a.runeOverrides, ch)
+			}
+		}
+	} else {
+		if handler != nil {
+			a.keyOverrides[key] = handler
+		} else {
+			if _, ok := a.keyOverrides[key]; ok {
+				delete(a.keyOverrides, key)
+			}
+		}
+	}
+	return a
 }
 
 // Run starts the application and thus the event loop. This function returns
@@ -58,21 +108,47 @@ func (a *Application) Run() error {
 
 	// Start event loop.
 	for {
+		a.RLock()
 		if a.screen == nil {
+			a.RUnlock()
 			break
 		}
 		event := a.screen.PollEvent()
+		a.RUnlock()
 		if event == nil {
 			break // The screen was finalized.
 		}
 		switch event := event.(type) {
 		case *tcell.EventKey:
-			if event.Key() == tcell.KeyCtrlC {
-				a.Stop() // Ctrl-C closes the application.
+			a.RLock()
+			p := a.focus
+			a.RUnlock()
+
+			// Intercept keys.
+			if event.Key() == tcell.KeyRune {
+				if handler, ok := a.runeOverrides[event.Rune()]; ok {
+					if !handler(p) {
+						break
+					}
+				}
+			} else {
+				if handler, ok := a.keyOverrides[event.Key()]; ok {
+					pr := p
+					if event.Key() == tcell.KeyCtrlC {
+						pr = nil
+					}
+					if !handler(pr) {
+						break
+					}
+				}
 			}
-			a.Lock()
-			p := a.focus // Pass other key events to the currently focused primitive.
-			a.Unlock()
+
+			// Ctrl-C closes the application.
+			if event.Key() == tcell.KeyCtrlC {
+				a.Stop()
+			}
+
+			// Pass other key events to the currently focused primitive.
 			if p != nil {
 				if handler := p.InputHandler(); handler != nil {
 					handler(event, func(p Primitive) {
@@ -91,6 +167,8 @@ func (a *Application) Run() error {
 
 // Stop stops the application, causing Run() to return.
 func (a *Application) Stop() {
+	a.RLock()
+	defer a.RUnlock()
 	if a.screen == nil {
 		return
 	}
@@ -101,8 +179,8 @@ func (a *Application) Stop() {
 // Draw refreshes the screen. It calls the Draw() function of the application's
 // root primitive and then syncs the screen buffer.
 func (a *Application) Draw() *Application {
-	a.Lock()
-	defer a.Unlock()
+	a.RLock()
+	defer a.RUnlock()
 
 	// Maybe we're not ready yet or not anymore.
 	if a.screen == nil || a.root == nil {
@@ -156,5 +234,7 @@ func (a *Application) SetFocus(p Primitive) *Application {
 // GetFocus returns the primitive which has the current focus. If none has it,
 // nil is returned.
 func (a *Application) GetFocus() Primitive {
+	a.RLock()
+	defer a.RUnlock()
 	return a.focus
 }
diff --git a/demos/basic.go b/demos/basic.go
index 6ed8e57..6ca2e16 100644
--- a/demos/basic.go
+++ b/demos/basic.go
@@ -19,6 +19,11 @@ func main() {
 	pages := tview.NewPages()
 	list := tview.NewList()
 
+	app.SetKeyCapture(tcell.KeyCtrlQ, 0, func(p tview.Primitive) bool {
+		app.Stop()
+		return false
+	})
+
 	form := tview.NewForm().
 		AddInputField("First name", "", 20, nil).
 		AddInputField("Last name", "", 20, nil).