Compare commits

...

5 Commits

2 changed files with 41 additions and 20 deletions

58
bot.py
View File

@ -1,14 +1,15 @@
#!/usr/bin/env python #!/usr/bin/env python
# -*- coding: utf-8 -*- # -*- coding: utf-8 -*-
import asyncio import asyncio
import html
from os import environ from os import environ
from collections import namedtuple from collections import namedtuple
from json import loads
import logging import logging
import requests import aiohttp
from telethon import TelegramClient, events from telethon import TelegramClient, events
from telethon.tl.custom import Button
from bs4 import BeautifulSoup from bs4 import BeautifulSoup
# Enable logging # Enable logging
@ -20,21 +21,25 @@ logger = logging.getLogger(__name__)
Xkcd = namedtuple('Xkcd', ['title', 'link', 'transcript', 'alt', 'number']) Xkcd = namedtuple('Xkcd', ['title', 'link', 'transcript', 'alt', 'number'])
URL_FMT_STR = "http://www.ohnorobot.com/index.php?Search=Search&comic=56&s={}" URL_FMT_STR = "http://www.ohnorobot.com/index.php?Search=Search&comic=56&s={}"
# FIXME HTML parsemode + escaping MSG_FMT_STR = '<a href="{link}">{number}</a>: <b>{title}</b>\n\n<i>{alt}</i>'
MSG_FMT_STR = "[{number}]({link}): **{title}**\n\n__{alt}__"
XKCD_JSON_FMT_STR = "https://xkcd.com/{}/info.0.json" XKCD_JSON_FMT_STR = "https://xkcd.com/{}/info.0.json"
MAX_SEARCH_RESULTS = 10 MAX_SEARCH_RESULTS = 10
bot = TelegramClient('xkcd', 6, 'eb06d4abfb49dc3eeb1aeb98ae0f581e') bot = TelegramClient('xkcd', 6, 'eb06d4abfb49dc3eeb1aeb98ae0f581e')
bot.parse_mode = 'html'
session = None # set later
me = None # set later
# blockquote element -> Xkcd # blockquote element -> Xkcd
def parse_blockquote(elem): async def parse_blockquote(elem):
children = list(elem.children) children = list(elem.children)
title = children[0].text title = children[0].text
link = 'https' + children[-1].text[4:] link = 'https' + children[-1].text[4:]
number = link.rsplit('/', 2)[1] number = link.rsplit('/', 2)[1]
info = loads(requests.get(XKCD_JSON_FMT_STR.format(number)).text) async with session.get(XKCD_JSON_FMT_STR.format(number)) as resp:
info = await resp.json()
alt = info['alt'] alt = info['alt']
# TODO markdown bold the <span> matches # TODO markdown bold the <span> matches
@ -47,16 +52,17 @@ def parse_blockquote(elem):
return Xkcd(title, link, text, alt, number) return Xkcd(title, link, text, alt, number)
# string -> [Xkcd] # string -> [Xkcd]
def get_xkcds(text): async def get_xkcds(text):
logger.info("getting %s", text) logger.info("getting %s", text)
if text == '': if text == '':
return [] return []
# TODO return newest when empty # TODO return newest when empty
soup = BeautifulSoup(requests.get(URL_FMT_STR.format(text)).text, "html.parser") 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] bqs = soup.find_all("blockquote")[:MAX_SEARCH_RESULTS]
logger.info(bqs) logger.info(bqs)
return (parse_blockquote(e) for e in 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 # Define a few command handlers. These usually take the two arguments bot and
@ -64,15 +70,19 @@ def get_xkcds(text):
@bot.on(events.NewMessage(pattern='/start$')) @bot.on(events.NewMessage(pattern='/start$'))
async def start(event): async def start(event):
"""Send a message when the command /start is issued.""" """Send a message when the command /start is issued."""
# TODO await event.respond(
await event.respond('Hi!') 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$')) @bot.on(events.NewMessage(pattern='/help$'))
async def help(event): async def help(event):
"""Send a message when the command /help is issued.""" """Send a message when the command /help is issued."""
# TODO await event.respond(
await event.respond('Help!') 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) @bot.on(events.InlineQuery)
@ -83,8 +93,13 @@ async def inlinequery(event):
result = await asyncio.gather(*(builder.article( result = await asyncio.gather(*(builder.article(
title=xkcd.title, title=xkcd.title,
url=xkcd.link, url=xkcd.link,
text=MSG_FMT_STR.format(number=xkcd.number, link=xkcd.link, title=xkcd.title, alt=xkcd.alt) text=MSG_FMT_STR.format(
) for xkcd in get_xkcds(event.text))) 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 # FIXME get_xkcds returns duplicates, which lead to the same result ID
# Build a dict by their ID to remove the duplicates # Build a dict by their ID to remove the duplicates
@ -92,11 +107,14 @@ async def inlinequery(event):
await event.answer(result) await event.answer(result)
def main(): async def main():
bot.start(bot_token=environ['TOKEN']) global session, me
with bot: async with aiohttp.ClientSession() as session:
bot.run_until_disconnected() await bot.start(bot_token=environ['TOKEN'])
async with bot:
me = await bot.get_me()
await bot.run_until_disconnected()
if __name__ == '__main__': if __name__ == '__main__':
main() asyncio.get_event_loop().run_until_complete(main())

3
requirements.txt Normal file
View File

@ -0,0 +1,3 @@
aiohttp~=3.5
beautifulsoup4~=4.7
telethon~=1.7