Rewrite issue cog

- Switch to a dataclass based communication
- Use GitHub headers when querying the repo_regex
- Properly handle non-200 return codes
- Use @whitelist_override

1 files changed, 119 insertions, 86 deletions

diff --git a/bot/exts/evergreen/issues.py b/bot/exts/evergreen/issues.py
index afaa7012..dce11678 100644
--- a/bot/exts/evergreen/issues.py
+++ b/bot/exts/evergreen/issues.py
@@ -2,12 +2,13 @@ import logging
 import random
 import re
 import typing as t
-from enum import Enum
+from dataclasses import dataclass
 
 import discord
 from discord.ext import commands, tasks
 
 from bot.constants import Categories, Colours, ERROR_REPLIES, Emojis, Tokens, WHITELISTED_CHANNELS
+from bot.utils.decorators import whitelist_override
 
 log = logging.getLogger(__name__)
 
@@ -17,7 +18,10 @@ BAD_RESPONSE = {
 }
 REQUEST_HEADERS = dict()
 
-REPOS_API = "https://api.github.com/orgs/{org}/repos"
+REPOSITORY_ENDPOINT = "https://api.github.com/orgs/{org}/repos"
+ISSUE_ENDPOINT = "https://api.github.com/repos/{user}/{repository}/issues/{number}"
+PR_MERGE_ENDPOINT = "https://api.github.com/repos/{user}/{repository}/pulls/{number}/merge"
+
 if GITHUB_TOKEN := Tokens.github:
     REQUEST_HEADERS["Authorization"] = f"token {GITHUB_TOKEN}"
 
@@ -35,11 +39,23 @@ CODE_BLOCK_RE = re.compile(
 MAXIMUM_ISSUES = 5
 
 
-class FetchIssueErrors(Enum):
-    """Errors returned in fetch issues."""
+@dataclass
+class FetchError:
+    """Dataclass representing an error while fetching an issue."""
+
+    return_code: int
+    message: str
+
 
-    value_error = "Numbers not found."
-    max_requests = "Max requests hit."
+@dataclass
+class IssueState:
+    """Dataclass representing the state of an issue."""
+
+    repository: str
+    number: int
+    url: str
+    title: str
+    icon_url: str
 
 
 class Issues(commands.Cog):
@@ -48,19 +64,27 @@ class Issues(commands.Cog):
     def __init__(self, bot: commands.Bot):
         self.bot = bot
         self.repos = []
+        self.repo_regex = None
         self.get_pydis_repos.start()
 
     @tasks.loop(minutes=30)
     async def get_pydis_repos(self) -> None:
-        """Get all python-discord repositories on github."""
-        async with self.bot.http_session.get(REPOS_API.format(org="python-discord")) as resp:
+        """
+        Get all python-discord repositories on github.
+
+        This task will update a pipe-separated list of repositories in self.repo_regex.
+        """
+        async with self.bot.http_session.get(
+                REPOSITORY_ENDPOINT.format(org="python-discord"),
+                headers=REQUEST_HEADERS
+        ) as resp:
             if resp.status == 200:
                 data = await resp.json()
                 for repo in data:
                     self.repos.append(repo["full_name"].split("/")[1])
                 self.repo_regex = "|".join(self.repos)
             else:
-                log.debug(f"Failed to get latest Pydis repositories. Status code {resp.status}")
+                log.warning(f"Failed to get latest Pydis repositories. Status code {resp.status}")
 
     @staticmethod
     def check_in_block(message: discord.Message, repo_issue: str) -> bool:
@@ -75,70 +99,85 @@ class Issues(commands.Cog):
 
     async def fetch_issues(
             self,
-            numbers: set,
+            number: int,
             repository: str,
             user: str
-    ) -> t.Union[FetchIssueErrors, str, list]:
-        """Retrieve issue(s) from a GitHub repository."""
-        links = []
-        if not numbers:
-            return FetchIssueErrors.value_error
-
-        if len(numbers) > MAXIMUM_ISSUES:
-            return FetchIssueErrors.max_requests
-
-        for number in numbers:
-            url = f"https://api.github.com/repos/{user}/{repository}/issues/{number}"
-            merge_url = f"https://api.github.com/repos/{user}/{repository}/pulls/{number}/merge"
-            log.trace(f"Querying GH issues API: {url}")
-            async with self.bot.http_session.get(url, headers=REQUEST_HEADERS) as r:
-                json_data = await r.json()
-
-            if r.status in BAD_RESPONSE:
-                log.warning(f"Received response {r.status} from: {url}")
-                return f"[{str(r.status)}] #{number} {BAD_RESPONSE.get(r.status)}"
-
-            # The initial API request is made to the issues API endpoint, which will return information
-            # if the issue or PR is present. However, the scope of information returned for PRs differs
-            # from issues: if the 'issues' key is present in the response then we can pull the data we
-            # need from the initial API call.
-            if "issues" in json_data.get("html_url"):
+    ) -> t.Union[IssueState, FetchError]:
+        """
+        Retrieve an issue from a GitHub repository.
+
+        returns IssueState on success, FetchError on failure.
+        """
+        url = ISSUE_ENDPOINT.format(user=user, repository=repository, number=number)
+        merge_url = PR_MERGE_ENDPOINT.format(user=user, repository=repository, number=number)
+        log.trace(f"Querying GH issues API: {url}")
+
+        async with self.bot.http_session.get(url, headers=REQUEST_HEADERS) as r:
+            json_data = await r.json()
+
+        if r.status == 403:
+            if "X-RateLimit-Remaining" in r.headers and r.headers["X-RateLimit-Remaining"] == "0":
+                log.info(f"Ratelimit reached while fetching {url}")
+                return FetchError(403, "Ratelimit reached, please retry in a few minutes.")
+            return FetchError(403, "Cannot access issue.")
+        elif r.status in (404, 410):
+            return FetchError(r.status, "Issue not found.")
+        elif r.status != 200:
+            return FetchError(r.status, "Error while fetching issue.")
+
+        # The initial API request is made to the issues API endpoint, which will return information
+        # if the issue or PR is present. However, the scope of information returned for PRs differs
+        # from issues: if the 'issues' key is present in the response then we can pull the data we
+        # need from the initial API call.
+        if "issues" in json_data.get("html_url"):
+            if json_data.get("state") == "open":
+                icon_url = Emojis.issue
+            else:
+                icon_url = Emojis.issue_closed
+
+        # If the 'issues' key is not contained in the API response and there is no error code, then
+        # we know that a PR has been requested and a call to the pulls API endpoint is necessary
+        # to get the desired information for the PR.
+        else:
+            log.trace(f"PR provided, querying GH pulls API for additional information: {merge_url}")
+            async with self.bot.http_session.get(merge_url) as m:
                 if json_data.get("state") == "open":
-                    icon_url = Emojis.issue
+                    icon_url = Emojis.pull_request
+                # When the status is 204 this means that the state of the PR is merged
+                elif m.status == 204:
+                    icon_url = Emojis.merge
                 else:
-                    icon_url = Emojis.issue_closed
-
-            # If the 'issues' key is not contained in the API response and there is no error code, then
-            # we know that a PR has been requested and a call to the pulls API endpoint is necessary
-            # to get the desired information for the PR.
-            else:
-                log.trace(f"PR provided, querying GH pulls API for additional information: {merge_url}")
-                async with self.bot.http_session.get(merge_url) as m:
-                    if json_data.get("state") == "open":
-                        icon_url = Emojis.pull_request
-                    # When the status is 204 this means that the state of the PR is merged
-                    elif m.status == 204:
-                        icon_url = Emojis.merge
-                    else:
-                        icon_url = Emojis.pull_request_closed
+                    icon_url = Emojis.pull_request_closed
 
-            issue_url = json_data.get("html_url")
-            links.append([icon_url, f"[{repository}] #{number} {json_data.get('title')}", issue_url])
+        issue_url = json_data.get("html_url")
 
-        return links
+        return IssueState(repository, number, issue_url, json_data.get('title', ''), icon_url)
 
     @staticmethod
-    def get_embed(result: list, user: str = "python-discord", repository: str = "") -> discord.Embed:
-        """Get Response Embed."""
-        description_list = ["{0} [{1}]({2})".format(*link) for link in result]
+    def format_embed(
+            results: t.List[t.Union[IssueState, FetchError]],
+            user: str,
+            repository: t.Optional[str] = None
+    ) -> discord.Embed:
+        """Take a list of IssueState or FetchError and format a Discord embed for them."""
+        description_list = []
+
+        for result in results:
+            if isinstance(result, IssueState):
+                description_list.append(f"{result.icon_url} [{result.title}]({result.url})")
+            elif isinstance(result, FetchError):
+                description_list.append(f"[{result.return_code}] {result.message}")
+
         resp = discord.Embed(
             colour=Colours.bright_green,
             description='\n'.join(description_list)
         )
 
-        resp.set_author(name="GitHub", url=f"https://github.com/{user}/{repository}")
+        embed_url = f"https://github.com/{user}/{repository}" if repository else f"https://github.com/{user}"
+        resp.set_author(name="GitHub", url=embed_url)
         return resp
 
+    @whitelist_override(channels=WHITELISTED_CHANNELS, categories=WHITELISTED_CATEGORIES)
     @commands.command(aliases=("pr",))
     async def issue(
             self,
@@ -148,62 +187,56 @@ class Issues(commands.Cog):
             user: str = "python-discord"
     ) -> None:
         """Command to retrieve issue(s) from a GitHub repository."""
-        if not(
-            ctx.channel.category.id in WHITELISTED_CATEGORIES
-            or ctx.channel.id in WHITELISTED_CHANNELS
-        ):
-            return
-
-        result = await self.fetch_issues(set(numbers), repository, user)
-
-        if result == FetchIssueErrors.value_error:
-            await ctx.invoke(self.bot.get_command('help'), 'issue')
+        # Remove duplicates
+        numbers = set(numbers)
 
-        elif result == FetchIssueErrors.max_requests:
+        if len(numbers) > MAXIMUM_ISSUES:
             embed = discord.Embed(
                 title=random.choice(ERROR_REPLIES),
                 color=Colours.soft_red,
                 description=f"Too many issues/PRs! (maximum of {MAXIMUM_ISSUES})"
             )
             await ctx.send(embed=embed)
+            await ctx.invoke(self.bot.get_command('help'), 'issue')
 
-        elif isinstance(result, list):
-            # Issue/PR format: emoji to show if open/closed/merged, number and the title as a singular link.
-            resp = self.get_embed(result, user, repository)
-            await ctx.send(embed=resp)
-
-        elif isinstance(result, str):
-            await ctx.send(result)
+        results = [await self.fetch_issues(number, repository, user) for number in numbers]
+        await ctx.send(embed=self.format_embed(results))
 
     @commands.Cog.listener()
     async def on_message(self, message: discord.Message) -> None:
         """Command to retrieve issue(s) from a GitHub repository using automatic linking if matching <repo>#<issue>."""
-        message_repo_issue_map = re.findall(fr"({self.repo_regex})#(\d+)", message.content)
+        if not self.repo_regex:
+            log.warning("repo_regex isn't ready, cannot look for issues.")
+            return
+
+        # `issues` will hold a list of two element tuples `(repository, issue_number)`
+        issues = re.findall(fr"({self.repo_regex})#(\d+)", message.content)
         links = []
 
-        if message_repo_issue_map:
+        if issues:
+            log.trace(f"Found {issues = }")
             # Remove duplicates
-            message_repo_issue_map = set(message_repo_issue_map)
+            issues = set(issues)
 
-            if len(message_repo_issue_map) > MAXIMUM_ISSUES:
+            if len(issues) > MAXIMUM_ISSUES:
                 embed = discord.Embed(
                     title=random.choice(ERROR_REPLIES),
                     color=Colours.soft_red,
                     description=f"Too many issues/PRs! (maximum of {MAXIMUM_ISSUES})"
                 )
-                await message.channel.send(embed=embed)
+                await message.channel.send(embed=embed).delete(delay=5)
                 return
 
-            for repo_issue in message_repo_issue_map:
+            for repo_issue in issues:
                 if not self.check_in_block(message, " ".join(repo_issue)):
-                    result = await self.fetch_issues({repo_issue[1]}, repo_issue[0], "python-discord")
-                    if isinstance(result, list):
-                        links.extend(result)
+                    result = await self.fetch_issues(repo_issue[1], repo_issue[0], "python-discord")
+                    if isinstance(result, IssueState):
+                        links.append(result)
 
         if not links:
             return
 
-        resp = self.get_embed(links, "python-discord")
+        resp = self.format_embed(links, "python-discord")
         await message.channel.send(embed=resp)