From 95c7d7a500e0a41833fb706233ad9ef782cbfef9 Mon Sep 17 00:00:00 2001
From: pamburus <mr.pavel.ivanov@gmail.com>
Date: Wed, 5 Jun 2024 19:27:25 +0200
Subject: [PATCH] new: added Application.QueueUpdateContext and
 Application.QueueUpdateDrawContext

---
 application.go | 46 ++++++++++++++++++++++++++++++++++++++++++----
 1 file changed, 42 insertions(+), 4 deletions(-)

diff --git a/application.go b/application.go
index bbd15d27..a8a6d488 100644
--- a/application.go
+++ b/application.go
@@ -1,6 +1,7 @@
 package tview
 
 import (
+	"context"
 	"strings"
 	"sync"
 	"time"
@@ -478,7 +479,7 @@ EventLoop:
 		case update := <-a.updates:
 			update.f()
 			if update.done != nil {
-				update.done <- struct{}{}
+				close(update.done)
 			}
 		}
 	}
@@ -843,12 +844,40 @@ func (a *Application) GetFocus() Primitive {
 //
 // This function returns after f has executed.
 func (a *Application) QueueUpdate(f func()) *Application {
-	ch := make(chan struct{})
-	a.updates <- queuedUpdate{f: f, done: ch}
-	<-ch
+	_ = a.QueueUpdateContext(context.Background(), f)
+
 	return a
 }
 
+// QueueUpdateContext is used to synchronize access to primitives from non-main
+// goroutines. The provided function will be executed as part of the event loop
+// and thus will not cause race conditions with other such update functions or
+// the Draw() function.
+//
+// Note that Draw() is not implicitly called after the execution of f as that
+// may not be desirable. You can call Draw() from f if the screen should be
+// refreshed after each update. Alternatively, use QueueUpdateDraw() to follow
+// up with an immediate refresh of the screen.
+//
+// This function returns with nil after f has executed or with an error
+// if context is done before that.
+func (a *Application) QueueUpdateContext(ctx context.Context, f func()) error {
+	ch := make(chan struct{})
+	select {
+	case <-ctx.Done():
+		return ctx.Err()
+	case a.updates <- queuedUpdate{f: f, done: ch}:
+	}
+
+	select {
+	case <-ctx.Done():
+		return ctx.Err()
+	case <-ch:
+	}
+
+	return nil
+}
+
 // QueueUpdateDraw works like QueueUpdate() except it refreshes the screen
 // immediately after executing f.
 func (a *Application) QueueUpdateDraw(f func()) *Application {
@@ -859,6 +888,15 @@ func (a *Application) QueueUpdateDraw(f func()) *Application {
 	return a
 }
 
+// QueueUpdateDrawContext works like QueueUpdateContext() except it refreshes the screen
+// immediately after executing f.
+func (a *Application) QueueUpdateDrawContext(ctx context.Context, f func()) error {
+	return a.QueueUpdateContext(ctx, func() {
+		f()
+		a.draw()
+	})
+}
+
 // QueueEvent sends an event to the Application event loop.
 //
 // It is not recommended for event to be nil.