Optimized iterator.prefetch_all()

tests/test_executor.py

@@ -1,7 +1,6 @@
 # flake8: noqa: F403, F405
 
 import time
-import types
 from dataclasses import asdict
 
 import pytest
@@ -108,7 +107,7 @@ def test_prefetch_off():
 
 
 def test_prefetch_on():
-    data = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
+    data = list(range(10 ** 4))
     client = mk_mock_client({"pag_on": [data, 200]}, paginated=["pag_on"])
     executor = Executor(client)
     p_pag = Path("pag_on", [int], pagination=True)
@@ -158,3 +157,13 @@ def test_take():
         next(it)
 
     assert it.take(2) == it.take_all() == []
+
+
+def test_chaining():
+    data = list(range(300))
+    client = mk_mock_client({"pag_on": [data, 200]}, paginated=["pag_on"])
+    executor = Executor(client)
+    p_pag = Path("pag_on", [int], pagination=True)
+
+    assert executor.run(path=p_pag, per_page=2).take_all() == data
+    assert executor.run(path=p_pag, per_page=2).prefetch_all().take_all() == data

trakt/core/executors.py

@@ -117,6 +117,7 @@ def find_matching_path(self) -> List[Path]:
 
 
 T = TypeVar("T")
+PER_PAGE_LIMIT = 100
 
 
 class PaginationIterator(Iterable[T]):
@@ -138,6 +139,7 @@ def __init__(
 
         self._exhausted = False
         self._queue: List[T] = []
+        self._yielded_items = 0
 
     def __iter__(self) -> PaginationIterator[T]:
         if self._exhausted:
@@ -148,6 +150,7 @@ def __iter__(self) -> PaginationIterator[T]:
         self._stop_at_page = self._max_pages
 
         self._queue = []
+        self._yielded_items = 0
 
         return self
 
@@ -158,32 +161,43 @@ def __next__(self) -> T:
 
             self._fetch_next_page()
 
+        self._yielded_items += 1
         return self._queue.pop(0)
 
-    def _fetch_next_page(self) -> None:
+    def _fetch_next_page(self, skip_first: int = 0) -> None:
         response, pagination = self._executor.exec_path_call(
             self._path,
             pagination=True,
             extra_quargs={"page": str(self._page), "limit": str(self._per_page)},
         )
 
-        for r in response:
+        for r in response[skip_first:]:
             self._queue.append(r)
 
         self._page += 1
         self._stop_at_page = int(pagination["page_count"])
         self.pages_total = self._stop_at_page
 
     def prefetch_all(self) -> PaginationIterator[T]:
-        """Prefetch all results. Faster than normal iteration."""
+        """Prefetch all results. Optimized."""
         iterator = cast(PaginationIterator[T], iter(self))
 
-        # TODO assert changing per_page doesn't break iteration
-        # old_per_page = self._per_page
-        # self._per_page = 100
+        if not self._has_next_page():
+            return iterator
+
+        # tweak per_page setting to make fetching as fast as possible
+        old_per_page = self._per_page
+        self._per_page = PER_PAGE_LIMIT
+
+        self._page = (self._yielded_items // PER_PAGE_LIMIT) + 1
+        to_skip = (self._yielded_items % PER_PAGE_LIMIT) + len(self._queue)
+
+        self._fetch_next_page(skip_first=to_skip)
+
         while self._has_next_page():
             self._fetch_next_page()
-        # self._per_page = old_per_page
+
+        self._per_page = old_per_page
 
         return iterator