xkcdtextbot/bot.py

135 lines
3.8 KiB
Python

#!/usr/bin/env python
# -*- coding: utf-8 -*-
from uuid import uuid4
from os import environ
from collections import namedtuple
from json import loads
import logging
from telegram import InlineQueryResultArticle
from telegram import InputTextMessageContent
from telegram import ParseMode
from telegram.ext import CommandHandler
from telegram.ext import InlineQueryHandler
from telegram.ext import Updater
import requests
from bs4 import BeautifulSoup
# Enable logging
logging.basicConfig(format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
level=logging.INFO)
logger = logging.getLogger(__name__)
Xkcd = namedtuple('Xkcd', ['title', 'link', 'transcript', 'alt', 'number'])
URL_FMT_STR = "http://www.ohnorobot.com/index.php?Search=Search&comic=56&s={}"
# FIXME HTML parsemode + escaping
MSG_FMT_STR = "[{number}]({link}): *{title}*\n\n_{alt}_"
XKCD_JSON_FMT_STR = "https://xkcd.com/{}/info.0.json"
MAX_SEARCH_RESULTS = 10
# blockquote element -> Xkcd
def parse_blockquote(elem):
children = list(elem.children)
title = children[0].text
link = 'https' + children[-1].text[4:]
number = link.rsplit('/', 2)[1]
info = loads(requests.get(XKCD_JSON_FMT_STR.format(number)).text)
alt = info['alt']
# TODO markdown bold the <span> matches
text = ''.join(
[e.text if hasattr(e, 'text') else e
for
e in children[1:-1]]
)
return Xkcd(title, link, text, alt, number)
# string -> [Xkcd]
def get_xkcds(text):
logger.info("getting %s", text)
if text == '':
return []
# TODO return newest when empty
soup = BeautifulSoup(requests.get(URL_FMT_STR.format(text)).text, "html.parser")
bqs = soup.find_all("blockquote")[:MAX_SEARCH_RESULTS]
logger.info(bqs)
return (parse_blockquote(e) for e in bqs)
# Define a few command handlers. These usually take the two arguments bot and
# update. Error handlers also receive the raised TelegramError object in error.
def start(bot, update):
"""Send a message when the command /start is issued."""
# TODO
update.message.reply_text('Hi!')
def help(bot, update):
"""Send a message when the command /help is issued."""
# TODO
update.message.reply_text('Help!')
def inlinequery(bot, update):
"""Handle the inline query."""
# TODO show transcript in result but not message?
query = update.inline_query.query
results = []
for xkcd in get_xkcds(query):
number = xkcd.number
link = xkcd.link
title = xkcd.title
alt = xkcd.alt
results.append(InlineQueryResultArticle(
id=uuid4(),
title=xkcd.title,
url=xkcd.link,
input_message_content=InputTextMessageContent(
MSG_FMT_STR.format(number=number, link=link, title=title, alt=alt),
parse_mode=ParseMode.MARKDOWN)))
update.inline_query.answer(results)
def error(bot, update, error):
"""Log Errors caused by Updates."""
logger.warning('Update "%s" caused error "%s"', update, error)
def main():
# Create the Updater and pass it your bot's token.
updater = Updater(environ["TOKEN"])
# Get the dispatcher to register handlers
dp = updater.dispatcher
# on different commands - answer in Telegram
dp.add_handler(CommandHandler("start", start))
dp.add_handler(CommandHandler("help", help))
# on noncommand i.e message - echo the message on Telegram
dp.add_handler(InlineQueryHandler(inlinequery))
# log all errors
dp.add_error_handler(error)
# Start the Bot
updater.start_polling()
# Block until the user presses Ctrl-C or the process receives SIGINT,
# SIGTERM or SIGABRT. This should be used most of the time, since
# start_polling() is non-blocking and will stop the bot gracefully.
updater.idle()
if __name__ == '__main__':
main()