From 38dfd747f68017a988f9c3e7340cb97fd032b5d3 Mon Sep 17 00:00:00 2001
From: Ivan Habunek <ivan@habunek.com>
Date: Sun, 12 Jan 2025 12:00:59 +0100
Subject: [PATCH] Add batched helper

---
 tests/test_utils.py    | 14 +++++++++++++-
 toot/utils/__init__.py | 19 ++++++++++++++++++-
 2 files changed, 31 insertions(+), 2 deletions(-)

diff --git a/tests/test_utils.py b/tests/test_utils.py
index bb5586e..6508f57 100644
--- a/tests/test_utils.py
+++ b/tests/test_utils.py
@@ -7,7 +7,7 @@ from toot.wcstring import wc_wrap, trunc, pad, fit_text
 from toot.tui.utils import LRUCache
 from PIL import Image
 from collections import namedtuple
-from toot.utils import urlencode_url
+from toot.utils import batched, urlencode_url
 
 
 def test_pad():
@@ -319,3 +319,15 @@ def test_urlencode_url():
     assert urlencode_url("https://www.example.com") == "https://www.example.com"
     assert urlencode_url("https://www.example.com/url%20with%20spaces") == "https://www.example.com/url%20with%20spaces"
 
+
+def test_batched():
+    assert list(batched("", 2)) == []
+    assert list(batched("a", 2)) == [["a"]]
+    assert list(batched("ab", 2)) == [["a", "b"]]
+    assert list(batched("abc", 2)) == [["a", "b"], ["c"]]
+    assert list(batched("abcd", 2)) == [["a", "b"], ["c", "d"]]
+    assert list(batched("abcde", 2)) == [["a", "b"], ["c", "d"], ["e"]]
+    assert list(batched("abcdef", 2)) == [["a", "b"], ["c", "d"], ["e", "f"]]
+
+    with pytest.raises(ValueError):
+        list(batched("foo", 0))
diff --git a/toot/utils/__init__.py b/toot/utils/__init__.py
index 6d2ae9f..2bfa3c7 100644
--- a/toot/utils/__init__.py
+++ b/toot/utils/__init__.py
@@ -9,7 +9,8 @@ import warnings
 
 from bs4 import BeautifulSoup
 from importlib.metadata import version
-from typing import Any, Dict, Generator, List, Optional
+from itertools import islice
+from typing import Any, Dict, Generator, Iterable, List, Optional, TypeVar
 from urllib.parse import urlparse, urlencode, quote, unquote
 
 
@@ -164,3 +165,19 @@ def get_version(name):
         return version(name)
     except Exception:
         return None
+
+
+T = TypeVar("T")
+
+def batched(iterable: Iterable[T], n: int) -> Generator[List[T], None, None]:
+    """Batch data from the iterable into lists of length n. The last batch may
+    be shorter than n."""
+    if n < 1:
+        raise ValueError("n must be positive")
+    iterator = iter(iterable)
+    while True:
+        batch = list(islice(iterator, n))
+        if batch:
+            yield batch
+        else:
+            break