Use the atomic update decorator everywhere

kittens/themes/main.py

@@ -482,8 +482,8 @@ 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)

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)

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:

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()

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)