diff --git a/bot/exts/fun/wikipedia.py b/bot/exts/fun/wikipedia.py new file mode 100644 index 00000000..ffeb5566 --- /dev/null +++ b/bot/exts/fun/wikipedia.py @@ -0,0 +1,89 @@ +import re +from datetime import datetime +from html import unescape +from typing import List, Optional + +from bot.bot import Bot +from bot.constants import Colours +from bot.utils.pagination import LinePaginator +from discord import Embed, TextChannel +from discord.ext import commands +from loguru import logger + + +SEARCH_API = ( + "https://en.wikipedia.org/w/api.php?action=query&list=search&prop=info&inprop=url&utf8=&" + "format=json&origin=*&srlimit={number_of_results}&srsearch={string}" +) +WIKI_THUMBNAIL = ( + "https://upload.wikimedia.org/wikipedia/en/thumb/8/80/Wikipedia-logo-v2.svg" + "/330px-Wikipedia-logo-v2.svg.png" +) +WIKI_SNIPPET_REGEX = r"(|<[^>]*>)" +WIKI_SEARCH_RESULT = "**[{name}]({url})**\n" "{description}\n" + + +class WikipediaSearch(commands.Cog): + """Get info from wikipedia.""" + + def __init__(self, bot: Bot): + self.bot = bot + + async def wiki_request( + self, channel: TextChannel, search: str + ) -> Optional[List[str]]: + """Search wikipedia search string and return formatted first 10 pages found.""" + url = SEARCH_API.format(number_of_results=10, string=search) + async with self.bot.http_session.get(url=url) as resp: + if resp.status == 200: + raw_data = await resp.json() + number_of_results = raw_data["query"]["searchinfo"]["totalhits"] + + if number_of_results: + results = raw_data["query"]["search"] + lines = [] + + for article in results: + line = WIKI_SEARCH_RESULT.format( + name=article["title"], + description=unescape( + re.sub(WIKI_SNIPPET_REGEX, "", article["snippet"]) + ), + url=f"https://en.wikipedia.org/?curid={article['pageid']}", + ) + lines.append(line) + + return lines + + else: + await channel.send( + "Sorry, we could not find a wikipedia article using that search term." + ) + return + else: + logger.info( + f"Unexpected response `{resp.status}` while searching wikipedia for `{search}`" + ) + await channel.send( + "Whoops, the Wikipedia API is having some issues right now. Try again later." + ) + return + + @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: + """Sends paginated top 10 results of Wikipedia search..""" + contents = await self.wiki_request(ctx.channel, search) + + if contents: + embed = Embed(title="Wikipedia Search Results", colour=Colours.green) + embed.set_thumbnail(url=WIKI_THUMBNAIL) + embed.timestamp = datetime.utcnow() + await LinePaginator.paginate(contents, ctx, embed) + + +def setup(bot: Bot) -> None: + """Wikipedia Cog load.""" + bot.add_cog(WikipediaSearch(bot))