def _default_include_presences( self, guild_id: snowflakes.Snowflake, include_presences: undefined.UndefinedOr[bool]) -> bool: if include_presences is not undefined.UNDEFINED: return include_presences shard_id = snowflakes.calculate_shard_id(self._app, guild_id) shard = self._app.shards[shard_id] return bool(shard.intents & intents_.Intents.GUILD_PRESENCES)
def shard(self) -> shard_.GatewayShard: if self.message.guild_id is not None: shard_id = snowflakes.calculate_shard_id(self.shard_service, self.message.guild_id) else: shard_id = 0 return self.shard_service.shards[shard_id]
async def list_requests_for_guild( self, guild: snowflakes.SnowflakeishOr[guilds.GatewayGuild], /) -> typing.Sequence[chunker.RequestInformation]: guild_id = snowflakes.Snowflake(guild) shard_id = snowflakes.calculate_shard_id(self._app, guild_id) if shard_id not in self._tracked: return () return tuple( copy.copy(event) for event in self._tracked[shard_id].values() if event.guild_id == guild_id)
def shard_id(self) -> typing.Optional[int]: """Return the ID of the shard this guild is served by. This may return `None` if the application does not have a gateway connection. """ try: # This is only sensible if there is a shard. shard_count = getattr(self.app, "shard_count") assert isinstance(shard_count, int) return snowflakes.calculate_shard_id(shard_count, self.id) except (TypeError, AttributeError, NameError): return None
def shard_id(self) -> typing.Optional[int]: """Return the shard ID for the shard. This may be `builtins.None` if the shard count is not known. """ try: shard_count = getattr(self.app, "shard_count") assert isinstance(shard_count, int) return snowflakes.calculate_shard_id(shard_count, self.guild_id) except (TypeError, AttributeError, NameError): pass return None
async def open(self) -> None: # as super().open will possibly flip self.started, we want to have its state when this was originally called. started = self._active await super().open() if not started: shard_id = snowflakes.calculate_shard_id(self._app, self._guild_id) await self._app.shards[shard_id].request_guild_members( guild=self._guild_id, include_presences=self._include_presences, query=self._query, limit=self._limit, users=self._users, nonce=self._nonce, )
async def about_command( ctx: tanjun.abc.Context, process: alluka.Injected[psutil.Process], ) -> None: """Get basic information about the current bot instance.""" start_date = datetime.datetime.fromtimestamp(process.create_time()) uptime = datetime.datetime.now() - start_date memory_usage: float = process.memory_full_info().uss / 1024**2 cpu_usage: float = process.cpu_percent() / psutil.cpu_count() memory_percent: float = process.memory_percent() if ctx.shards: shard_id = snowflakes.calculate_shard_id( ctx.shards.shard_count, ctx.guild_id) if ctx.guild_id else 0 name = f"Reinhard: Shard {shard_id} of {ctx.shards.shard_count}" else: name = "Reinhard: REST Server" description = ( "An experimental pythonic Hikari bot.\n " "The source can be found on [Github](https://github.com/FasterSpeeding/Reinhard)." ) embed = (hikari.Embed( description=description, colour=utility.embed_colour() ).set_author(name=name, url=hikari.__url__).add_field( name="Uptime", value=str(uptime), inline=True ).add_field( name="Process", value= f"{memory_usage:.2f} MB ({memory_percent:.0f}%)\n{cpu_usage:.2f}% CPU", inline=True, ).set_footer( icon="http://i.imgur.com/5BFecvA.png", text= f"Made with Hikari v{hikari.__version__} (python {platform.python_version()})", )) await ctx.respond(embed=embed, component=utility.delete_row(ctx))
class StatefulGuildChunkerImpl(chunker.GuildChunker): """Guild chunker implementation. Parameters ---------- app : hikari.traits.BotAware The object of the bot aware app this is bound to. limit : builtins.int The maximum amount of requests that this chunker should store information about for each shard. """ __slots__: typing.Sequence[str] = ("_app", "_limit", "_tracked") def __init__(self, app: traits.BotAware, limit: int = 200) -> None: self._app = app self._limit = limit self._tracked: typing.MutableMapping[int, mapping.CMRIMutableMapping[ str, _TrackedRequests]] = {} def _default_include_presences( self, guild_id: snowflakes.Snowflake, include_presences: undefined.UndefinedOr[bool]) -> bool: if include_presences is not undefined.UNDEFINED: return include_presences shard_id = snowflakes.calculate_shard_id(self._app, guild_id) shard = self._app.shards[shard_id] return shard.intents is None or bool( shard.intents & intents_.Intents.GUILD_PRESENCES) def fetch_members_for_guild( self, guild: snowflakes.SnowflakeishOr[guilds.GatewayGuild], *, timeout: typing.Union[int, float, None], limit: typing.Optional[int] = None, include_presences: undefined.UndefinedOr[bool] = undefined.UNDEFINED, query_limit: int = 0, query: str = "", users: undefined.UndefinedOr[typing.Sequence[snowflakes.SnowflakeishOr[ users_.User]]] = undefined.UNDEFINED, ) -> event_stream.Streamer[shard_events.MemberChunkEvent]: guild_id = snowflakes.Snowflake(guild) return ChunkStream( app=self._app, guild_id=guild_id, timeout=timeout, limit=limit, include_presences=self._default_include_presences( guild_id, include_presences), query_limit=query_limit, query=query, users=users, ) async def get_request_status( self, nonce: str, /) -> typing.Optional[chunker.RequestInformation]: try: shard_id = int(nonce.split(".", 1)[0]) except ValueError: return None else: return copy.copy(self._tracked[shard_id].get( nonce)) if shard_id in self._tracked else None async def list_requests_for_shard( self, shard: typing.Union[gateway_shard.GatewayShard, int], /) -> typing.Sequence[chunker.RequestInformation]: shard_id = shard if isinstance(shard, int) else shard.id if shard_id not in self._tracked: return () return tuple( copy.copy(chunk) for chunk in self._tracked[shard_id].copy().values()) async def list_requests_for_guild( self, guild: snowflakes.SnowflakeishOr[guilds.GatewayGuild], /) -> typing.Sequence[chunker.RequestInformation]: guild_id = snowflakes.Snowflake(guild) shard_id = snowflakes.calculate_shard_id(self._app, guild_id) if shard_id not in self._tracked: return () return tuple( copy.copy(event) for event in self._tracked[shard_id].values() if event.guild_id == guild_id) async def consume_chunk_event(self, event: shard_events.MemberChunkEvent, /) -> None: if (event.shard.id not in self._tracked or event.nonce is None or event.nonce not in self._tracked[event.shard.id]): return self._tracked[event.shard.id][event.nonce].update(event) async def request_guild_members( self, guild: snowflakes.SnowflakeishOr[guilds.GatewayGuild], /, include_presences: undefined.UndefinedOr[bool] = undefined.UNDEFINED, limit: int = 0, query: str = "", users: undefined.UndefinedOr[typing.Sequence[snowflakes.SnowflakeishOr[ users_.User]]] = undefined.UNDEFINED, ) -> str: guild_id = snowflakes.Snowflake(guild) shard_id = snowflakes.calculate_shard_id(self._app, guild_id) nonce = f"{shard_id}.{_random_nonce()}" if shard_id not in self._tracked: self._tracked[shard_id] = mapping.CMRIMutableMapping( limit=self._limit) tracker = _TrackedRequests(guild_id=guild_id, nonce=nonce) self._tracked[shard_id][nonce] = tracker await self._app.shards[shard_id].request_guild_members( guild=guild_id, include_presences=self._default_include_presences( guild_id, include_presences), limit=limit, nonce=nonce, query=query, users=users, ) return nonce
def test_calculate_shard_id_with_app(guild_id, expected_id): mock_app = mock.Mock(bot.BotApp, shard_count=8) assert snowflakes.calculate_shard_id(mock_app, guild_id) == expected_id
def test_calculate_shard_id_with_shard_count(guild_id, expected_id): assert snowflakes.calculate_shard_id(4, guild_id) == expected_id
async def connect_to( self, guild: snowflakes.SnowflakeishOr[guilds.PartialGuild], channel: snowflakes.SnowflakeishOr[channels.GuildVoiceChannel], voice_connection_type: typing.Type[_VoiceConnectionT], *, deaf: bool = False, mute: bool = False, **kwargs: typing.Any, ) -> _VoiceConnectionT: self._check_if_alive() guild_id = snowflakes.Snowflake(guild) if guild_id in self._connections: raise errors.VoiceError( "Already in a voice channel for that guild. Disconnect before attempting to connect again" ) shard_id = snowflakes.calculate_shard_id(self._app, guild_id) try: shard = self._app.shards[shard_id] except KeyError: raise errors.VoiceError( f"Cannot connect to shard {shard_id} as it is not present in this application" ) from None user = self._app.cache.get_me() if not user: user = await self._app.rest.fetch_my_user() _LOGGER.log( ux.TRACE, "attempting to connect to voice channel %s in %s via shard %s", channel, guild, shard_id) await shard.update_voice_state(guild, channel, self_deaf=deaf, self_mute=mute) _LOGGER.log( ux.TRACE, "waiting for voice events for connecting to voice channel %s in %s via shard %s", channel, guild, shard_id, ) state_event, server_event = await asyncio.gather( # Voice state update: self._app.event_manager.wait_for( voice_events.VoiceStateUpdateEvent, timeout=None, predicate=self._init_state_update_predicate(guild_id, user.id), ), # Server update: self._app.event_manager.wait_for( voice_events.VoiceServerUpdateEvent, timeout=None, predicate=self._init_server_update_predicate(guild_id), ), ) # We will never receive the first endpoint as `None` assert server_event.endpoint is not None _LOGGER.debug( "joined voice channel %s in guild %s via shard %s using endpoint %s. Session will be %s. " "Delegating to voice websocket", state_event.state.channel_id, state_event.state.guild_id, shard_id, server_event.endpoint, state_event.state.session_id, ) try: voice_connection = await voice_connection_type.initialize( channel_id=snowflakes.Snowflake(channel), endpoint=server_event.endpoint, guild_id=guild_id, on_close=self._on_connection_close, owner=self, session_id=state_event.state.session_id, shard_id=shard_id, token=server_event.token, user_id=user.id, **kwargs, ) except Exception: _LOGGER.debug( "error occurred in initialization, leaving voice channel %s in guild %s", channel, guild) try: await asyncio.wait_for(shard.update_voice_state(guild, None), timeout=5.0) except asyncio.TimeoutError: pass raise self._connections[guild_id] = voice_connection return voice_connection
async def connect_to( self, guild: snowflakes.SnowflakeishOr[guilds.PartialGuild], channel: snowflakes.SnowflakeishOr[channels.GuildVoiceChannel], voice_connection_type: typing.Type[_VoiceConnectionT], *, deaf: bool = False, mute: bool = False, **kwargs: typing.Any, ) -> _VoiceConnectionT: guild_id = snowflakes.Snowflake(guild) shard_id = snowflakes.calculate_shard_id(self._app, guild_id) if shard_id is None: raise errors.VoiceError( "Cannot connect to voice. Ensure the application is configured as a gateway zookeeper and try again." ) if guild_id in self._connections: raise errors.VoiceError( "The bot is already in a voice channel for this guild. Close the other connection first, or " "request that the application moves to the new voice channel instead." ) try: shard = self._app.shards[shard_id] except KeyError: raise errors.VoiceError( f"Cannot connect to shard {shard_id}, it is not present in this application." ) from None if not shard.is_alive: # Not sure if I can think of a situation this will happen in... really. # Unless the user sleeps for a bit then tries to connect, and in the mean time the # shard has disconnected. # TODO: make shards declare if they are in the process of reconnecting, if they are, make them wait # for a little bit. raise errors.VoiceError(f"Cannot connect to shard {shard_id}, the shard is not online.") _LOGGER.log(ux.TRACE, "attempting to connect to voice channel %s in %s via shard %s", channel, guild, shard_id) user = self._app.cache.get_me() if user is None: user = await self._app.rest.fetch_my_user() await asyncio.wait_for(shard.update_voice_state(guild, channel, self_deaf=deaf, self_mute=mute), timeout=5.0) _LOGGER.log( ux.TRACE, "waiting for voice events for connecting to voice channel %s in %s via shard %s", channel, guild, shard_id, ) state_event, server_event = await asyncio.wait_for( asyncio.gather( # Voice state update: self._events.wait_for( voice_events.VoiceStateUpdateEvent, timeout=None, predicate=self._init_state_update_predicate(guild_id, user.id), ), # Server update: self._events.wait_for( voice_events.VoiceServerUpdateEvent, timeout=None, predicate=self._init_server_update_predicate(guild_id), ), ), timeout=10.0, ) _LOGGER.debug( "joined voice channel %s in guild %s via shard %s using endpoint %s. Session will be %s. " "Delegating to voice websocket", state_event.state.channel_id, state_event.state.guild_id, shard_id, server_event.endpoint, state_event.state.session_id, ) try: voice_connection = await voice_connection_type.initialize( channel_id=snowflakes.Snowflake(channel), endpoint=server_event.endpoint, guild_id=guild_id, on_close=self._on_connection_close, owner=self, session_id=state_event.state.session_id, shard_id=shard_id, token=server_event.token, user_id=user.id, **kwargs, ) except Exception: _LOGGER.debug( "error occurred in initialization, leaving voice channel %s in guild %s again", channel, guild ) await asyncio.wait_for(shard.update_voice_state(guild, None), timeout=5.0) raise self._connections[guild_id] = voice_connection return voice_connection
def _get_shard(self, guild: snowflakes.SnowflakeishOr[guilds.PartialGuild]) -> gateway_shard.GatewayShard: guild = snowflakes.Snowflake(guild) if shard := self._shards.get(snowflakes.calculate_shard_id(self.shard_count, guild)): return shard