Merge branch 'master' into ext-managment

3 files changed, 118 insertions, 1 deletions

diff --git a/bot/__init__.py b/bot/__init__.py
index 6976e089..a9a0865e 100644
--- a/bot/__init__.py
+++ b/bot/__init__.py
@@ -37,7 +37,7 @@ os.makedirs(log_dir, exist_ok=True)
 
 # File handler rotates logs every 5 MB
 file_handler = logging.handlers.RotatingFileHandler(
-    log_file, maxBytes=5*(2**20), backupCount=10)
+    log_file, maxBytes=5 * (2**20), backupCount=10)
 file_handler.setLevel(logging.TRACE if Client.debug else logging.DEBUG)
 
 # Console handler prints to terminal
diff --git a/bot/constants.py b/bot/constants.py
index f3424673..935b90e0 100644
--- a/bot/constants.py
+++ b/bot/constants.py
@@ -269,3 +269,6 @@ POSITIVE_REPLIES = [
     "Aye aye, cap'n!",
     "I'll allow it.",
 ]
+
+class Wikipedia:
+    total_chance = 3
diff --git a/bot/exts/evergreen/wikipedia.py b/bot/exts/evergreen/wikipedia.py
new file mode 100644
index 00000000..be36e2c4
--- /dev/null
+++ b/bot/exts/evergreen/wikipedia.py
@@ -0,0 +1,114 @@
+import asyncio
+import datetime
+import logging
+from typing import List, Optional
+
+from aiohttp import client_exceptions
+from discord import Color, Embed, Message
+from discord.ext import commands
+
+from bot.constants import Wikipedia
+
+log = logging.getLogger(__name__)
+
+SEARCH_API = "https://en.wikipedia.org/w/api.php?action=query&list=search&srsearch={search_term}&format=json"
+WIKIPEDIA_URL = "https://en.wikipedia.org/wiki/{title}"
+
+
+class WikipediaSearch(commands.Cog):
+    """Get info from wikipedia."""
+
+    def __init__(self, bot: commands.Bot):
+        self.bot = bot
+        self.http_session = bot.http_session
+
+    @staticmethod
+    def formatted_wiki_url(index: int, title: str) -> str:
+        """Formating wikipedia link with index and title."""
+        return f'`{index}` [{title}]({WIKIPEDIA_URL.format(title=title.replace(" ", "_"))})'
+
+    async def search_wikipedia(self, search_term: str) -> Optional[List[str]]:
+        """Search wikipedia and return the first 10 pages found."""
+        pages = []
+        async with self.http_session.get(SEARCH_API.format(search_term=search_term)) as response:
+            try:
+                data = await response.json()
+
+                search_results = data["query"]["search"]
+
+                # Ignore pages with "may refer to"
+                for search_result in search_results:
+                    log.info("trying to append titles")
+                    if "may refer to" not in search_result["snippet"]:
+                        pages.append(search_result["title"])
+            except client_exceptions.ContentTypeError:
+                pages = None
+
+        log.info("Finished appending titles")
+        return pages
+
+    @commands.cooldown(1, 10, commands.BucketType.user)
+    @commands.command(name="wikipedia", aliases=["wiki"])
+    async def wikipedia_search_command(self, ctx: commands.Context, *, search: str) -> None:
+        """Return list of results containing your search query from wikipedia."""
+        titles = await self.search_wikipedia(search)
+
+        def check(message: Message) -> bool:
+            return message.author.id == ctx.author.id and message.channel == ctx.channel
+
+        if not titles:
+            await ctx.send("Sorry, we could not find a wikipedia article using that search term")
+            return
+
+        async with ctx.typing():
+            log.info("Finished appending titles to titles_no_underscore list")
+
+            s_desc = "\n".join(self.formatted_wiki_url(index, title) for index, title in enumerate(titles, start=1))
+            embed = Embed(colour=Color.blue(), title=f"Wikipedia results for `{search}`", description=s_desc)
+            embed.timestamp = datetime.datetime.utcnow()
+            await ctx.send(embed=embed)
+        embed = Embed(colour=Color.green(), description="Enter number to choose")
+        msg = await ctx.send(embed=embed)
+        titles_len = len(titles)  # getting length of list
+
+        for retry_count in range(1, Wikipedia.total_chance + 1):
+            retries_left = Wikipedia.total_chance - retry_count
+            if retry_count < Wikipedia.total_chance:
+                error_msg = f"You have `{retries_left}/{Wikipedia.total_chance}` chances left"
+            else:
+                error_msg = 'Please try again by using `.wiki` command'
+            try:
+                message = await ctx.bot.wait_for('message', timeout=60.0, check=check)
+                response_from_user = await self.bot.get_context(message)
+
+                if response_from_user.command:
+                    return
+
+                response = int(message.content)
+                if response < 0:
+                    await ctx.send(f"Sorry, but you can't give negative index, {error_msg}")
+                elif response == 0:
+                    await ctx.send(f"Sorry, please give an integer between `1` to `{titles_len}`, {error_msg}")
+                else:
+                    await ctx.send(WIKIPEDIA_URL.format(title=titles[response - 1].replace(" ", "_")))
+                    break
+
+            except asyncio.TimeoutError:
+                embed = Embed(colour=Color.red(), description=f"Time's up {ctx.author.mention}")
+                await msg.edit(embed=embed)
+                break
+
+            except ValueError:
+                await ctx.send(f"Sorry, but you cannot do that, I will only accept an positive integer, {error_msg}")
+
+            except IndexError:
+                await ctx.send(f"Sorry, please give an integer between `1` to `{titles_len}`, {error_msg}")
+
+            except Exception as e:
+                log.info(f"Caught exception {e}, breaking out of retry loop")
+                break
+
+
+def setup(bot: commands.Bot) -> None:
+    """Wikipedia Cog load."""
+    bot.add_cog(WikipediaSearch(bot))