Edm0nd/edmond/plugins/wikipedia.py

import time
from typing import cast

try:
    import wikipedia

    DEPENDENCIES_FOUND = True
except ImportError:
    DEPENDENCIES_FOUND = False

from edmond.plugin import Plugin
from edmond.plugins.plus import PlusPlugin


class WikipediaPlugin(Plugin):

    REQUIRED_CONFIGS = [
        "commands",
        "ambiguous_response",
        "empty_response",
        "lang",
    ]
    NUM_RETRIES = 3

    def __init__(self, bot):
        super().__init__(bot)
        if not self.is_ready:
            return
        wikipedia.set_lang(self.config["lang"])

    def on_pubmsg(self, event):
        if not self.should_handle_command(event.arguments[0]):
            return False
        # "science"
        if self.command.ident == self.config["commands"][0]:
            self.tell_random_summary(event)
        # "definition"
        elif self.command.ident == self.config["commands"][1]:
            self.tell_definition(event)
        return True

    def tell_random_summary(self, event):
        page = None
        retries = self.NUM_RETRIES
        while retries > 0:
            try:
                page = wikipedia.page(title=wikipedia.random())
                break
            except Exception as exc:
                # The wikipedia package can raise a lot of different stuff,
                # so we sort of have to catch broadly.
                self.bot.log_d(f"Wikipedia exception: {exc}")
            retries -= 1
        if page:
            reply = WikipediaPlugin.limit_text_length(page.summary)
            self.register_url_for_plus(page.url, event.target)
            self.bot.say(event.target, summary)

    def tell_definition(self, event):
        page = None
        reply = ""
        retries = self.NUM_RETRIES
        while retries > 0:
            try:
                page = wikipedia.page(title=self.command.content)
                break
            except wikipedia.exceptions.DisambiguationError:
                reply = self.config["ambiguous_response"]
                break
            except wikipedia.exceptions.PageError:
                reply = self.config["empty_response"]
                break
            except:
                reply = self.bot.config["error_message"]
                # Keep trying after a slight delay.
                time.sleep(1)
            retries -= 1
        if page:
            reply = WikipediaPlugin.limit_text_length(page.summary)
            self.register_url_for_plus(page.url, event.target)
        self.bot.say(event.target, reply)

    def register_url_for_plus(self, url: str, target: str):
        if plus_plugin := self.bot.get_plugin("plus"):
            def handler(plus_event):
                self.bot.say(plus_event.target, url)
            cast(PlusPlugin, plus_plugin).add_handler(target, handler)

    @staticmethod
    def limit_text_length(text, max_length=200):
        """Limit text size to 200 characters max."""
        words = text.split(" ")
        cut_text = ""
        while words:
            next_word = words.pop(0)
            if len(cut_text) + len(next_word) + 1 >= max_length:
                break
            cut_text += next_word + " "
        if len(cut_text) < len(text):
            cut_text = cut_text[:-1] + "…"
        else:
            cut_text = cut_text.rstrip()
        return cut_text