feat: add CacheSettings.only_my_member (hikari-py#1679)

Co-authored-by: davfsa <[email protected]>

changes/1679.feature.md

@@ -0,0 +1 @@
+Add `CacheSettings.only_my_member` to only cache the bot member.

hikari/api/config.py

@@ -77,7 +77,7 @@ class CacheComponents(enums.Flag):
     """Enables the guild stickers cache."""
 
     GUILD_THREADS = 1 << 12
-    """Enabled the guild threads cache."""
+    """Enables the guild threads cache."""
 
     ALL = (
         GUILDS
@@ -199,3 +199,14 @@ class CacheSettings(abc.ABC):
     @abc.abstractmethod
     def components(self) -> CacheComponents:
         """Cache components to use."""
+
+    @property
+    @abc.abstractmethod
+    def only_my_member(self) -> bool:
+        """Reduce the members cache to only the bot itself.
+
+        Useful when only the bot member is required (eg. permission checks).
+        This will have no effect if the members cache is not enabled.
+
+        Defaults to `False`.
+        """

hikari/impl/config.py

@@ -29,6 +29,7 @@
     "ProxySettings",
     "HTTPTimeoutSettings",
     "HTTPSettings",
+    "CacheComponents",
     "CacheSettings",
 )
 
@@ -292,6 +293,10 @@ class that will **NOT** enforce SSL verification. This is then stored
     """
 
 
+# Re-export
+CacheComponents = config.CacheComponents
+
+
 @attrs_extensions.with_copy
 @attrs.define(kw_only=True, weakref_slot=False)
 class CacheSettings(config.CacheSettings):
@@ -320,3 +325,12 @@ class CacheSettings(config.CacheSettings):
 
     Defaults to `50`.
     """
+
+    only_my_member: bool = attrs.field(default=False)
+    """Reduce the members cache to only the bot itself.
+
+    Useful when only the bot member is required (eg. permission checks).
+    This will have no effect if the members cache is not enabled.
+
+    Defaults to `False`.
+    """

hikari/impl/event_manager.py

@@ -343,8 +343,12 @@ async def on_guild_create(  # noqa: C901, CFQ001 - Function too complex and too
             if members:
                 # TODO: do we really want to invalidate these all after an outage.
                 self._cache.clear_members_for_guild(guild_id)
-                for member in members.values():
-                    self._cache.set_member(member)
+                if not self._cache.settings.only_my_member:
+                    for member in members.values():
+                        self._cache.set_member(member)
+                else:
+                    my_member = members[shard.get_user_id()]
+                    self._cache.set_member(my_member)
 
             if presences:
                 self._cache.clear_presences_for_guild(guild_id)
@@ -361,17 +365,34 @@ async def on_guild_create(  # noqa: C901, CFQ001 - Function too complex and too
                 for thread in threads.values():
                     self._cache.set_thread(thread)
 
+        # We only want to chunk if we are allowed and need to:
+        #   Allowed?
+        #       All the following must be true:
+        #           1. `auto_chunk_members` is true (the user wants us to).
+        #           2. We have the necessary intents (`GUILD_MEMBERS`).
+        #           3. The guild is marked as "large" or we do not have `GUILD_PRESENCES` intent
+        #              Discord will only send every other member objects on the `GUILD_CREATE`
+        #              payload if presence intents are also declared, so if this isn't the case then we also
+        #              want to chunk small guilds.
+        #
+        #   Need to?
+        #       One of the following must be true:
+        #           1. We have a cache, and it requires it (it is enabled for `MEMBERS`), but we are
+        #              not limited to only our own member (which is included in the `GUILD_CREATE` payload).
+        #           2. The user is waiting for the member chunks (there is an event listener for it).
         presences_declared = self._intents & intents_.Intents.GUILD_PRESENCES
 
-        # When intents are enabled Discord will only send other member objects on the guild create
-        # payload if presence intents are also declared, so if this isn't the case then we also want
-        # to chunk small guilds.
         if (
             self._auto_chunk_members
             and self._intents & intents_.Intents.GUILD_MEMBERS
             and (payload.get("large") or not presences_declared)
             and (
-                self._cache_enabled_for(config.CacheComponents.MEMBERS)
+                (
+                    self._cache
+                    and self._cache_enabled_for(config.CacheComponents.MEMBERS)
+                    and not self._cache.settings.only_my_member
+                )
+                # This call is a bit expensive, so best to do it last
                 or self._enabled_for_event(shard_events.MemberChunkEvent)
             )
         ):

hikari/impl/gateway_bot.py

@@ -183,15 +183,22 @@ class GatewayBot(traits.GatewayBotAware):
         Defaults to `True`. If `False`, then no member chunks
         will be requested automatically, even if there are reasons to do so.
 
-        All following statements must be true to automatically request chunks:
-
-        1. `auto_chunk_members` is `True`.
-        2. The members intent is enabled.
-        3. The server is marked as "large" or the presences intent is not enabled
-           (since Discord only sends other members when presences are declared,
-           we should also chunk small guilds if the presences are not declared).
-        4. The members cache is enabled or there are listeners for the
-           `MemberChunkEvent`.
+        We only want to chunk if we are allowed and need to:
+
+        - Allowed?
+            All the following must be true:
+                1. `auto_chunk_members` is true (the user wants us to).
+                2. We have the necessary intents (`GUILD_MEMBERS`).
+                3. The guild is marked as "large" or we do not have `GUILD_PRESENCES` intent
+                   Discord will only send every other member objects on the `GUILD_CREATE`
+                   payload if presence intents are also declared, so if this isn't the case then we also
+                   want to chunk small guilds.
+
+        - Needed?
+            One of the following must be true:
+                1. We have a cache, and it requires it (it is enabled for `MEMBERS`), but we are
+                   not limited to only our own member (which is included in the `GUILD_CREATE` payload).
+                2. The user is waiting for the member chunks (there is an event listener for it).
     logs : typing.Union[None, str, int, typing.Dict[str, typing.Any], os.PathLike]
         The flavour to set the logging to.
 

tests/hikari/impl/test_event_manager.py

@@ -551,15 +551,19 @@ async def test_on_guild_create_when_not_dispatching_and_not_caching(
 
         event_manager_impl.dispatch.assert_not_called()
 
-    @pytest.mark.parametrize("include_unavailable", [True, False])
+    @pytest.mark.parametrize(
+        ("include_unavailable", "only_my_member"), [(True, True), (True, False), (False, True), (False, False)]
+    )
     @pytest.mark.asyncio()
     async def test_on_guild_create_when_not_dispatching_and_caching(
-        self, event_manager_impl, shard, event_factory, entity_factory, include_unavailable
+        self, event_manager_impl, shard, event_factory, entity_factory, include_unavailable, only_my_member
     ):
         payload = {"unavailable": False} if include_unavailable else {}
         event_manager_impl._intents = intents.Intents.NONE
         event_manager_impl._cache_enabled_for = mock.Mock(return_value=True)
         event_manager_impl._enabled_for_event = mock.Mock(return_value=False)
+        event_manager_impl._cache.settings.only_my_member = only_my_member
+        shard.get_user_id.return_value = 1
         gateway_guild = entity_factory.deserialize_gateway_guild.return_value
         gateway_guild.channels.return_value = {1: "channel1", 2: "channel2"}
         gateway_guild.emojis.return_value = {1: "emoji1", 2: "emoji2"}
@@ -595,13 +599,17 @@ async def test_on_guild_create_when_not_dispatching_and_caching(
         event_manager_impl._cache.clear_roles_for_guild.assert_called_once_with(gateway_guild.id)
         event_manager_impl._cache.set_role.assert_has_calls([mock.call("role1"), mock.call("role2")])
         event_manager_impl._cache.clear_members_for_guild.assert_called_once_with(gateway_guild.id)
-        event_manager_impl._cache.set_member.assert_has_calls([mock.call("member1"), mock.call("member2")])
+        if only_my_member:
+            event_manager_impl._cache.set_member.assert_called_once_with("member1")
+            shard.get_user_id.assert_has_calls([mock.call(), mock.call()])
+        else:
+            event_manager_impl._cache.set_member.assert_has_calls([mock.call("member1"), mock.call("member2")])
+            shard.get_user_id.assert_called_once_with()
         event_manager_impl._cache.clear_presences_for_guild.assert_called_once_with(gateway_guild.id)
         event_manager_impl._cache.set_presence.assert_has_calls([mock.call("presence1"), mock.call("presence2")])
         event_manager_impl._cache.clear_voice_states_for_guild.assert_called_once_with(gateway_guild.id)
         event_manager_impl._cache.set_voice_state.assert_has_calls([mock.call("voice1"), mock.call("voice2")])
         request_guild_members.assert_not_called()
-        shard.get_user_id.assert_called_once_with()
 
         event_manager_impl.dispatch.assert_not_called()
 
@@ -633,25 +641,58 @@ async def test_on_guild_create_when_stateless(
         stateless_event_manager_impl.dispatch.assert_not_called()
 
     @pytest.mark.asyncio()
-    async def test_on_guild_create_when_members_declared_and_member_cache_enabled(
-        self, stateless_event_manager_impl, shard, event_factory, entity_factory
+    async def test_on_guild_create_when_members_declared_and_member_cache_enabled_but_only_my_member_not_enabled(
+        self, event_manager_impl, shard, event_factory, entity_factory
     ):
+        def cache_enabled_for_members_only(component):
+            return component == config.CacheComponents.MEMBERS
+
         shard.id = 123
-        stateless_event_manager_impl._intents = intents.Intents.GUILD_MEMBERS
-        stateless_event_manager_impl._cache_enabled_for = mock.Mock(return_value=True)
-        stateless_event_manager_impl._enabled_for_event = mock.Mock(return_value=False)
+        event_manager_impl._cache.settings.only_my_member = False
+        event_manager_impl._intents = intents.Intents.GUILD_MEMBERS
+        event_manager_impl._cache_enabled_for = cache_enabled_for_members_only
+        event_manager_impl._enabled_for_event = mock.Mock(return_value=False)
+        gateway_guild = entity_factory.deserialize_gateway_guild.return_value
+        gateway_guild.id = 456
+        gateway_guild.members.return_value = {1: "member1", 2: "member2"}
         mock_request_guild_members = mock.Mock()
 
         with mock.patch.object(asyncio, "create_task") as create_task:
             with mock.patch.object(event_manager, "_fixed_size_nonce", return_value="abc"):
                 with mock.patch.object(event_manager, "_request_guild_members", new=mock_request_guild_members):
-                    await stateless_event_manager_impl.on_guild_create(shard, {"id": 456, "large": False})
+                    await event_manager_impl.on_guild_create(shard, {"id": 456, "large": False})
 
         mock_request_guild_members.assert_called_once_with(shard, 456, include_presences=False, nonce="123.abc")
         create_task.assert_called_once_with(
             mock_request_guild_members.return_value, name="123:456 guild create members request"
         )
 
+    @pytest.mark.asyncio()
+    async def test_on_guild_create_when_members_declared_and_member_cache_but_only_my_member_enabled(
+        self, event_manager_impl, shard, event_factory, entity_factory
+    ):
+        def cache_enabled_for_members_only(component):
+            return component == config.CacheComponents.MEMBERS
+
+        shard.id = 123
+        shard.get_user_id.return_value = 1
+        event_manager_impl._cache.settings.only_my_member = True
+        event_manager_impl._intents = intents.Intents.GUILD_MEMBERS
+        event_manager_impl._cache_enabled_for = cache_enabled_for_members_only
+        event_manager_impl._enabled_for_event = mock.Mock(return_value=False)
+        gateway_guild = entity_factory.deserialize_gateway_guild.return_value
+        gateway_guild.members.return_value = {1: "member1", 2: "member2"}
+
+        mock_request_guild_members = mock.Mock()
+
+        with mock.patch.object(asyncio, "create_task") as create_task:
+            with mock.patch.object(event_manager, "_fixed_size_nonce", return_value="abc"):
+                with mock.patch.object(event_manager, "_request_guild_members", new=mock_request_guild_members):
+                    await event_manager_impl.on_guild_create(shard, {"id": 456, "large": False})
+
+        mock_request_guild_members.assert_not_called()
+        create_task.assert_not_called()
+
     @pytest.mark.asyncio()
     async def test_on_guild_create_when_members_declared_and_enabled_for_member_chunk_event(
         self, stateless_event_manager_impl, shard, event_factory, entity_factory