121 lines
3.6 KiB
Python
121 lines
3.6 KiB
Python
#!/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())
|