diff options
Diffstat (limited to 'bot/exts/evergreen/wikipedia.py')
| -rw-r--r-- | bot/exts/evergreen/wikipedia.py | 164 | 
1 files changed, 72 insertions, 92 deletions
| diff --git a/bot/exts/evergreen/wikipedia.py b/bot/exts/evergreen/wikipedia.py index be36e2c4..068c4f43 100644 --- a/bot/exts/evergreen/wikipedia.py +++ b/bot/exts/evergreen/wikipedia.py @@ -1,114 +1,94 @@ -import asyncio -import datetime  import logging +import re +from datetime import datetime +from html import unescape  from typing import List, Optional -from aiohttp import client_exceptions -from discord import Color, Embed, Message +from discord import Color, Embed, TextChannel  from discord.ext import commands -from bot.constants import Wikipedia +from bot.bot import Bot +from bot.utils import LinePaginator  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}" +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: commands.Bot): +    def __init__(self, bot: 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 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 -    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 +                else: +                    await channel.send( +                        "Sorry, we could not find a wikipedia article using that search term." +                    ) +                    return +            else: +                log.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: -        """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: +        """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=Color.blurple() +            ) +            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)) | 
