|
| 1 | +import enum |
| 2 | +import inspect |
| 3 | +from pathlib import Path |
| 4 | +from typing import NamedTuple, TYPE_CHECKING |
| 5 | + |
| 6 | +from discord import Embed |
| 7 | +from discord.ext import commands |
| 8 | +from discord.utils import escape_markdown |
| 9 | + |
| 10 | +if TYPE_CHECKING: |
| 11 | + from pydis_core import BotBase as Bot |
| 12 | + |
| 13 | + |
| 14 | +class TagIdentifierStub(NamedTuple): |
| 15 | + """A minmally functioning stub representing a tag identifier.""" |
| 16 | + |
| 17 | + group: str | None |
| 18 | + name: str |
| 19 | + |
| 20 | + @classmethod |
| 21 | + def from_string(cls, string: str) -> "TagIdentifierStub": |
| 22 | + """Create a TagIdentifierStub from a string.""" |
| 23 | + split_string = string.split(" ", maxsplit=2) |
| 24 | + if len(split_string) == 1: |
| 25 | + return cls(None, split_string[0]) |
| 26 | + return cls(split_string[0], split_string[1]) |
| 27 | + |
| 28 | + |
| 29 | +class SourceType(enum.StrEnum): |
| 30 | + """The types of source objects recognized by the source command.""" |
| 31 | + |
| 32 | + help_command = enum.auto() |
| 33 | + command = enum.auto() |
| 34 | + cog = enum.auto() |
| 35 | + tag = enum.auto() |
| 36 | + extension_not_loaded = enum.auto() |
| 37 | + |
| 38 | + |
| 39 | +class SourceCode(commands.Cog, description="Displays information about the bot's source code."): |
| 40 | + """ |
| 41 | + Pre-built cog to display source code links for commands and cogs (and if applicable, tags). |
| 42 | +
|
| 43 | + To use this cog, instantiate it with the bot instance and the base GitHub repository URL. |
| 44 | +
|
| 45 | + Args: |
| 46 | + bot (:obj:`pydis_core.BotBase`): The bot instance. |
| 47 | + github_repo: The base URL to the GitHub repository (e.g. `https://github.com/python-discord/bot`). |
| 48 | + """ |
| 49 | + |
| 50 | + def __init__(self, bot: "Bot", github_repo: str) -> None: |
| 51 | + self.bot = bot |
| 52 | + self.github_repo = github_repo |
| 53 | + |
| 54 | + @commands.command(name="source", aliases=("src",)) |
| 55 | + async def source_command( |
| 56 | + self, |
| 57 | + ctx: commands.Context, |
| 58 | + *, |
| 59 | + source_item: str | None = None, |
| 60 | + ) -> None: |
| 61 | + """Display information and a GitHub link to the source code of a command, tag, or cog.""" |
| 62 | + if not source_item: |
| 63 | + embed = Embed(title=f"{self.bot.user.name}'s GitHub Repository") |
| 64 | + embed.add_field(name="Repository", value=f"[Go to GitHub]({self.github_repo})") |
| 65 | + embed.set_thumbnail(url="https://avatars1.githubusercontent.com/u/9919") |
| 66 | + await ctx.send(embed=embed) |
| 67 | + return |
| 68 | + |
| 69 | + obj, source_type = await self._get_source_object(ctx, source_item) |
| 70 | + embed = await self._build_embed(obj, source_type) |
| 71 | + await ctx.send(embed=embed) |
| 72 | + |
| 73 | + @staticmethod |
| 74 | + async def _get_source_object(ctx: commands.Context, argument: str) -> tuple[object, SourceType]: |
| 75 | + """Convert argument into the source object and source type.""" |
| 76 | + if argument.lower() == "help": |
| 77 | + return ctx.bot.help_command, SourceType.help_command |
| 78 | + |
| 79 | + cog = ctx.bot.get_cog(argument) |
| 80 | + if cog: |
| 81 | + return cog, SourceType.cog |
| 82 | + |
| 83 | + cmd = ctx.bot.get_command(argument) |
| 84 | + if cmd: |
| 85 | + return cmd, SourceType.command |
| 86 | + |
| 87 | + tags_cog = ctx.bot.get_cog("Tags") |
| 88 | + show_tag = True |
| 89 | + |
| 90 | + if not tags_cog: |
| 91 | + show_tag = False |
| 92 | + else: |
| 93 | + identifier = TagIdentifierStub.from_string(argument.lower()) |
| 94 | + if identifier in tags_cog.tags: |
| 95 | + return identifier, SourceType.tag |
| 96 | + |
| 97 | + escaped_arg = escape_markdown(argument) |
| 98 | + |
| 99 | + raise commands.BadArgument( |
| 100 | + f"Unable to convert '{escaped_arg}' to valid command{', tag,' if show_tag else ''} or Cog." |
| 101 | + ) |
| 102 | + |
| 103 | + def _get_source_link(self, source_item: object, source_type: SourceType) -> tuple[str, str, int | None]: |
| 104 | + """ |
| 105 | + Build GitHub link of source item, return this link, file location and first line number. |
| 106 | +
|
| 107 | + Raise BadArgument if `source_item` is a dynamically-created object (e.g. via internal eval). |
| 108 | + """ |
| 109 | + if source_type == SourceType.command: |
| 110 | + source_item = inspect.unwrap(source_item.callback) |
| 111 | + src = source_item.__code__ |
| 112 | + filename = src.co_filename |
| 113 | + elif source_type == SourceType.tag: |
| 114 | + tags_cog = self.bot.get_cog("Tags") |
| 115 | + filename = tags_cog.tags[source_item].file_path |
| 116 | + else: |
| 117 | + src = type(source_item) |
| 118 | + try: |
| 119 | + filename = inspect.getsourcefile(src) |
| 120 | + except TypeError: |
| 121 | + raise commands.BadArgument("Cannot get source for a dynamically-created object.") |
| 122 | + |
| 123 | + if source_type != SourceType.tag: |
| 124 | + try: |
| 125 | + lines, first_line_no = inspect.getsourcelines(src) |
| 126 | + except OSError: |
| 127 | + raise commands.BadArgument("Cannot get source for a dynamically-created object.") |
| 128 | + |
| 129 | + lines_extension = f"#L{first_line_no}-L{first_line_no+len(lines)-1}" |
| 130 | + else: |
| 131 | + first_line_no = None |
| 132 | + lines_extension = "" |
| 133 | + |
| 134 | + # Handle tag file location differently than others to avoid errors in some cases |
| 135 | + if not first_line_no: |
| 136 | + file_location = Path(filename) |
| 137 | + else: |
| 138 | + file_location = Path(filename).relative_to(Path.cwd()).as_posix() |
| 139 | + |
| 140 | + url = f"{self.github_repo}/blob/main/{file_location}{lines_extension}" |
| 141 | + |
| 142 | + return url, file_location, first_line_no or None |
| 143 | + |
| 144 | + async def _build_embed(self, source_object: object, source_type: SourceType) -> Embed | None: |
| 145 | + """Build embed based on source object.""" |
| 146 | + url, location, first_line = self._get_source_link(source_object, source_type) |
| 147 | + |
| 148 | + if source_type == SourceType.help_command: |
| 149 | + title = "Help Command" |
| 150 | + description = source_object.__doc__.splitlines()[1] |
| 151 | + elif source_type == SourceType.command: |
| 152 | + description = source_object.short_doc |
| 153 | + title = f"Command: {source_object.qualified_name}" |
| 154 | + elif source_type == SourceType.tag: |
| 155 | + title = f"Tag: {source_object}" |
| 156 | + description = "" |
| 157 | + else: |
| 158 | + title = f"Cog: {source_object.qualified_name}" |
| 159 | + description = source_object.description.splitlines()[0] |
| 160 | + |
| 161 | + embed = Embed(title=title, description=description) |
| 162 | + embed.add_field(name="Source Code", value=f"[Go to GitHub]({url})") |
| 163 | + line_text = f":{first_line}" if first_line else "" |
| 164 | + embed.set_footer(text=f"{location}{line_text}") |
| 165 | + |
| 166 | + return embed |
0 commit comments