Return iterators for paginated endpoints

rocketchat_API/APISections/base.py

@@ -1,7 +1,11 @@
 import re
+
+from functools import wraps
+
 from json import JSONDecodeError
 from typing import Any
 
+
 import requests
 
 from rocketchat_API.APIExceptions.RocketExceptions import (
@@ -11,6 +15,59 @@
 )
 
 
+def paginated(data_key):
+    """
+    Decorator that converts a paginated API method into an iterator.
+
+    Args:
+        data_key: The key in the API response that contains the list of items
+                  (e.g., 'groups', 'channels', 'users')
+
+    Returns:
+        A decorator that wraps the original method to yield items one by one,
+        automatically handling pagination with offset and count parameters.
+
+    Example:
+        @paginated('groups')
+        def groups_list_all(self, **kwargs):
+            return self.call_api_get("groups.listAll", kwargs=kwargs)
+    """
+
+    def decorator(func):
+        def _generator(self, first_data, offset, count, **kwargs):
+            """Inner generator that yields items from paginated API responses."""
+            data = first_data
+            while True:
+                items = data.get(data_key, [])
+                if not items:
+                    break
+
+                for item in items:
+                    yield item
+
+                # If we got fewer items than requested, we've reached the end
+                if len(items) < count:
+                    break
+
+                offset += count
+                # Call the original function with pagination parameters
+                data = func(self, offset=offset, count=count, **kwargs)
+
+        @wraps(func)
+        def wrapper(self, **kwargs):
+            offset = kwargs.pop("offset", 0)
+            count = kwargs.pop("count", 50)
+
+            # Call the original function eagerly to propagate any exceptions
+            first_data = func(self, offset=offset, count=count, **kwargs)
+
+            return _generator(self, first_data, offset, count, **kwargs)
+
+        return wrapper
+
+    return decorator
+
+
 def json_or_error(r: requests.Response) -> Any:
     if r.status_code > 399 or r.status_code < 200:
         raise RocketBadStatusCodeException(r.status_code, r.text)

rocketchat_API/APISections/channels.py

@@ -1,14 +1,16 @@
 from rocketchat_API.APIExceptions.RocketExceptions import RocketMissingParamException
-from rocketchat_API.APISections.base import RocketChatBase
+from rocketchat_API.APISections.base import RocketChatBase, paginated
 
 
 class RocketChatChannels(RocketChatBase):
+    @paginated("channels")
     def channels_list(self, **kwargs):
         """Retrieves all of the channels from the server."""
         return self.call_api_get("channels.list", kwargs=kwargs)
 
+    @paginated("channels")
     def channels_list_joined(self, **kwargs):
-        """Lists all of the channels the calling user has joined"""
+        """Lists all of the channels the calling user has joined."""
         return self.call_api_get("channels.list.joined", kwargs=kwargs)
 
     def channels_info(self, room_id=None, channel=None, **kwargs):
@@ -19,6 +21,7 @@ def channels_info(self, room_id=None, channel=None, **kwargs):
             return self.call_api_get("channels.info", roomName=channel, kwargs=kwargs)
         raise RocketMissingParamException("room_id or channel required")
 
+    @paginated("messages")
     def channels_history(self, room_id, **kwargs):
         """Retrieves the messages from a channel."""
         return self.call_api_get("channels.history", roomId=room_id, kwargs=kwargs)
@@ -205,6 +208,7 @@ def channels_delete(self, room_id=None, channel=None, **kwargs):
             )
         raise RocketMissingParamException("room_id or channel required")
 
+    @paginated("members")
     def channels_members(self, room_id=None, channel=None, **kwargs):
         """Lists all channel users."""
         if room_id:
@@ -225,6 +229,7 @@ def channels_roles(self, room_id=None, room_name=None, **kwargs):
             )
         raise RocketMissingParamException("room_id or room_name required")
 
+    @paginated("files")
     def channels_files(self, room_id=None, room_name=None, **kwargs):
         """Retrieves the files from a channel."""
         if room_id:
@@ -235,6 +240,7 @@ def channels_files(self, room_id=None, room_name=None, **kwargs):
             )
         raise RocketMissingParamException("room_id or room_name required")
 
+    @paginated("mentions")
     def channels_get_all_user_mentions_by_channel(self, room_id, **kwargs):
         """Gets all the mentions of a channel."""
         return self.call_api_get(
@@ -251,10 +257,7 @@ def channels_counters(self, room_id=None, room_name=None, **kwargs):
             )
         raise RocketMissingParamException("room_id or room_name required")
 
-    def channels_online(self, _id=None):
+    def channels_online(self, _id):
         """Lists all online users of a channel if the channel's id is provided, otherwise it gets all online users of
         all channels."""
-        if _id:
-            return self.call_api_get("channels.online", _id=_id)
-        else:
-            raise RocketMissingParamException("_id or query required")
+        return self.call_api_get("channels.online", _id=_id)

rocketchat_API/APISections/chat.py

@@ -1,5 +1,5 @@
 from rocketchat_API.APIExceptions.RocketExceptions import RocketMissingParamException
-from rocketchat_API.APISections.base import RocketChatBase
+from rocketchat_API.APISections.base import RocketChatBase, paginated
 
 
 class RocketChatChat(RocketChatBase):
@@ -59,6 +59,7 @@ def chat_react(self, msg_id, emoji="smile", **kwargs):
             "chat.react", messageId=msg_id, emoji=emoji, kwargs=kwargs
         )
 
+    @paginated("messages")
     def chat_search(self, room_id, search_text, **kwargs):
         """Search for messages in a channel by id and text message."""
         return self.call_api_get(
@@ -71,6 +72,7 @@ def chat_get_message_read_receipts(self, message_id, **kwargs):
             "chat.getMessageReadReceipts", messageId=message_id, kwargs=kwargs
         )
 
+    @paginated("messages")
     def chat_get_starred_messages(self, room_id, **kwargs):
         """Retrieve starred messages."""
         return self.call_api_get(
@@ -90,6 +92,7 @@ def chat_follow_message(self, mid, **kwargs):
         """Follows a chat message to the message's channel."""
         return self.call_api_post("chat.followMessage", mid=mid, kwargs=kwargs)
 
+    @paginated("messages")
     def chat_get_thread_messages(self, thread_msg_id, **kwargs):
         """Get thread messages."""
         return self.call_api_get(
@@ -98,6 +101,7 @@ def chat_get_thread_messages(self, thread_msg_id, **kwargs):
             kwargs=kwargs,
         )
 
+    @paginated("messages")
     def chat_get_mentioned_messages(self, room_id, **kwargs):
         """Get the messages in which you are mentioned (users are mentioned with the @ symbol)."""
         return self.call_api_get(

rocketchat_API/APISections/dm.py

@@ -1,18 +1,21 @@
 from rocketchat_API.APIExceptions.RocketExceptions import RocketMissingParamException
-from rocketchat_API.APISections.base import RocketChatBase
+from rocketchat_API.APISections.base import RocketChatBase, paginated
 
 
 class RocketChatDM(RocketChatBase):
+    @paginated("ims")
     def dm_list(self, **kwargs):
-        """List the private im chats for logged user"""
+        """List the private im chats for logged user."""
         return self.call_api_get("dm.list", kwargs=kwargs)
 
+    @paginated("ims")
     def dm_list_everyone(self, **kwargs):
         """List all direct message the caller in the server."""
         return self.call_api_get("dm.list.everyone", kwargs=kwargs)
 
+    @paginated("messages")
     def dm_history(self, room_id, **kwargs):
-        """Retrieves the history for a private im chat"""
+        """Retrieves the history for a private im chat."""
         return self.call_api_get("dm.history", roomId=room_id, kwargs=kwargs)
 
     def dm_create(self, username, **kwargs):
@@ -37,8 +40,9 @@ def dm_members(self, room_id):
         """Retrieves members of a direct message."""
         return self.call_api_get("dm.members", roomId=room_id)
 
+    @paginated("messages")
     def dm_messages(self, room_id=None, username=None, **kwargs):
-        """Retrieves direct messages from the server by username"""
+        """Retrieves direct messages from the server."""
         if room_id:
             return self.call_api_get("dm.messages", roomId=room_id, kwargs=kwargs)
 
@@ -57,6 +61,7 @@ def dm_set_topic(self, room_id, topic, **kwargs):
             "dm.setTopic", roomId=room_id, topic=topic, kwargs=kwargs
         )
 
+    @paginated("files")
     def dm_files(self, room_id=None, user_name=None, **kwargs):
         """Retrieves the files from a direct message."""
         if room_id:

rocketchat_API/APISections/groups.py

@@ -1,19 +1,22 @@
 from rocketchat_API.APIExceptions.RocketExceptions import RocketMissingParamException
-from rocketchat_API.APISections.base import RocketChatBase
+from rocketchat_API.APISections.base import RocketChatBase, paginated
 
 
 class RocketChatGroups(RocketChatBase):
+    @paginated("groups")
     def groups_list_all(self, **kwargs):
         """
         List all the private groups on the server.
         The calling user must have the 'view-room-administration' right
         """
         return self.call_api_get("groups.listAll", kwargs=kwargs)
 
+    @paginated("groups")
     def groups_list(self, **kwargs):
         """List the private groups the caller is part of."""
         return self.call_api_get("groups.list", kwargs=kwargs)
 
+    @paginated("messages")
     def groups_history(self, room_id, **kwargs):
         """Retrieves the messages from a private group."""
         return self.call_api_get("groups.history", roomId=room_id, kwargs=kwargs)
@@ -172,6 +175,7 @@ def groups_delete(self, room_id=None, group=None, **kwargs):
             return self.call_api_post("groups.delete", roomName=group, kwargs=kwargs)
         raise RocketMissingParamException("room_id or group required")
 
+    @paginated("members")
     def groups_members(self, room_id=None, group=None, **kwargs):
         """Lists all group users."""
         if room_id:
@@ -188,6 +192,7 @@ def groups_roles(self, room_id=None, room_name=None, **kwargs):
             return self.call_api_get("groups.roles", roomName=room_name, kwargs=kwargs)
         raise RocketMissingParamException("room_id or room_name required")
 
+    @paginated("files")
     def groups_files(self, room_id=None, room_name=None, **kwargs):
         """Retrieves the files from a private group."""
         if room_id:

rocketchat_API/APISections/integrations.py

@@ -1,7 +1,7 @@
 from rocketchat_API.APIExceptions.RocketExceptions import (
     RocketUnsuportedIntegrationType,
 )
-from rocketchat_API.APISections.base import RocketChatBase
+from rocketchat_API.APISections.base import RocketChatBase, paginated
 
 
 class RocketChatIntegrations(RocketChatBase):
@@ -52,12 +52,14 @@ def integrations_get(self, integration_id, **kwargs):
             "integrations.get", integrationId=integration_id, kwargs=kwargs
         )
 
+    @paginated("history")
     def integrations_history(self, integration_id, **kwargs):
         """Lists all history of the specified integration."""
         return self.call_api_get(
             "integrations.history", id=integration_id, kwargs=kwargs
         )
 
+    @paginated("integrations")
     def integrations_list(self, **kwargs):
         """Lists all of the integrations on the server."""
         return self.call_api_get("integrations.list", kwargs=kwargs)

rocketchat_API/APISections/livechat.py

@@ -1,11 +1,13 @@
-from rocketchat_API.APISections.base import RocketChatBase
+from rocketchat_API.APISections.base import RocketChatBase, paginated
 
 
 class RocketChatLivechat(RocketChatBase):
+    @paginated("rooms")
     def livechat_rooms(self, **kwargs):
         """Retrieves a list of livechat rooms."""
         return self.call_api_get("livechat/rooms", kwargs=kwargs)
 
+    @paginated("inquiries")
     def livechat_inquiries_list(self, **kwargs):
         """Lists all of the open livechat inquiries."""
         return self.call_api_get("livechat/inquiries.list", kwargs=kwargs)
@@ -16,6 +18,7 @@ def livechat_inquiries_take(self, inquiry_id, **kwargs):
             "livechat/inquiries.take", inquiryId=inquiry_id, kwargs=kwargs
         )
 
+    @paginated("users")
     def livechat_get_users(self, user_type, **kwargs):
         """Get a list of agents or managers."""
         return self.call_api_get("livechat/users/{}".format(user_type), kwargs=kwargs)

rocketchat_API/APISections/roles.py

@@ -1,4 +1,4 @@
-from rocketchat_API.APISections.base import RocketChatBase
+from rocketchat_API.APISections.base import RocketChatBase, paginated
 
 
 class RocketChatRoles(RocketChatBase):
@@ -25,8 +25,9 @@ def roles_remove_user_from_role(self, role_name, username, **kwargs):
             kwargs=kwargs,
         )
 
+    @paginated("users")
     def roles_get_users_in_role(self, role, **kwargs):
-        """Gets the users that belongs to a role. It supports the Offset and Count Only."""
+        """Gets the users that belongs to a role."""
         return self.call_api_get("roles.getUsersInRole", role=role, kwargs=kwargs)
 
     def roles_sync(self, updated_since):

rocketchat_API/APISections/rooms.py

@@ -2,7 +2,7 @@
 import os
 
 from rocketchat_API.APIExceptions.RocketExceptions import RocketMissingParamException
-from rocketchat_API.APISections.base import RocketChatBase
+from rocketchat_API.APISections.base import RocketChatBase, paginated
 
 
 class RocketChatRooms(RocketChatBase):
@@ -53,6 +53,7 @@ def rooms_info(self, room_id=None, room_name=None):
             return self.call_api_get("rooms.info", roomName=room_name)
         raise RocketMissingParamException("room_id or roomName required")
 
+    @paginated("rooms")
     def rooms_admin_rooms(self, **kwargs):
         """Retrieves all rooms (requires the view-room-administration permission)."""
         return self.call_api_get("rooms.adminRooms", kwargs=kwargs)

rocketchat_API/APISections/teams.py

@@ -1,5 +1,5 @@
 from rocketchat_API.APIExceptions.RocketExceptions import RocketMissingParamException
-from rocketchat_API.APISections.base import RocketChatBase
+from rocketchat_API.APISections.base import RocketChatBase, paginated
 
 ID_OR_TEAM_NAME_REQUIRED = "team_id or team_name required"
 
@@ -19,6 +19,7 @@ def teams_delete(self, team_id=None, team_name=None, **kwargs):
             return self.call_api_post("teams.delete", teamName=team_name, kwargs=kwargs)
         raise RocketMissingParamException(ID_OR_TEAM_NAME_REQUIRED)
 
+    @paginated("teams")
     def teams_list_all(self, **kwargs):
         """
         List all the teams on the server.
@@ -38,6 +39,7 @@ def teams_info(self, team_id=None, team_name=None, **kwargs):
             return self.call_api_get("teams.info", teamName=team_name, kwargs=kwargs)
         raise RocketMissingParamException(ID_OR_TEAM_NAME_REQUIRED)
 
+    @paginated("members")
     def teams_members(
         self, team_id=None, team_name=None, name="", username="", **kwargs
     ):
@@ -96,6 +98,7 @@ def teams_update_member(self, team_id=None, team_name=None, member=None, **kwarg
             )
         raise RocketMissingParamException(ID_OR_TEAM_NAME_REQUIRED)
 
+    @paginated("rooms")
     def teams_list_rooms(
         self, team_id=None, team_name=None, room_type="", name="", **kwargs
     ):

rocketchat_API/APISections/users.py

@@ -2,7 +2,7 @@
 import os
 
 from rocketchat_API.APIExceptions.RocketExceptions import RocketMissingParamException
-from rocketchat_API.APISections.base import RocketChatBase
+from rocketchat_API.APISections.base import RocketChatBase, paginated
 
 
 class RocketChatUsers(RocketChatBase):
@@ -18,6 +18,7 @@ def users_info(self, user_id=None, username=None, **kwargs):
             return self.call_api_get("users.info", username=username, kwargs=kwargs)
         raise RocketMissingParamException("userID or username required")
 
+    @paginated("users")
     def users_list(self, **kwargs):
         """All of the users and their information, limited to permissions."""
         return self.call_api_get("users.list", kwargs=kwargs)