Improve image fetching in ThreadABC

examples/fetch.py

@@ -1,4 +1,3 @@
-import itertools
 import fbchat
 
 session = fbchat.Session.login("<email>", "<password>")
@@ -62,7 +61,9 @@ print("thread's name: {}".format(thread.name))
 # Here should be an example of `getUnread`
 
 
-# Print image url for 20 last images from thread.
+# Print image url for up to 20 last images from thread.
-images = thread.fetch_images()
+images = list(thread.fetch_images(limit=20))
-for image in itertools.islice(image, 20):
+for image in images:
-    print(image.large_preview_url)
+    if isinstance(image, fbchat.ImageAttachment):
+        url = c.fetch_image_url(image.id)
+        print(url)

fbchat/_file.py

@@ -91,8 +91,6 @@ class ImageAttachment(Attachment):
 
     @classmethod
     def _from_list(cls, data):
-        data = data["node"]
-
         previews = {
             Image._from_uri_or_none(data["image"]),
             Image._from_uri(data["image1"]),
@@ -156,7 +154,6 @@ class VideoAttachment(Attachment):
 
     @classmethod
     def _from_list(cls, data):
-        data = data["node"]
         previews = {
             Image._from_uri(data["image"]),
             Image._from_uri(data["image1"]),

fbchat/_thread.py

@@ -338,35 +338,48 @@ class ThreadABC(metaclass=abc.ABCMeta):
 
         return messages
 
-    def fetch_images(self):
+    def _fetch_images(self, limit, after):
-        """Fetch images/videos posted in the thread."""
+        data = {"id": self.id, "first": limit, "after": after}
-        # TODO: Return proper searchable iterator
-        data = {"id": self.id, "first": 48}
         (j,) = self.session._graphql_requests(
             _graphql.from_query_id("515216185516880", data)
         )
-        while True:
-            try:
-                i = j[self.id]["message_shared_media"]["edges"][0]
-            except IndexError:
-                if j[self.id]["message_shared_media"]["page_info"].get("has_next_page"):
-                    data["after"] = j[self.id]["message_shared_media"]["page_info"].get(
-                        "end_cursor"
-                    )
-                    (j,) = self.session._graphql_requests(
-                        _graphql.from_query_id("515216185516880", data)
-                    )
-                    continue
-                else:
-                    break
 
-            if i["node"].get("__typename") == "MessageImage":
+        result = j[self.id]["message_shared_media"]
-                yield _file.ImageAttachment._from_list(i)
+
-            elif i["node"].get("__typename") == "MessageVideo":
+        print(len(result["edges"]))
-                yield _file.VideoAttachment._from_list(i)
+
+        rtn = []
+        for edge in result["edges"]:
+            node = edge["node"]
+            type_ = node["__typename"]
+            if type_ == "MessageImage":
+                rtn.append(_file.ImageAttachment._from_list(node))
+            elif type_ == "MessageVideo":
+                rtn.append(_file.VideoAttachment._from_list(node))
             else:
-                yield _attachment.Attachment(id=i["node"].get("legacy_attachment_id"))
+                log.warning("Unknown image type %s, data: %s", type_, edge)
-            del j[self.id]["message_shared_media"]["edges"][0]
+                rtn.append(None)
+
+        # result["page_info"]["has_next_page"] is not correct when limit > 12
+        return (result["page_info"]["end_cursor"], rtn)
+
+    def fetch_images(self, limit: int) -> Iterable[_attachment.Attachment]:
+        """Fetch images/videos posted in the thread.
+
+        Args:
+            limit: Max. number of images to retrieve. If ``None``, all images will be
+                retrieved.
+        """
+        cursor = None
+        # The max limit on this request is unknown, so we set it reasonably high
+        # This way `limit=None` also still works
+        for limit in _util.get_limits(limit, max_limit=1000):
+            cursor, images = self._fetch_images(limit, cursor)
+            if not images:
+                return  # No more data to fetch
+            for image in images:
+                if image:
+                    yield image
 
     def set_nickname(self, user_id: str, nickname: str):
         """Change the nickname of a user in the thread.

tests/test_file.py

@@ -46,7 +46,7 @@ def test_imageattachment_from_list():
                 height=988,
             ),
         },
-    ) == ImageAttachment._from_list({"node": data})
+    ) == ImageAttachment._from_list(data)
 
 
 def test_videoattachment_from_list():
@@ -88,7 +88,7 @@ def test_videoattachment_from_list():
                 height=368,
             ),
         },
-    ) == VideoAttachment._from_list({"node": data})
+    ) == VideoAttachment._from_list(data)
 
 
 def test_graphql_to_attachment_empty():