fulvia/modules/wikipedia.py

#!/usr/bin/env python3
"""
Display excerpts from wikipedia.
"""
import requests

from module import commands, example, url_callback


def wiki_search(query, num):
	"""
	Searches en.wikipedia for the given query, and returns the specified
	number of results.
	"""
	search_url = "http://en.wikipedia.org/w/api.php?format=json&action=query" \
			   + f"&list=search&srlimit={num}&srprop=timestamp&srwhat=text" \
			   + "&srsearch="
	search_url += query
	res = requests.get(search_url, verify=True)
	res.raise_for_status()
	data = res.json()
	if 'query' in data:
		data = data['query']['search']
		return [r['title'] for r in data]
	else:
		return None


def say_snippet(bot, query, show_url=True):
	page_name = query.replace('_', ' ')
	query = query.replace(' ', '_')
	snippet = wiki_snippet(query)
	msg = f'[\x0304WIKIPEDIA\x03] \x0310{page_name}\x03 | \x0312"{snippet}"'
	if show_url:
		msg = msg + f"\x03 | \x0307https://en.wikipedia.org/wiki/{query}"
	msg = msg[:270] + '...' # arbitrary cutoff
	bot.msg(msg)


def wiki_snippet(query):
	"""
	Retrives a snippet of the specified length from the given page on
	en.wikipedia.
	"""
	snippet_url = "https://en.wikipedia.org/w/api.php?format=json" \
				+ "&action=query&prop=extracts&exintro&explaintext" \
				+ "&exchars=300&redirects&titles="
	snippet_url += query
	res = requests.get(snippet_url, verify=True)
	res.raise_for_status()
	snippet = res.json()
	snippet = snippet['query']['pages']

	# For some reason, the API gives the page *number* as the key, so we just
	# grab the first page number in the results.
	snippet = snippet[list(snippet.keys())[0]]

	return snippet['extract']


@url_callback(".wikipedia.org/wiki/")
def wiki_info(bot, url):
	"""
	Retrives a snippet of the specified length from the given page on the given
	server.
	"""
	_, _, query = url.partition("wiki/")
	if not query:
		return
	say_snippet(bot, query, False)


@commands('wikipedia', 'wiki')
@example('.wiki San Francisco')
def wikipedia(bot, trigger):
	"""Search wikipedia and return a snippet of the results."""
	if len(trigger.args) < 2:
		return bot.reply("What do you want me to look up?")

	query = ' '.join(trigger.args[1:])
	if not query:
		return bot.reply("What do you want me to look up?")

	data = wiki_search(query, 1)
	if not data:
		return bot.reply("I can't find any results for that.")
	else:
		data = data[0]
	say_snippet(bot, data)