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
112
113
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))
|
