xkcdtextbot/bot.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
import asyncio
import html
from os import environ
from collections import namedtuple

import logging

import aiohttp
from telethon import TelegramClient, events
from telethon.tl.custom import Button
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={}"
MSG_FMT_STR = '<a href="{link}">{number}</a>: <b>{title}</b>\n\n<i>{alt}</i>'
XKCD_JSON_FMT_STR = "https://xkcd.com/{}/info.0.json"
MAX_SEARCH_RESULTS = 10

bot = TelegramClient('xkcd', 6, 'eb06d4abfb49dc3eeb1aeb98ae0f581e')
bot.parse_mode = 'html'

session = None  # set later
me = None  # set later

# blockquote element -> Xkcd
async def parse_blockquote(elem):
    children = list(elem.children)
    title = children[0].text
    link = 'https' + children[-1].text[4:]

    number = link.rsplit('/', 2)[1]
    async with session.get(XKCD_JSON_FMT_STR.format(number)) as resp:
        info = await resp.json()
    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]
async def get_xkcds(text):
    logger.info("getting %s", text)
    if text == '':
        return []
    # TODO return newest when empty
    async with session.get(URL_FMT_STR.format(text)) as resp:
        soup = BeautifulSoup(await resp.text(), "html.parser")

    bqs = soup.find_all("blockquote")[:MAX_SEARCH_RESULTS]
    logger.info(bqs)
    return await asyncio.gather(*(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.
@bot.on(events.NewMessage(pattern='/start$'))
async def start(event):
    """Send a message when the command /start is issued."""
    await event.respond(
        f"Hello! I'm {me.username} and I search for XKCD when used inline.",
        buttons=Button.switch_inline('Try it!', 'cheaply')
    )


@bot.on(events.NewMessage(pattern='/help$'))
async def help(event):
    """Send a message when the command /help is issued."""
    await event.respond(
        f"I only work inline, and it is my job to search for XKCD comics!",
        buttons=Button.switch_inline('Try it!', 'cheaply')
    )


@bot.on(events.InlineQuery)
async def inlinequery(event):
    """Handle the inline query."""
    # TODO show transcript in result but not message?
    builder = event.builder
    result = await asyncio.gather(*(builder.article(
        title=xkcd.title,
        url=xkcd.link,
        text=MSG_FMT_STR.format(
            number=xkcd.number,
            link=xkcd.link,
            title=html.escape(xkcd.title),
            alt=html.escape(xkcd.alt)
        )
    ) for xkcd in await get_xkcds(event.text)))

    # FIXME get_xkcds returns duplicates, which lead to the same result ID
    #       Build a dict by their ID to remove the duplicates
    result = list({r.id: r for r in result}.values())
    await event.answer(result)


async def main():
    global session, me
    async with aiohttp.ClientSession() as session:
        await bot.start(bot_token=environ['TOKEN'])
        async with bot:
            me = await bot.get_me()
            await bot.run_until_disconnected()


if __name__ == '__main__':
    asyncio.get_event_loop().run_until_complete(main())