From 7e17ed21cede3b2836aec2345a586ac5f7ce50df Mon Sep 17 00:00:00 2001
From: Kovid Goyal <kovid@kovidgoyal.net>
Date: Mon, 27 Sep 2021 12:00:33 +0530
Subject: [PATCH] Use the atomic update decorator everywhere

---
 kittens/themes/main.py        | 20 ++++++++++----------
 kittens/transfer/send.py      |  3 ++-
 kittens/tui/handler.py        | 15 ++++++---------
 kittens/unicode_input/main.py |  2 +-
 kitty/types.py                |  8 +++++++-
 5 files changed, 26 insertions(+), 22 deletions(-)

diff --git a/kittens/themes/main.py b/kittens/themes/main.py
index 44ff931c0..e0597360c 100644
--- a/kittens/themes/main.py
+++ b/kittens/themes/main.py
@@ -482,17 +482,17 @@ class ThemesHandler(Handler):
         elif self.state is State.accepting:
             self.on_accepting_key_event(key_event, in_bracketed_paste)
 
+    @Handler.atomic_update
     def draw_screen(self) -> None:
-        with self.pending_update():
-            self.cmd.clear_screen()
-            self.enforce_cursor_state()
-            self.cmd.set_line_wrapping(False)
-            if self.state is State.fetching:
-                self.draw_fetching_screen()
-            elif self.state in (State.browsing, State.searching):
-                self.draw_browsing_screen()
-            elif self.state is State.accepting:
-                self.draw_accepting_screen()
+        self.cmd.clear_screen()
+        self.enforce_cursor_state()
+        self.cmd.set_line_wrapping(False)
+        if self.state is State.fetching:
+            self.draw_fetching_screen()
+        elif self.state in (State.browsing, State.searching):
+            self.draw_browsing_screen()
+        elif self.state is State.accepting:
+            self.draw_accepting_screen()
 
     def on_resize(self, screen_size: ScreenSize) -> None:
         self.screen_size = screen_size
diff --git a/kittens/transfer/send.py b/kittens/transfer/send.py
index f719a61b6..dfb321a9e 100644
--- a/kittens/transfer/send.py
+++ b/kittens/transfer/send.py
@@ -566,8 +566,9 @@ class Send(Handler):
             self.failed_files.append(file)
         self.asyncio_loop.call_soon(self.refresh_progress)
 
+    @Handler.atomic_update
     def draw_progress(self) -> None:
-        with self.pending_update(), without_line_wrap(self.write):
+        with without_line_wrap(self.write):
             for df in self.done_files:
                 sc = styled('✔', fg='green') if not df.err_msg else styled('✘', fg='red')
                 self.draw_progress_for_current_file(df, spinner_char=sc, is_complete=True)
diff --git a/kittens/tui/handler.py b/kittens/tui/handler.py
index a7cb98d70..20a940692 100644
--- a/kittens/tui/handler.py
+++ b/kittens/tui/handler.py
@@ -5,10 +5,11 @@
 
 from types import TracebackType
 from typing import (
-    Any, Callable, ContextManager, Dict, Optional, Sequence, Type, Union, TYPE_CHECKING
+    TYPE_CHECKING, Any, Callable, ContextManager, Dict, Optional, Sequence,
+    Type, Union, cast
 )
 
-from kitty.types import ParsedShortcut
+from kitty.types import DecoratedFunc, ParsedShortcut
 from kitty.typing import (
     AbstractEventLoop, BossType, Debug, ImageManagerType, KeyActionType,
     KeyEventType, LoopType, MouseEvent, ScreenSize, TermManagerType
@@ -16,7 +17,6 @@ from kitty.typing import (
 
 from .operations import pending_update
 
-
 if TYPE_CHECKING:
     from kitty.file_transmission import FileTransmissionCommand
 
@@ -144,18 +144,15 @@ class Handler:
     def suspend(self) -> ContextManager[TermManagerType]:
         return self._term_manager.suspend()
 
-    def pending_update(self) -> ContextManager[None]:
-        return pending_update(self.write)
-
     @classmethod
-    def with_pending_update(cls, func: Callable) -> Callable:
+    def atomic_update(cls, func: DecoratedFunc) -> DecoratedFunc:
         from functools import wraps
 
         @wraps(func)
         def f(*a: Any, **kw: Any) -> Any:
-            with a[0].pending_update():
+            with pending_update(a[0].write):
                 return func(*a, **kw)
-        return f
+        return cast(DecoratedFunc, f)
 
 
 class HandleResult:
diff --git a/kittens/unicode_input/main.py b/kittens/unicode_input/main.py
index bcfc60827..233784d6b 100644
--- a/kittens/unicode_input/main.py
+++ b/kittens/unicode_input/main.py
@@ -391,7 +391,7 @@ class UnicodeInput(Handler):
             text += ' ' * extra
         self.print(styled(text, reverse=True))
 
-    @Handler.with_pending_update
+    @Handler.atomic_update
     def draw_screen(self) -> None:
         self.write(clear_screen())
         self.draw_title_bar()
diff --git a/kitty/types.py b/kitty/types.py
index 61dcda1de..33dcdb2ae 100644
--- a/kitty/types.py
+++ b/kitty/types.py
@@ -3,7 +3,9 @@
 # License: GPLv3 Copyright: 2021, Kovid Goyal <kovid at kovidgoyal.net>
 
 from functools import update_wrapper
-from typing import TYPE_CHECKING, Callable, Generic, NamedTuple, TypeVar, Union
+from typing import (
+    TYPE_CHECKING, Any, Callable, Generic, NamedTuple, TypeVar, Union
+)
 
 _T = TypeVar('_T')
 
@@ -116,3 +118,7 @@ def ac(group: ActionGroup, doc: str) -> Callable[[_T], _T]:
         setattr(f, 'action_spec', ActionSpec(group, doc))
         return f
     return w
+
+
+_BaseDecoratedFunc = Callable[..., Any]
+DecoratedFunc = TypeVar('DecoratedFunc', bound=_BaseDecoratedFunc)