aboutsummaryrefslogtreecommitdiffstats
path: root/bot/exts/evergreen/wikipedia.py
blob: c1fff87373d38d3a642d8db48902211a1533906a (plain) (blame)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
import asyncio
import datetime
import logging
from typing import List

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 WikipediaCog(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) -> List[str]:
        """Search wikipedia and return the first 10 pages found."""
        async with self.http_session.get(SEARCH_API.format(search_term=search_term)) as response:
            data = await response.json()

        pages = []

        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"])

        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(WikipediaCog(bot))