Repository: hikari-py/hikari
Pull request #1119 (2022-04-01 01:11:39)

Implement pagination for fetch_bans method
### Summary
<!-- Small summary of the merge request -->
This pull request implements pagination for `RESTClientImpl.fetch_bans` method.

- It now returns a `LazyIterator[GuildBan]`.
- `GuildBanIterator` added to special_endpoints.

**This is a breaking change.**

### Checklist
<!-- Make sure to tick all the following boxes by putting an `x` in between (like this `[x]`) -->
- [x] I have run `nox` and all the pipelines have passed.
- [x] I have made unittests according to the code I have added/modified/deleted.



Patch:
diff --git a/changes/1119.bugfix.md b/changes/1119.bugfix.md
new file mode 100644
index 00000000..f3e4007e
--- /dev/null
+++ b/changes/1119.bugfix.md
@@ -0,0 +1,1 @@
+`RESTClientImpl.fetch_bans` now return a `LazyIterator` to allow pagination of values.
diff --git a/hikari/api/rest.py b/hikari/api/rest.py
index 55dcc855..062b21bf 100644
--- a/hikari/api/rest.py
+++ b/hikari/api/rest.py
@@ -5498,21 +5498,42 @@ async def fetch_ban(
         """
 
     @abc.abstractmethod
-    async def fetch_bans(
+    def fetch_bans(
         self,
         guild: snowflakes.SnowflakeishOr[guilds.PartialGuild],
-    ) -> typing.Sequence[guilds.GuildBan]:
+        /,
+        *,
+        newest_first: bool = False,
+        start_at: undefined.UndefinedOr[snowflakes.SearchableSnowflakeishOr[users.PartialUser]] = undefined.UNDEFINED,
+    ) -> iterators.LazyIterator[guilds.GuildBan]:
         """Fetch the bans of a guild.
 
+        !!! note
+            This call is not a coroutine function, it returns a special type of
+            lazy iterator that will perform API calls as you iterate across it.
+            See `hikari.iterators` for the full API for this iterator type.
+
         Parameters
         ----------
         guild : hikari.snowflakes.SnowflakeishOr[hikari.guilds.PartialGuild]
             The guild to fetch the bans from. This may be the
             object or the ID of an existing guild.
 
+        Other Parameters
+        ----------------
+        newest_first : builtins.bool
+            Whether to fetch the newest first or the oldest first.
+
+            Defaults to `builtins.False`.
+        start_at : undefined.UndefinedOr[snowflakes.SearchableSnowflakeishOr[users.PartialUser]]
+            If provided, will start at this snowflake. If you provide
+            a datetime object, it will be transformed into a snowflake. This
+            may also be a scheduled event object object. In this case, the
+            date the object was first created will be used.
+
         Returns
         -------
-        typing.Sequence[hikari.guilds.GuildBan]
+        hikari.iterators.LazyIterator[hikari.guilds.GuildBan]
             The requested bans.
 
         Raises
diff --git a/hikari/impl/rest.py b/hikari/impl/rest.py
index 926a3895..ec3af346 100644
--- a/hikari/impl/rest.py
+++ b/hikari/impl/rest.py
@@ -2765,13 +2765,24 @@ async def fetch_ban(
         assert isinstance(response, dict)
         return self._entity_factory.deserialize_guild_member_ban(response)
 
-    async def fetch_bans(
-        self, guild: snowflakes.SnowflakeishOr[guilds.PartialGuild]
-    ) -> typing.Sequence[guilds.GuildBan]:
-        route = routes.GET_GUILD_BANS.compile(guild=guild)
-        response = await self._request(route)
-        assert isinstance(response, list)
-        return [self._entity_factory.deserialize_guild_member_ban(ban_payload) for ban_payload in response]
+    def fetch_bans(
+        self,
+        guild: snowflakes.SnowflakeishOr[guilds.PartialGuild],
+        /,
+        *,
+        newest_first: bool = False,
+        start_at: undefined.UndefinedOr[snowflakes.SearchableSnowflakeishOr[users.PartialUser]] = undefined.UNDEFINED,
+    ) -> iterators.LazyIterator[guilds.GuildBan]:
+        if start_at is undefined.UNDEFINED:
+            start_at = snowflakes.Snowflake.max() if newest_first else snowflakes.Snowflake.min()
+        elif isinstance(start_at, datetime.datetime):
+            start_at = snowflakes.Snowflake.from_datetime(start_at)
+        else:
+            start_at = int(start_at)
+
+        return special_endpoints_impl.GuildBanIterator(
+            self._entity_factory, self._request, guild, newest_first, str(start_at)
+        )
 
     async def fetch_roles(
         self,
diff --git a/hikari/impl/special_endpoints.py b/hikari/impl/special_endpoints.py
index 772d1559..a2aa18f7 100644
--- a/hikari/impl/special_endpoints.py
+++ b/hikari/impl/special_endpoints.py
@@ -609,6 +609,58 @@ async def _next_chunk(self) -> typing.Optional[typing.Generator[applications.Own
         return (self._entity_factory.deserialize_own_guild(g) for g in chunk)
 
 
+# We use an explicit forward reference for this, since this breaks potential
+# circular import issues (once the file has executed, using those resources is
+# not an issue for us).
+class GuildBanIterator(iterators.BufferedLazyIterator["guilds.GuildBan"]):
+    """Iterator implementation for retrieving guild bans."""
+
+    __slots__: typing.Sequence[str] = (
+        "_entity_factory",
+        "_guild_id",
+        "_request_call",
+        "_route",
+        "_first_id",
+        "_newest_first",
+    )
+
+    def __init__(
+        self,
+        entity_factory: entity_factory_.EntityFactory,
+        request_call: typing.Callable[
+            ..., typing.Coroutine[None, None, typing.Union[None, data_binding.JSONObject, data_binding.JSONArray]]
+        ],
+        guild: snowflakes.SnowflakeishOr[guilds.PartialGuild],
+        newest_first: bool,
+        first_id: str,
+    ) -> None:
+        super().__init__()
+        self._guild_id = snowflakes.Snowflake(str(int(guild)))
+        self._route = routes.GET_GUILD_BANS.compile(guild=guild)
+        self._request_call = request_call
+        self._entity_factory = entity_factory
+        self._first_id = first_id
+        self._newest_first = newest_first
+
+    async def _next_chunk(self) -> typing.Optional[typing.Generator[guilds.GuildBan, typing.Any, None]]:
+        query = data_binding.StringMapBuilder()
+        query.put("before" if self._newest_first else "after", self._first_id)
+        query.put("limit", 1000)
+
+        chunk = await self._request_call(compiled_route=self._route, query=query)
+        assert isinstance(chunk, list)
+
+        if not chunk:
+            return None
+
+        if self._newest_first:
+            # These are always returned in ascending order by `.user.id`.
+            chunk.reverse()
+
+        self._first_id = chunk[-1]["user"]["id"]
+        return (self._entity_factory.deserialize_guild_member_ban(b) for b in chunk)
+
+
 # We use an explicit forward reference for this, since this breaks potential
 # circular import issues (once the file has executed, using those resources is
 # not an issue for us).


Test patch:
diff --git a/tests/hikari/impl/test_rest.py b/tests/hikari/impl/test_rest.py
index 9ad20649..21ef655a 100644
--- a/tests/hikari/impl/test_rest.py
+++ b/tests/hikari/impl/test_rest.py
@@ -1043,6 +1043,59 @@ def test_unban_member(self, rest_client):
         assert reason is mock_unban_user.return_value
         mock_unban_user.assert_called_once_with(123, 321, reason="ayaya")
 
+    def test_fetch_bans(self, rest_client: rest.RESTClientImpl):
+        with mock.patch.object(special_endpoints, "GuildBanIterator") as iterator_cls:
+            iterator = rest_client.fetch_bans(187, newest_first=True, start_at=StubModel(65652342134))
+
+        iterator_cls.assert_called_once_with(
+            rest_client._entity_factory,
+            rest_client._request,
+            187,
+            True,
+            "65652342134",
+        )
+        assert iterator is iterator_cls.return_value
+
+    def test_fetch_bans_when_datetime_for_start_at(self, rest_client: rest.RESTClientImpl):
+        start_at = datetime.datetime(2022, 3, 6, 12, 1, 58, 415625, tzinfo=datetime.timezone.utc)
+        with mock.patch.object(special_endpoints, "GuildBanIterator") as iterator_cls:
+            iterator = rest_client.fetch_bans(9000, newest_first=True, start_at=start_at)
+
+        iterator_cls.assert_called_once_with(
+            rest_client._entity_factory,
+            rest_client._request,
+            9000,
+            True,
+            "950000286338908160",
+        )
+        assert iterator is iterator_cls.return_value
+
+    def test_fetch_bans_when_start_at_undefined(self, rest_client: rest.RESTClientImpl):
+        with mock.patch.object(special_endpoints, "GuildBanIterator") as iterator_cls:
+            iterator = rest_client.fetch_bans(8844)
+
+        iterator_cls.assert_called_once_with(
+            rest_client._entity_factory,
+            rest_client._request,
+            8844,
+            False,
+            str(snowflakes.Snowflake.min()),
+        )
+        assert iterator is iterator_cls.return_value
+
+    def test_fetch_bans_when_start_at_undefined_and_newest_first(self, rest_client: rest.RESTClientImpl):
+        with mock.patch.object(special_endpoints, "GuildBanIterator") as iterator_cls:
+            iterator = rest_client.fetch_bans(3848, newest_first=True)
+
+        iterator_cls.assert_called_once_with(
+            rest_client._entity_factory,
+            rest_client._request,
+            3848,
+            True,
+            str(snowflakes.Snowflake.max()),
+        )
+        assert iterator is iterator_cls.return_value
+
     def test_command_builder(self, rest_client):
         with warnings.catch_warnings():
             warnings.simplefilter("ignore", category=DeprecationWarning)
@@ -4194,21 +4247,6 @@ async def test_fetch_ban(self, rest_client):
         rest_client._request.assert_awaited_once_with(expected_route)
         rest_client._entity_factory.deserialize_guild_member_ban.assert_called_once_with({"id": "789"})
 
-    async def test_fetch_bans(self, rest_client):
-        ban1 = StubModel(456)
-        ban2 = StubModel(789)
-        expected_route = routes.GET_GUILD_BANS.compile(guild=123)
-        rest_client._request = mock.AsyncMock(return_value=[{"id": "456"}, {"id": "789"}])
-        rest_client._entity_factory.deserialize_guild_member_ban = mock.Mock(side_effect=[ban1, ban2])
-
-        assert await rest_client.fetch_bans(StubModel(123)) == [ban1, ban2]
-
-        rest_client._request.assert_awaited_once_with(expected_route)
-        assert rest_client._entity_factory.deserialize_guild_member_ban.call_count == 2
-        rest_client._entity_factory.deserialize_guild_member_ban.assert_has_calls(
-            [mock.call({"id": "456"}), mock.call({"id": "789"})]
-        )
-
     async def test_fetch_roles(self, rest_client):
         role1 = StubModel(456)
         role2 = StubModel(789)
diff --git a/tests/hikari/impl/test_special_endpoints.py b/tests/hikari/impl/test_special_endpoints.py
index 9a482727..971dc1b3 100644
--- a/tests/hikari/impl/test_special_endpoints.py
+++ b/tests/hikari/impl/test_special_endpoints.py
@@ -165,6 +165,133 @@ async def test_aiter_when_empty_chunk(self, newest_first: bool):
         mock_request.assert_awaited_once_with(compiled_route=expected_route, query=query)
 
 
+class TestGuildBanIterator:
+    @pytest.mark.asyncio()
+    async def test_aiter(self):
+        expected_route = routes.GET_GUILD_BANS.compile(guild=10000)
+        mock_entity_factory = mock.Mock()
+        mock_payload_1 = {"user": {"id": "45234"}}
+        mock_payload_2 = {"user": {"id": "452745"}}
+        mock_payload_3 = {"user": {"id": "45237656"}}
+        mock_payload_4 = {"user": {"id": "452345666"}}
+        mock_payload_5 = {"user": {"id": "4523456744"}}
+        mock_result_1 = mock.Mock()
+        mock_result_2 = mock.Mock()
+        mock_result_3 = mock.Mock()
+        mock_result_4 = mock.Mock()
+        mock_result_5 = mock.Mock()
+        mock_entity_factory.deserialize_guild_member_ban.side_effect = [
+            mock_result_1,
+            mock_result_2,
+            mock_result_3,
+            mock_result_4,
+            mock_result_5,
+        ]
+        mock_request = mock.AsyncMock(
+            side_effect=[[mock_payload_1, mock_payload_2, mock_payload_3], [mock_payload_4, mock_payload_5], []]
+        )
+        iterator = special_endpoints.GuildBanIterator(
+            entity_factory=mock_entity_factory,
+            request_call=mock_request,
+            guild=10000,
+            newest_first=False,
+            first_id="0",
+        )
+
+        result = await iterator
+
+        assert result == [mock_result_1, mock_result_2, mock_result_3, mock_result_4, mock_result_5]
+        mock_entity_factory.deserialize_guild_member_ban.assert_has_calls(
+            [
+                mock.call(mock_payload_1),
+                mock.call(mock_payload_2),
+                mock.call(mock_payload_3),
+                mock.call(mock_payload_4),
+                mock.call(mock_payload_5),
+            ]
+        )
+        mock_request.assert_has_awaits(
+            [
+                mock.call(compiled_route=expected_route, query={"after": "0", "limit": "1000"}),
+                mock.call(compiled_route=expected_route, query={"after": "45237656", "limit": "1000"}),
+                mock.call(compiled_route=expected_route, query={"after": "4523456744", "limit": "1000"}),
+            ]
+        )
+
+    @pytest.mark.asyncio()
+    async def test_aiter_when_newest_first(self):
+        expected_route = routes.GET_GUILD_BANS.compile(guild=10000)
+        mock_entity_factory = mock.Mock()
+        mock_payload_1 = {"user": {"id": "432234"}}
+        mock_payload_2 = {"user": {"id": "1233211"}}
+        mock_payload_3 = {"user": {"id": "12332112"}}
+        mock_payload_4 = {"user": {"id": "1233"}}
+        mock_payload_5 = {"user": {"id": "54334"}}
+        mock_result_1 = mock.Mock()
+        mock_result_2 = mock.Mock()
+        mock_result_3 = mock.Mock()
+        mock_result_4 = mock.Mock()
+        mock_result_5 = mock.Mock()
+        mock_entity_factory.deserialize_guild_member_ban.side_effect = [
+            mock_result_1,
+            mock_result_2,
+            mock_result_3,
+            mock_result_4,
+            mock_result_5,
+        ]
+        mock_request = mock.AsyncMock(
+            side_effect=[[mock_payload_1, mock_payload_2, mock_payload_3], [mock_payload_4, mock_payload_5], []]
+        )
+        iterator = special_endpoints.GuildBanIterator(
+            entity_factory=mock_entity_factory,
+            request_call=mock_request,
+            guild=10000,
+            newest_first=True,
+            first_id="321123321",
+        )
+
+        result = await iterator
+
+        assert result == [mock_result_1, mock_result_2, mock_result_3, mock_result_4, mock_result_5]
+        mock_entity_factory.deserialize_guild_member_ban.assert_has_calls(
+            [
+                mock.call(mock_payload_3),
+                mock.call(mock_payload_2),
+                mock.call(mock_payload_1),
+                mock.call(mock_payload_5),
+                mock.call(mock_payload_4),
+            ]
+        )
+        mock_request.assert_has_awaits(
+            [
+                mock.call(compiled_route=expected_route, query={"before": "321123321", "limit": "1000"}),
+                mock.call(compiled_route=expected_route, query={"before": "432234", "limit": "1000"}),
+                mock.call(compiled_route=expected_route, query={"before": "1233", "limit": "1000"}),
+            ]
+        )
+
+    @pytest.mark.parametrize("newest_first", [True, False])
+    @pytest.mark.asyncio()
+    async def test_aiter_when_empty_chunk(self, newest_first: bool):
+        expected_route = routes.GET_GUILD_BANS.compile(guild=10000)
+        mock_entity_factory = mock.Mock()
+        mock_request = mock.AsyncMock(return_value=[])
+        iterator = special_endpoints.GuildBanIterator(
+            entity_factory=mock_entity_factory,
+            request_call=mock_request,
+            guild=10000,
+            newest_first=newest_first,
+            first_id="54234123123",
+        )
+
+        result = await iterator
+
+        assert result == []
+        mock_entity_factory.deserialize_guild_member_ban.assert_not_called()
+        query = {"before" if newest_first else "after": "54234123123", "limit": "1000"}
+        mock_request.assert_awaited_once_with(compiled_route=expected_route, query=query)
+
+
 class TestScheduledEventUserIterator:
     @pytest.mark.asyncio()
     async def test_aiter(self):