Repository: AnonymouX47/term-image (language: python)

Issue:
Support for terminal synchronized output
The attached patch is an update for UrwidImageScreen to support the [terminal sychronized output spec](https://gist.github.com/christianparpart/d8a62cc1ab659194337d73e399004036)

[sync.patch](https://github.com/AnonymouX47/term-image/files/10887436/sync.patch)

I thought it might help the Kitty scroll flicker issue and maybe it does - a little? But the flicker definitely is still there.

Give it a try and see for yourself if it helps anything.  Should be harmless on terminals that do not support the spec - but I know that recent versions of Kitty and iTerm2 do.
Implement synchronized output for urwid images
Partially implements #79
See also: #73

- Moves automatic image clearing into `UrwidImageScreen.draw_screen()`.
- Removes `UrwidImageJanitor` in favour of `UrwidImageScreen`.
- Adds support for synchronized output to `UrwidImageScreen`.
  - Reduces/Eliminates image flickering/tearing on supporting terminal emulators.


Discussion / hints:
I see. This is actually potentially very useful, even beyond just the urwid widget. I wasn't aware of this mode.

By the way, about your patch:
```py
def write(self, data):
    super().write("\033[?2026h" + data + "\033[?2026l")
```
1. Shouldn't be in `Screen.write()`.
   - This is called multiple times when drawing the screen. So, the sequences will surround every chunk of data written and won't make any difference.
   - the output may or may not be flushed multiple times (even though `Screen.flush()` is called only once when drawing the screen, buffer size and line buffering may trigger a flush).
3. Should be in [`Screen.draw_screen()`](https://github.com/urwid/urwid/blob/95c15ce0b991b28f178498dea18f4e89d22d783c/urwid/raw_display.py#L695) instead.
   - Calls to `Screen.write()` just before and after a call to `super().draw_screen()`

I'll try it out with `Screen.draw_screen()` and I have high hopes for better results. 😃 By the way, if this works out well, would you be interested in opening the PR that implements this?

Thanks so much.
I've tried it out and yes, it improves (or completely eliminates on some e.g Wezterm) the flickering situation on terminal emulators that support the feature.

I observed that Kitty still flickers when images get large (amount of bytes for the escape sequence, not necessarily size in pixels) enough (not sure what the exact threshold is) but it's still better than without the sequence. Anyways, as long as there are no large (by size in pixels) and (not "or") pixel-dense images on-screen, it should be just fine.

Looking at the [original spec](https://gitlab.com/gnachman/iterm2/-/wikis/synchronized-updates-spec) for the feature, it states:

> ### Optional Conditions
> Terminal emulators may choose to end an atomic update when:
> - The number of bytes received since the update began exceeds some threshold
> - ...

I suspect the reason for kitty still flickering is either due to the condition above or a timeout (also specified in the spec).
While trying this out, I realized that it can't work along with `UrwidImageJanitor` because it writes the image deletion sequences before `Screen.draw_screen()` is called. I adapted `UrwidImageJanitor.render()` into `UrwidImageScreen` and then I thought, "Why not just let `UrwidImageScreen` handle it all?" i.e clearing images and all.

I don't know why it never occurred to me until now but the `Screen` actually seems more like the right place for "clearing images of the **screen**". Also, the fact that it actually isn't ideal (even though it works :grin:) for a widget to explicitly write to output.

I'm considering moving this functionality (automatically clearing images) into `UrwidImageScreen` and scrapping `UrwidImageJanitor` entirely, after all, that's all it does. The only reason I might want to leave `UrwidImageJanitor` is to allow a user to wrap a single widget that is not the topmost widget, probably for the sake of performance but that would nullify synchronized output. **What say you?**

By the way, I'd also like to add support for this in `BaseImage.draw()`, it'll improve animations in particular... but that will be later.
Thanks for trying this out. I was so happy to see there was an easily overridable method in `Screen.write()` to implement this support, that I didn't bother looking for something higher level like `draw_screen()` 🙄 .  That's definitely a better place to put it.  

And you are probably correct that the remaining flicker with Kitty is due to an aggressive timeout setting or bytes-transferred setting.  I might dig into that code looking for what they do, if I have some time. 

I think that the right thing to do is kill `UrwidImageJanitor` and move the screen clearing code to `UrwidImageScreen`.   Simpler is better. Keep `UrwidImageJanitor` only if performance is unacceptable.
Now that you've implemented the code for testing purposes, do you still need me to write a PR?
>  I might dig into that code looking for what they do, if I have some time.

Which code? 🤔  If it's Kitty's, [this](https://github.com/kovidgoyal/kitty/commit/5768c54c5b5763e4bbb300726b8ff71b40c128f8) might be a starting point, though the feature was already implemented before that.

> I think that the right thing to do is kill UrwidImageJanitor and move the screen clearing code to UrwidImageScreen. Simpler is better. Keep UrwidImageJanitor only if performance is unacceptable.

Ok, good. IMO, the performance is not noticeable (to a human) even for a UI that looks like:

![Screenshot_2023-03-04_17-44-44](https://user-images.githubusercontent.com/61663146/222918281-5dba022d-d119-4c43-987c-1101097b963b.png)

Note that `LineBox`es are actually composed of quite a number of widgets/canvases. So, I guess I'll kill it for now and if anything comes up in the future, I can easily bring it back from the grave.

> Now that you've implemented the code for testing purposes, do you still need me to write a PR?

I just thought since you brought it up... but you're right. I shouldn't bother you. 😃



Fix patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 58871ba4..6d8c41dc 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -28,8 +28,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   - `term_image.widget` subpackage
   - `term_image.widget.UrwidImage`
   - `term_image.widget.UrwidImageCanvas`
-  - `term_image.widget.UrwidImageJanitor`
   - `term_image.widget.UrwidImageScreen`
+    - Support for terminal-synchronized output ([#80]).
 
 ### Changed
 - **(BREAKING!)** Redefined `KittyImage.clear()` ([97eceab]).
@@ -50,6 +50,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 [#73]: https://github.com/AnonymouX47/term-image/pull/73
 [#74]: https://github.com/AnonymouX47/term-image/pull/74
 [#78]: https://github.com/AnonymouX47/term-image/pull/78
+[#80]: https://github.com/AnonymouX47/term-image/pull/80
 [b4533d5]: https://github.com/AnonymouX47/term-image/commit/b4533d5697d41fe0742c2ac895077da3b8d889dc
 [97eceab]: https://github.com/AnonymouX47/term-image/commit/97eceab77e7448a18281aa6edb3fa8ec9e6564c5
 [807a9ec]: https://github.com/AnonymouX47/term-image/commit/807a9ecad717e46621a5214dbf849369d3afbc0b
diff --git a/src/term_image/utils.py b/src/term_image/utils.py
index 5ac10d93..c6d1b023 100644
--- a/src/term_image/utils.py
+++ b/src/term_image/utils.py
@@ -692,6 +692,17 @@ def _process_run_wrapper(self, *args, set_tty_lock: bool = True, **kwargs):
 FG_FMT_b = FG_FMT.encode()
 COLOR_RESET = f"{CSI}m"
 COLOR_RESET_b = COLOR_RESET.encode()
+DECSET = f"{CSI}?%dh"
+DECSET_b = DECSET.encode()
+DECRST = f"{CSI}?%dl"
+DECRST_b = DECRST.encode()
+
+# Terminal Synchronized Output
+# See https://gist.github.com/christianparpart/d8a62cc1ab659194337d73e399004036
+BEGIN_SYNCED_UPDATE = DECSET % 2026
+BEGIN_SYNCED_UPDATE_b = BEGIN_SYNCED_UPDATE.encode()
+END_SYNCED_UPDATE = DECRST % 2026
+END_SYNCED_UPDATE_b = END_SYNCED_UPDATE.encode()
 
 # Private internal variables
 _query_timeout = DEFAULT_QUERY_TIMEOUT
diff --git a/src/term_image/widget/__init__.py b/src/term_image/widget/__init__.py
index eee22e9f..a8a4af0c 100644
--- a/src/term_image/widget/__init__.py
+++ b/src/term_image/widget/__init__.py
@@ -9,12 +9,7 @@
 except ImportError:
     pass
 else:
-    from .urwid import UrwidImage, UrwidImageCanvas, UrwidImageJanitor, UrwidImageScreen
+    from .urwid import UrwidImage, UrwidImageCanvas, UrwidImageScreen
 
     del urwid
-    __all__ += [
-        "UrwidImage",
-        "UrwidImageCanvas",
-        "UrwidImageJanitor",
-        "UrwidImageScreen",
-    ]
+    __all__ += ["UrwidImage", "UrwidImageCanvas", "UrwidImageScreen"]
diff --git a/src/term_image/widget/urwid.py b/src/term_image/widget/urwid.py
index bc6c9573..25a2efe6 100644
--- a/src/term_image/widget/urwid.py
+++ b/src/term_image/widget/urwid.py
@@ -2,7 +2,7 @@
 
 from __future__ import annotations
 
-__all__ = ("UrwidImage", "UrwidImageCanvas", "UrwidImageJanitor")
+__all__ = ("UrwidImage", "UrwidImageCanvas", "UrwidImageScreen")
 
 from typing import Optional, Tuple
 
@@ -10,7 +10,7 @@
 
 from ..exceptions import UrwidImageError
 from ..image import BaseImage, ITerm2Image, KittyImage, Size, TextImage, kitty
-from ..utils import COLOR_RESET_b, ESC_b
+from ..utils import BEGIN_SYNCED_UPDATE, END_SYNCED_UPDATE, COLOR_RESET_b, ESC_b
 
 # NOTE: Any new "private" attribute of any subclass of an urwid class should be
 # prepended with "_ti" to prevent clashes with names used by urwid itself.
@@ -500,49 +500,66 @@ def _ti_calc_trim(
         return new_pad_side1, trim_image_side1, trim_image_side2, new_pad_side2
 
 
-class UrwidImageJanitor(urwid.WidgetWrap):
-    """A widget wrapper that monitors images of some :ref:`graphics-based
-    <graphics-based>` render styles and clears them off the screen when necessary.
+class UrwidImageScreen(urwid.raw_display.Screen):
+    """A screen that supports drawing images.
 
-    Args:
-        widget: A widget containing image widgets (possibly recursively).
+    It monitors images of some :ref:`graphics-based <graphics-based>` render styles
+    and clears them off the screen when necessary (e.g at startup, when scrolling,
+    upon terminal resize and at exit).
 
-    NOTE:
-        For this widget to function properly, ensure that the position of its top-left
-        corner, **relative to the screen** can **never** change.
+    It also synchronizes output on terminal emulators that support the feature to
+    reduce/eliminate image flickering/tearing.
 
-        It's most advisable to use this widget to wrap the topmost widget but it may
-        be used to wrap any other widget and will function properly as long as the
-        condition mentioned earlier is met.
+    See the `baseclass
+    <http://urwid.org/reference/display_modules.html#urwid.raw_display.Screen>`_
+    for futher description.
     """
 
-    no_cache = ["render"]
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._ti_screen_canv = None
+        self._ti_image_cviews = frozenset()
 
-    def __init__(self, widget: urwid.Widget) -> None:
-        if not isinstance(widget, urwid.Widget):
-            raise TypeError(f"Invalid type for 'widget' (got: {type(widget).__name__})")
+    def draw_screen(self, maxres, canvas):
+        self.write(BEGIN_SYNCED_UPDATE)
+        try:
+            if canvas is not self._ti_screen_canv:
+                self._ti_screen_canv = canvas
+                self._ti_clear_images()
+            return super().draw_screen(maxres, canvas)
+        finally:
+            self.write(END_SYNCED_UPDATE)
+            self.flush()
 
-        super().__init__(widget)
-        self._ti_image_cviews = frozenset()
+    def _start(self, *args, **kwargs):
+        ret = super()._start(*args, **kwargs)
+        if KittyImage._forced_support or KittyImage.is_supported():
+            self.write(kitty.DELETE_ALL_IMAGES)
 
-    widget = property(lambda self: self._w, doc="The wrapped widget")
+        return ret
 
-    def render(self, size: Tuple[int, int], focus: bool = False) -> urwid.Canvas:
-        main_canv = super().render(size, focus)
+    def _stop(self):
+        if KittyImage._forced_support or KittyImage.is_supported():
+            self.write(kitty.DELETE_ALL_IMAGES)
 
+        return super()._stop()
+
+    def _ti_clear_images(self):
         if not (
             KittyImage._forced_support
             or KittyImage.is_supported()
             or ITerm2Image.is_supported()
             and ITerm2Image._TERM == "konsole"
         ):
-            return main_canv
+            return
 
-        if not isinstance(main_canv, urwid.CompositeCanvas):
+        screen_canv = self._ti_screen_canv
+
+        if not isinstance(screen_canv, urwid.CompositeCanvas):
             if self._ti_image_cviews:
                 UrwidImage.clear_all()
                 self._ti_image_cviews.clear()
-            return main_canv
+            return
 
         def process_shard_tails():
             nonlocal col
@@ -555,27 +572,28 @@ def process_shard_tails():
                     del shard_tails[col]
                 col += cols
 
-        images = set()
+        image_cviews = set()
         shard_tails = {}
         row = 1
 
-        for n_rows, cviews in main_canv.shards:
+        for n_rows, cviews in screen_canv.shards:
             col = 1
             for cview in cviews:
                 process_shard_tails()
                 *trim, cols, rows, _, canv = cview
 
-                try:
-                    widget = canv.widget_info[0]
-                except TypeError:
-                    pass
-                else:
-                    if isinstance(canv, UrwidImageCanvas) and (
-                        isinstance(widget._ti_image, KittyImage)
-                        or isinstance(widget._ti_image, ITerm2Image)
-                        and ITerm2Image._TERM == "konsole"
-                    ):
-                        images.add((canv, row, col, *trim, cols, rows))
+                if isinstance(canv, UrwidImageCanvas):
+                    try:
+                        widget = canv.widget_info[0]
+                    except TypeError:
+                        pass
+                    else:
+                        if (
+                            isinstance(widget._ti_image, KittyImage)
+                            or isinstance(widget._ti_image, ITerm2Image)
+                            and ITerm2Image._TERM == "konsole"
+                        ):
+                            image_cviews.add((canv, row, col, *trim, cols, rows))
 
                 if rows > n_rows:
                     shard_tails[col] = (*trim, cols, rows - n_rows, canv)
@@ -583,39 +601,18 @@ def process_shard_tails():
             process_shard_tails()
             row += n_rows
 
-        for canv, *_ in self._ti_image_cviews - images:
+        kitty_widgets = []
+        for canv, *_ in self._ti_image_cviews - image_cviews:
             widget = canv.widget_info[0]
             if isinstance(widget._ti_image, KittyImage):
-                widget.clear()
+                kitty_widgets.append(widget)
             else:
                 UrwidImage.clear_all()
                 # Multiple `clear_all()`s messes up the canvas disguise
                 # Also, a single `clear_all()` takes care of all images
                 break
-        self._ti_image_cviews = frozenset(images)
-
-        return main_canv
-
-
-class UrwidImageScreen(urwid.raw_display.Screen):
-    """A screen that clears all visible images of the
-    :py:class:`kitty <term_image.image.KittyImage>` render style immediately after
-    it starts and immediately before it stops.
-
-    See the `baseclass
-    <http://urwid.org/reference/display_modules.html#urwid.raw_display.Screen>`_
-    for futher description.
-    """
-
-    def _start(self, *args, **kwargs):
-        ret = super()._start(*args, **kwargs)
-        if KittyImage._forced_support or KittyImage.is_supported():
-            self.write(kitty.DELETE_ALL_IMAGES)
-
-        return ret
-
-    def _stop(self):
-        if KittyImage._forced_support or KittyImage.is_supported():
-            self.write(kitty.DELETE_ALL_IMAGES)
+        else:
+            for widget in kitty_widgets:
+                widget.clear()
 
-        return super()._stop()
+        self._ti_image_cviews = frozenset(image_cviews)


Test patch:
diff --git a/tests/test_widget/test_urwid.py b/tests/test_widget/test_urwid.py
index 4a0ef566..2627f2f8 100644
--- a/tests/test_widget/test_urwid.py
+++ b/tests/test_widget/test_urwid.py
@@ -18,13 +18,14 @@
     kitty,
 )
 from term_image.image.common import _ALPHA_THRESHOLD
-from term_image.utils import COLOR_RESET, BG_FMT_b, COLOR_RESET_b
-from term_image.widget import (
-    UrwidImage,
-    UrwidImageCanvas,
-    UrwidImageJanitor,
-    UrwidImageScreen,
+from term_image.utils import (
+    BEGIN_SYNCED_UPDATE,
+    COLOR_RESET,
+    END_SYNCED_UPDATE,
+    BG_FMT_b,
+    COLOR_RESET_b,
 )
+from term_image.widget import UrwidImage, UrwidImageCanvas, UrwidImageScreen
 
 _size = (30, 15)
 
@@ -1077,6 +1078,55 @@ def test_horizontal_trim_first_color_ColoredText(self):
                     assert line.startswith(BG_FMT_b % (16, 32, 48) + b" ")
 
 
+class TestScreen:
+    class TestStartStop:
+        def test_supported(self):
+            buf = io.StringIO()
+            screen = UrwidImageScreen(sys.__stdin__, buf)
+
+            screen.start()
+            start_output = buf.getvalue()
+            buf.seek(0)
+            buf.truncate()
+            screen.stop()
+            stop_output = buf.getvalue()
+
+            assert start_output.endswith(kitty.DELETE_ALL_IMAGES)
+            assert stop_output.startswith(kitty.DELETE_ALL_IMAGES)
+
+        def test_not_supported(self):
+            buf = io.StringIO()
+            screen = UrwidImageScreen(sys.__stdin__, buf)
+
+            KittyImage._supported = False
+            try:
+                screen.start()
+                start_output = buf.getvalue()
+                buf.seek(0)
+                buf.truncate()
+                screen.stop()
+                stop_output = buf.getvalue()
+
+                assert kitty.DELETE_ALL_IMAGES not in start_output
+                assert kitty.DELETE_ALL_IMAGES not in stop_output
+            finally:
+                KittyImage._supported = True
+
+    def test_synced_output(self):
+        widget = urwid.SolidFill("x")
+        buf = io.StringIO()
+        screen = UrwidImageScreen(sys.__stdin__, buf)
+        screen.start()
+
+        buf.seek(0)
+        buf.truncate()
+        screen.draw_screen(_size, widget.render(_size))
+        output = buf.getvalue()
+
+        assert output.startswith(BEGIN_SYNCED_UPDATE)
+        assert output.endswith(END_SYNCED_UPDATE)
+
+
 block_image_w = UrwidImage(BlockImage(python_img))
 kitty_image_w = UrwidImage(KittyImage(python_img))
 iterm2_image_w = UrwidImage(ITerm2Image(python_img))
@@ -1101,50 +1151,36 @@ def test_horizontal_trim_first_color_ColoredText(self):
     5,
 )
 
-
-class TestJanitor:
-    def test_args(self):
-        for value in (None, python_file, python_img, python_image):
-            with pytest.raises(TypeError, match="'widget'"):
-                UrwidImageJanitor(value)
-
-    def test_widget_property(self):
-        janitor = UrwidImageJanitor(widget)
-        assert janitor.widget is widget is janitor._w
-
-    def test_same_canvas_as_widget(self):
-        janitor = UrwidImageJanitor(widget)
-        widget_canv = widget.render(_size)
-        widget._invalidate()
-        canv = janitor.render(_size)
-        assert canv.text == widget_canv.text
+buf = io.StringIO()
+screen = UrwidImageScreen(sys.__stdin__, buf)
+screen.start()
 
 
-class TestJanitorRender:
+class TestScreenClearImages:
     def test_image_cviews(self):
-        janitor = UrwidImageJanitor(widget)
-        assert janitor._ti_image_cviews == frozenset()
-
-        janitor.render(_size)
-        assert isinstance(janitor._ti_image_cviews, frozenset)
+        assert screen._ti_image_cviews == frozenset()
+        screen.draw_screen(_size, widget.render(_size))
+        assert isinstance(screen._ti_image_cviews, frozenset)
 
     def test_move_top_widget(self):
         _TERM = ITerm2Image._TERM
         ITerm2Image._TERM = "konsole"
         try:
             # Setup
-            janitor = UrwidImageJanitor(widget)
+            buf.seek(0)
+            buf.truncate()
             widget.top_w = top_w
             kitty_list_box.shift_focus(_size, 0)
             iterm2_list_box.shift_focus(_size, 0)
 
             widget.top = 2
             widget._invalidate()
-            canv = janitor.render(_size)  # noqa: F841
+            canv = widget.render(_size)
             # print(b"\n".join(canv.text).decode())
+            screen.draw_screen(_size, canv)
 
-            assert len(janitor._ti_image_cviews) == 6
-            assert {(row, col) for _, row, col, *_ in janitor._ti_image_cviews} == {
+            assert len(screen._ti_image_cviews) == 6
+            assert {(row, col) for _, row, col, *_ in screen._ti_image_cviews} == {
                 (2, 2),
                 (2, 17),
                 (3, 2),
@@ -1153,17 +1189,18 @@ def test_move_top_widget(self):
                 (9, 17),
             }
 
-            prev_image_cviews = janitor._ti_image_cviews
+            prev_image_cviews = screen._ti_image_cviews
             widget.top += 1
             widget._invalidate()
-            canv = janitor.render(_size)  # noqa: F841
+            canv = widget.render(_size)
             # print(b"\n".join(canv.text).decode())
+            screen.draw_screen(_size, canv)
 
-            assert len(janitor._ti_image_cviews) == 6
-            assert prev_image_cviews != janitor._ti_image_cviews
+            assert len(screen._ti_image_cviews) == 6
+            assert prev_image_cviews != screen._ti_image_cviews
             # The two lower cviews should be unchanged
-            assert len(prev_image_cviews & janitor._ti_image_cviews) == 2
-            assert {(row, col) for _, row, col, *_ in janitor._ti_image_cviews} == {
+            assert len(prev_image_cviews & screen._ti_image_cviews) == 2
+            assert {(row, col) for _, row, col, *_ in screen._ti_image_cviews} == {
                 (2, 2),
                 (2, 17),
                 (4, 2),
@@ -1172,17 +1209,18 @@ def test_move_top_widget(self):
                 (9, 17),
             }
 
-            prev_image_cviews = janitor._ti_image_cviews
+            prev_image_cviews = screen._ti_image_cviews
             widget.top = 5
             widget._invalidate()
-            canv = janitor.render(_size)  # noqa: F841
+            canv = widget.render(_size)
             # print(b"\n".join(canv.text).decode())
+            screen.draw_screen(_size, canv)
 
-            assert len(janitor._ti_image_cviews) == 8
-            assert prev_image_cviews != janitor._ti_image_cviews
+            assert len(screen._ti_image_cviews) == 8
+            assert prev_image_cviews != screen._ti_image_cviews
             # All cviews should be changed
-            assert len(prev_image_cviews & janitor._ti_image_cviews) == 0
-            assert {(row, col) for _, row, col, *_ in janitor._ti_image_cviews} == {
+            assert len(prev_image_cviews & screen._ti_image_cviews) == 0
+            assert {(row, col) for _, row, col, *_ in screen._ti_image_cviews} == {
                 (2, 2),
                 (2, 17),
                 (6, 2),
@@ -1193,17 +1231,18 @@ def test_move_top_widget(self):
                 (11, 17),
             }
 
-            prev_image_cviews = janitor._ti_image_cviews
+            prev_image_cviews = screen._ti_image_cviews
             widget.top = 8
             widget._invalidate()
-            canv = janitor.render(_size)  # noqa: F841
+            canv = widget.render(_size)
             # print(b"\n".join(canv.text).decode())
+            screen.draw_screen(_size, canv)
 
-            assert len(janitor._ti_image_cviews) == 6
-            assert prev_image_cviews != janitor._ti_image_cviews
+            assert len(screen._ti_image_cviews) == 6
+            assert prev_image_cviews != screen._ti_image_cviews
             # All cviews should be changed
-            assert len(prev_image_cviews & janitor._ti_image_cviews) == 0
-            assert {(row, col) for _, row, col, *_ in janitor._ti_image_cviews} == {
+            assert len(prev_image_cviews & screen._ti_image_cviews) == 0
+            assert {(row, col) for _, row, col, *_ in screen._ti_image_cviews} == {
                 (2, 2),
                 (2, 17),
                 (9, 2),
@@ -1222,24 +1261,26 @@ def test_scroll_listboxes(self):
         ITerm2Image._TERM = "konsole"
         try:
             # Setup
-            janitor = UrwidImageJanitor(widget)
+            buf.seek(0)
+            buf.truncate()
             widget.top_w = top_w
             widget.top = 5
             kitty_list_box.shift_focus(_size, 0)
             iterm2_list_box.shift_focus(_size, 0)
             widget._invalidate()
-            janitor.render(_size)
+            screen.draw_screen(_size, widget.render(_size))
 
-            prev_image_cviews = janitor._ti_image_cviews
+            prev_image_cviews = screen._ti_image_cviews
             kitty_list_box.shift_focus(_size, -3)
-            canv = janitor.render(_size)  # noqa: F841
+            canv = widget.render(_size)
             # print(b"\n".join(canv.text).decode())
+            screen.draw_screen(_size, canv)
 
-            assert len(janitor._ti_image_cviews) == 8
-            assert prev_image_cviews != janitor._ti_image_cviews
+            assert len(screen._ti_image_cviews) == 8
+            assert prev_image_cviews != screen._ti_image_cviews
             # All images (4 cviews) on the right should remain unchanged
-            assert len(prev_image_cviews & janitor._ti_image_cviews) == 4
-            assert {(row, col) for _, row, col, *_ in janitor._ti_image_cviews} == {
+            assert len(prev_image_cviews & screen._ti_image_cviews) == 4
+            assert {(row, col) for _, row, col, *_ in screen._ti_image_cviews} == {
                 (2, 2),
                 (2, 17),
                 (6, 2),
@@ -1251,7 +1292,7 @@ def test_scroll_listboxes(self):
             }
             assert {
                 (row, col)
-                for _, row, col, *_ in prev_image_cviews & janitor._ti_image_cviews
+                for _, row, col, *_ in prev_image_cviews & screen._ti_image_cviews
             } == {
                 (2, 17),
                 (6, 21),
@@ -1259,16 +1300,17 @@ def test_scroll_listboxes(self):
                 (11, 17),
             }
 
-            prev_image_cviews = janitor._ti_image_cviews
+            prev_image_cviews = screen._ti_image_cviews
             iterm2_list_box.shift_focus(_size, -9)
-            canv = janitor.render(_size)  # noqa: F841
+            canv = widget.render(_size)
             # print(b"\n".join(canv.text).decode())
+            screen.draw_screen(_size, canv)
 
-            assert len(janitor._ti_image_cviews) == 8
-            assert prev_image_cviews != janitor._ti_image_cviews
+            assert len(screen._ti_image_cviews) == 8
+            assert prev_image_cviews != screen._ti_image_cviews
             # All images (4 cviews) on the left should remain unchanged
-            assert len(prev_image_cviews & janitor._ti_image_cviews) == 4
-            assert {(row, col) for _, row, col, *_ in janitor._ti_image_cviews} == {
+            assert len(prev_image_cviews & screen._ti_image_cviews) == 4
+            assert {(row, col) for _, row, col, *_ in screen._ti_image_cviews} == {
                 (2, 2),
                 (2, 17),
                 (6, 2),
@@ -1280,7 +1322,7 @@ def test_scroll_listboxes(self):
             }
             assert {
                 (row, col)
-                for _, row, col, *_ in prev_image_cviews & janitor._ti_image_cviews
+                for _, row, col, *_ in prev_image_cviews & screen._ti_image_cviews
             } == {
                 (2, 2),
                 (6, 2),
@@ -1296,48 +1338,52 @@ def test_change_top_widget(self):
         ITerm2Image._TERM = "konsole"
         try:
             # Setup
-            janitor = UrwidImageJanitor(widget)
+            buf.seek(0)
+            buf.truncate()
             widget.top_w = top_w
             widget.top = 5
             kitty_list_box.shift_focus(_size, 0)
             iterm2_list_box.shift_focus(_size, 0)
-            janitor.render(_size)
+            screen.draw_screen(_size, widget.render(_size))
 
-            prev_image_cviews = janitor._ti_image_cviews
+            # kitty image in the top widget
+            prev_image_cviews = screen._ti_image_cviews
             widget.top_w = kitty_image_w
             widget._invalidate()
-            canv = janitor.render(_size)  # noqa: F841
+            canv = widget.render(_size)
             # print(b"\n".join(canv.text).decode())
+            screen.draw_screen(_size, canv)
 
-            assert len(janitor._ti_image_cviews) == 9
-            assert prev_image_cviews != janitor._ti_image_cviews
-            # One new kitty image in the top widget
-            assert len(janitor._ti_image_cviews - prev_image_cviews) == 1
-            assert (*(janitor._ti_image_cviews - prev_image_cviews),)[0][1:3] == (6, 11)
+            assert len(screen._ti_image_cviews) == 9
+            assert prev_image_cviews != screen._ti_image_cviews
+            assert len(screen._ti_image_cviews - prev_image_cviews) == 1
+            assert (*(screen._ti_image_cviews - prev_image_cviews),)[0][1:3] == (6, 11)
 
-            prev_image_cviews = janitor._ti_image_cviews
+            # block image in the top widget
+            prev_image_cviews = screen._ti_image_cviews
             widget.top_w = block_image_w
             widget._invalidate()
-            canv = janitor.render(_size)  # noqa: F841
+            canv = widget.render(_size)
             # print(b"\n".join(canv.text).decode())
+            screen.draw_screen(_size, canv)
 
-            assert len(janitor._ti_image_cviews) == 8
-            assert prev_image_cviews != janitor._ti_image_cviews
-            # The image in the top widget is now of block style
-            assert len(prev_image_cviews - janitor._ti_image_cviews) == 1
-            assert (*(prev_image_cviews - janitor._ti_image_cviews),)[0][1:3] == (6, 11)
+            assert len(screen._ti_image_cviews) == 8
+            assert prev_image_cviews != screen._ti_image_cviews
+            assert len(prev_image_cviews - screen._ti_image_cviews) == 1
+            assert (*(prev_image_cviews - screen._ti_image_cviews),)[0][1:3] == (6, 11)
 
-            prev_image_cviews = janitor._ti_image_cviews
+            # iterm2 image in the top widget
+            prev_image_cviews = screen._ti_image_cviews
             widget.top_w = iterm2_image_w
             widget._invalidate()
-            canv = janitor.render(_size)  # noqa: F841
+            canv = widget.render(_size)
             # print(b"\n".join(canv.text).decode())
+            screen.draw_screen(_size, canv)
 
-            assert len(janitor._ti_image_cviews) == 9
-            assert prev_image_cviews != janitor._ti_image_cviews
-            # The image in the top widget is now of iterm2 style
-            assert len(janitor._ti_image_cviews - prev_image_cviews) == 1
-            assert (*(janitor._ti_image_cviews - prev_image_cviews),)[0][1:3] == (6, 11)
+            assert len(screen._ti_image_cviews) == 9
+            assert prev_image_cviews != screen._ti_image_cviews
+            assert len(screen._ti_image_cviews - prev_image_cviews) == 1
+            assert (*(screen._ti_image_cviews - prev_image_cviews),)[0][1:3] == (6, 11)
 
         finally:
             ITerm2Image._TERM = _TERM
@@ -1347,19 +1393,21 @@ def test_iterm2_not_on_konsole(self):
         ITerm2Image._TERM = "wezterm"
         try:
             # Setup
-            janitor = UrwidImageJanitor(widget)
+            buf.seek(0)
+            buf.truncate()
             widget.top_w = top_w
             widget.top = 5
             kitty_list_box.shift_focus(_size, 0)
             iterm2_list_box.shift_focus(_size, 0)
 
             widget._invalidate()
-            canv = janitor.render(_size)  # noqa: F841
+            canv = widget.render(_size)  # noqa: F841
             # print(b"\n".join(canv.text).decode())
+            screen.draw_screen(_size, canv)
 
             # The images on the right are of iterm2 style but not on konsole
-            assert len(janitor._ti_image_cviews) == 4
-            assert {(row, col) for _, row, col, *_ in janitor._ti_image_cviews} == {
+            assert len(screen._ti_image_cviews) == 4
+            assert {(row, col) for _, row, col, *_ in screen._ti_image_cviews} == {
                 (2, 2),
                 (6, 2),
                 (9, 2),
@@ -1368,39 +1416,3 @@ def test_iterm2_not_on_konsole(self):
 
         finally:
             ITerm2Image._TERM = _TERM
-
-
-class TestScreen:
-    def test_start(self):
-        screen = UrwidImageScreen(input=sys.__stdin__)
-        buf = io.StringIO()
-        screen.write = buf.write
-
-        screen.start()
-        start_output = buf.getvalue()
-        buf.seek(0)
-        buf.truncate()
-        screen.stop()
-        stop_output = buf.getvalue()
-
-        assert start_output.endswith(kitty.DELETE_ALL_IMAGES)
-        assert stop_output.startswith(kitty.DELETE_ALL_IMAGES)
-
-    def test_not_supported(self):
-        screen = UrwidImageScreen(input=sys.__stdin__)
-        buf = io.StringIO()
-        screen.write = buf.write
-
-        KittyImage._supported = False
-        try:
-            screen.start()
-            start_output = buf.getvalue()
-            buf.seek(0)
-            buf.truncate()
-            screen.stop()
-            stop_output = buf.getvalue()
-
-            assert kitty.DELETE_ALL_IMAGES not in start_output
-            assert kitty.DELETE_ALL_IMAGES not in stop_output
-        finally:
-            KittyImage._supported = True