aboutsummaryrefslogtreecommitdiffstats
path: root/bot/exts/evergreen/wikipedia.py
diff options
context:
space:
mode:
Diffstat (limited to 'bot/exts/evergreen/wikipedia.py')
-rw-r--r--bot/exts/evergreen/wikipedia.py164
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))