Images are now cached in memory (10MB) with a LRU eviction policy

tests/test_utils.py

@@ -1,8 +1,11 @@
 from argparse import ArgumentTypeError
 import pytest
-
+import sys
 from toot.console import duration
 from toot.wcstring import wc_wrap, trunc, pad, fit_text
+from toot.tui.utils import ImageCache
+from PIL import Image
+from collections import namedtuple
 
 
 def test_pad():
@@ -201,3 +204,108 @@ def test_duration():
 
     with pytest.raises(ArgumentTypeError):
         duration("banana")
+
+
+def test_cache_null():
+    """Null dict is null."""
+    cache = ImageCache(cache_max_bytes=1024)
+    assert cache.__len__() == 0
+
+
+Case = namedtuple("Case", ["cache_len", "len", "init"])
+
+img = Image.new('RGB', (100, 100))
+img_size = sys.getsizeof(img.tobytes())
+
+@pytest.mark.parametrize(
+    "case",
+    [
+        Case(9, 0, []),
+        Case(9, 1, [("one", img)]),
+        Case(9, 2, [("one", img), ("two", img)]),
+        Case(2, 2, [("one", img), ("two", img)]),
+        Case(1, 1, [("one", img), ("two", img)]),
+    ],
+)
+@pytest.mark.parametrize("method", ["assign", "init"])
+def test_cache_init(case, method):
+    """Check that the # of elements is right, given # given and cache_len."""
+    if method == "init":
+        cache = ImageCache(case.init, cache_max_bytes=img_size * case.cache_len)
+    elif method == "assign":
+        cache = ImageCache(cache_max_bytes=img_size * case.cache_len)
+        for (key, val) in case.init:
+            cache[key] = val
+    else:
+        assert False
+
+    # length is max(#entries, cache_len)
+    assert cache.__len__() == case.len
+
+    # make sure the first entry is the one ejected
+    if case.cache_len > 1 and case.init:
+        assert "one" in cache.keys()
+    else:
+        assert "one" not in cache.keys()
+
+
+@pytest.mark.parametrize("method", ["init", "assign"])
+def test_cache_overflow_default(method):
+    """Test default overflow logic."""
+    if method == "init":
+        cache = ImageCache([("one", img), ("two", img), ("three", img)], cache_max_bytes=img_size * 2)
+    elif method == "assign":
+        cache = ImageCache(cache_max_bytes=img_size * 2)
+        cache["one"] = img
+        cache["two"] = img
+        cache["three"] = img
+    else:
+        assert False
+
+    assert "one" not in cache.keys()
+    assert "two" in cache.keys()
+    assert "three" in cache.keys()
+
+@pytest.mark.parametrize("mode", ["get", "set"])
+@pytest.mark.parametrize("add_third", [False, True])
+def test_cache_lru_overflow(mode, add_third):
+    img = Image.new('RGB', (100, 100))
+    img_size = sys.getsizeof(img.tobytes())
+
+    """Test that key access resets LRU logic."""
+
+    cache = ImageCache([("one", img), ("two", img)], cache_max_bytes=img_size * 2)
+
+    if mode == "get":
+        dummy = cache["one"]
+    elif mode == "set":
+        cache["one"] = img
+    else:
+        assert False
+
+    if add_third:
+        cache["three"] = img
+
+        assert "one" in cache.keys()
+        assert "two" not in cache.keys()
+        assert "three" in cache.keys()
+    else:
+        assert "one" in cache.keys()
+        assert "two" in cache.keys()
+        assert "three" not in cache.keys()
+
+
+def test_cache_keyerror():
+    cache = ImageCache()
+    with pytest.raises(KeyError):
+        cache["foo"]
+
+
+def test_cache_miss_doesnt_eject():
+    cache = ImageCache([("one", img), ("two", img)], cache_max_bytes=img_size * 3)
+    with pytest.raises(KeyError):
+        cache["foo"]
+
+    assert len(cache) == 2
+    assert "one" in cache.keys()
+    assert "two" in cache.keys()

toot/tui/app.py

@@ -15,7 +15,7 @@ from .overlays import ExceptionStackTrace, GotoMenu, Help, StatusSource, StatusL
 from .overlays import StatusDeleteConfirmation, Account
 from .poll import Poll
 from .timeline import Timeline
-from .utils import get_max_toot_chars, parse_content_links, show_media, copy_to_clipboard
+from .utils import get_max_toot_chars, parse_content_links, show_media, copy_to_clipboard, ImageCache
 
 from PIL import Image
 from term_image.widget import UrwidImageScreen
@@ -663,13 +663,13 @@ class TUI(urwid.Frame):
                 return
 
             if not hasattr(timeline, "images"):
-                timeline.images = dict()
+                timeline.images = ImageCache()  # use the default 10MB image cache for now
             try:
                 img = Image.open(requests.get(path, stream=True).raw)
                 if img.format == 'PNG' and img.mode != 'RGBA':
                     img = img.convert("RGBA")
                 timeline.images[str(hash(path))] = img
-            except:  # noqa E722
+            except Exception:
                 pass  # ignore errors; if we can't load an image, just show blank
 
         def _done(loop):

toot/tui/timeline.py

@@ -314,8 +314,12 @@ class Timeline(urwid.Columns):
         assert self.statuses[index].id == status.id  # Sanity check
 
         # get the image and replace the placeholder with a graphics widget
+        img = None
         if hasattr(self, "images"):
-            img = self.images.get(str(hash(path)))
+            try:
+                img = self.images[(str(hash(path)))]
+            except KeyError:
+                pass
         if img:
             try:
                 render_img = add_corners(img, 10) if self.can_render_pixels else img
@@ -383,7 +387,10 @@ class StatusDetails(urwid.Pile):
 
         img = None
         if hasattr(self.timeline, "images"):
-            img = self.timeline.images.get(str(hash(path)))
+            try:
+                img = self.timeline.images[(str(hash(path)))]
+            except KeyError:
+                pass
         if img:
             render_img = add_corners(img, 10) if self.timeline.can_render_pixels else img
             return (urwid.BoxAdapter(

toot/tui/utils.py

@@ -2,9 +2,10 @@ import base64
 import re
 import shutil
 import subprocess
+import sys
 import urwid
 import math
-
+from collections import OrderedDict
 from functools import reduce
 from html.parser import HTMLParser
 from typing import List
@@ -180,3 +181,33 @@ def deep_get(adict: dict, path: List[str], default=None):
         path,
         adict
     )
+
+
+class ImageCache(OrderedDict):
+    """Dict with a limited size, ejecting LRUs as needed.
+        Default max size = 10Mb"""
+
+    def __init__(self, *args, cache_max_bytes: int = 1024 * 1024 * 10, **kwargs):
+        assert cache_max_bytes > 0
+        self.total_value_size = 0
+        self.cache_max_bytes = cache_max_bytes
+
+        super().__init__(*args, **kwargs)
+
+    def __setitem__(self, key: str, value: Image):
+        if key in self:
+            self.total_value_size -= sys.getsizeof(super().__getitem__(key).tobytes())
+        self.total_value_size += sys.getsizeof(value.tobytes())
+        super().__setitem__(key, value)
+        super().move_to_end(key)
+
+        while self.total_value_size > self.cache_max_bytes:
+            old_key, value = next(iter(self.items()))
+            sz = sys.getsizeof(value.tobytes())
+            super().__delitem__(old_key)
+            self.total_value_size -= sz
+
+    def __getitem__(self, key: str):
+        val = super().__getitem__(key)
+        super().move_to_end(key)
+        return val