import re
import logging
import asyncio
import aiohttp
import json
from functools import partialmethod
from . chat import Chat, Sender
__author__ = "Stepan Zastupov"
__copyright__ = "Copyright 2015, Stepan Zastupov"
__license__ = "MIT"
API_URL = "https://api.telegram.org"
API_TIMEOUT = 60
RETRY_TIMEOUT = 30
RETRY_CODES = [429, 500, 502, 503, 504]
BOTAN_URL = "https://api.botan.io/track"
MESSAGE_TYPES = [
"location", "photo", "document", "audio", "voice", "sticker", "contact"
]
logger = logging.getLogger("aiotg")
[docs]class Bot:
"""Telegram bot framework designed for asyncio
:param str api_token: Telegram bot token, ask @BotFather for this
:param int api_timeout: Timeout for long polling
:param str botan_token: Token for http://botan.io
:param str name: Bot name
"""
_running = False
_offset = 0
def __init__(self, api_token, api_timeout=API_TIMEOUT,
botan_token=None, name=None):
self.api_token = api_token
self.api_timeout = api_timeout
self.botan_token = botan_token
self.name = name
def no_handle(mt):
return lambda chat, msg: logger.debug("no handle for %s", mt)
self._handlers = {mt: no_handle(mt) for mt in MESSAGE_TYPES}
self._commands = []
self._default = lambda c, m: None
self._inline = lambda iq: None
self._callback = lambda c, cq: None
@asyncio.coroutine
[docs] def loop(self):
"""
Return bot's main loop as coroutine. Use with asyncio.
:Example:
>>> loop = asyncio.get_event_loop()
>>> loop.run_until_complete(bot.loop())
or
>>> loop = asyncio.get_event_loop()
>>> loop.create_task(bot.loop())
"""
self._running = True
while self._running:
updates = yield from self.api_call(
'getUpdates',
offset=self._offset + 1,
timeout=self.api_timeout
)
self._process_updates(updates)
[docs] def run(self):
"""
Convenience method for running bots
:Example:
>>> if __name__ == '__main__':
>>> bot.run()
"""
loop = asyncio.get_event_loop()
try:
loop.run_until_complete(self.loop())
except KeyboardInterrupt:
self.stop()
[docs] def command(self, regexp):
"""
Register a new command
:param str regexp: Regular expression matching the command to register
:Example:
>>> @bot.command(r"/echo (.+)")
>>> def echo(chat, match):
>>> return chat.reply(match.group(1))
"""
def decorator(fn):
self._commands.append((regexp, fn))
return fn
return decorator
[docs] def default(self, callback):
"""
Set callback for default command that is called on unrecognized
commands for 1-to-1 chats
:Example:
>>> @bot.default
>>> def echo(chat, message):
>>> return chat.reply(message["text"])
"""
self._default = callback
return callback
[docs] def inline(self, callback):
"""
Set callback for inline queries
:Example:
>>> @bot.inline
>>> def echo(iq):
>>> return iq.answer([
>>> {"type": "text", "title": "test", "id", "0"}
>>> ])
"""
self._inline = callback
return callback
[docs] def callback(self, callback):
"""
Set callback for callback queries
:Example:
>>> @bot.callback
>>> def echo(chat, cq):
>>> return cq.answer()
"""
self._callback = callback
return callback
[docs] def handle(self, msg_type):
"""
Set handler for specific message type
:Example:
>>> @bot.handle("audio")
>>> def handle(chat, audio):
>>> pass
"""
def wrap(callback):
self._handlers[msg_type] = callback
return callback
return wrap
[docs] def channel(self, channel_name):
"""
Construct a Chat object used to post to channel
:param str channel_name: Channel name
"""
return Chat(self, channel_name, "channel")
[docs] def private(self, user_id):
"""
Construct a Chat object used to post direct messages
:param str user_id: User id
"""
return Chat(self, user_id, "private")
[docs] def group(self, group_id):
"""
Construct a Chat object used to post group messages
:param str group_id: Group chat id
"""
return Chat(self, group_id, "group")
[docs] async def api_call(self, method, **params):
"""
Call Telegram API.
See https://core.telegram.org/bots/api for reference.
:param str method: Telegram API method
:param params: Arguments for the method call
"""
url = "{0}/bot{1}/{2}".format(API_URL, self.api_token, method)
logger.debug("api_call %s, %s", method, params)
response = await aiohttp.post(url, data=params)
if response.status == 200:
return (await response.json())
elif response.status in RETRY_CODES:
logger.info("Server returned %d, retrying in %d sec.",
response.status, RETRY_TIMEOUT)
await response.release()
await asyncio.sleep(RETRY_TIMEOUT)
return (await self.api_call(method, **params))
else:
if response.headers['content-type'] == 'application/json':
err_msg = (await response.json())["description"]
else:
err_msg = await response.read()
logger.error(err_msg)
raise RuntimeError(err_msg)
_send_message = partialmethod(api_call, "sendMessage")
[docs] def send_message(self, chat_id, text, **options):
"""
Send a text message to chat
:param int chat_id: ID of the chat to send the message to
:param str text: Text to send
:param options: Additional sendMessage options
(see https://core.telegram.org/bots/api#sendmessage)
"""
return self._send_message(chat_id=chat_id, text=text, **options)
_edit_message_text = partialmethod(api_call, "editMessageText")
[docs] def edit_message_text(self, chat_id, message_id, text, **options):
"""
Edit a text message in a chat
:param int chat_id: ID of the chat the message to edit is in
:param int message_id: ID of the message to edit
:param str text: Text to edit the message to
:param options: Additional API options
"""
return self._edit_message_text(
chat_id=chat_id,
message_id=message_id,
text=text,
**options
)
@asyncio.coroutine
[docs] def get_file(self, file_id):
"""
Get basic information about a file and prepare it for downloading.
:param int file_id: File identifier to get information about
:return: File object (see https://core.telegram.org/bots/api#file)
"""
json = yield from self.api_call("getFile", file_id=file_id)
return json["result"]
[docs] def download_file(self, file_path, range=None):
"""
Dowload a file from Telegram servers
"""
headers = {"range": range} if range else None
url = "{0}/file/bot{1}/{2}".format(API_URL, self.api_token, file_path)
return aiohttp.get(url, headers=headers)
[docs] def track(self, message, name="Message"):
"""
Track message using http://botan.io
Set botak_token to make it work
"""
if self.botan_token:
asyncio.ensure_future(self._track(message, name))
[docs] def stop(self):
self._running = False
@asyncio.coroutine
def _track(self, message, name):
response = yield from aiohttp.post(
BOTAN_URL,
params={
"token": self.botan_token,
"uid": message["from"]["id"],
"name": name
},
data=json.dumps(message),
headers={'content-type': 'application/json'}
)
if response.status != 200:
logger.info("error submiting stats %d", response.status)
yield from response.release()
def _process_message(self, message):
chat = Chat.from_message(self, message)
for mt in MESSAGE_TYPES:
if mt in message:
self.track(message, mt)
return self._handlers[mt](chat, message[mt])
if "text" not in message:
return
for patterns, handler in self._commands:
m = re.search(patterns, message["text"], re.I)
if m:
self.track(message, handler.__name__)
return handler(chat, m)
# No match, run default if it's a 1to1 chat
if not chat.is_group():
self.track(message, "default")
return self._default(chat, message)
def _process_inline_query(self, query):
iq = InlineQuery(self, query)
return self._inline(iq)
def _process_callback_query(self, query):
chat = Chat.from_message(self, query["message"])
cq = CallbackQuery(self, query)
return self._callback(chat, cq)
def _process_updates(self, updates):
if not updates["ok"]:
logger.error("getUpdates error: %s", updates.get("description"))
return
for update in updates["result"]:
logger.debug("update %s", update)
self._offset = max(self._offset, update["update_id"])
coro = None
if "message" in update:
coro = self._process_message(update["message"])
elif "inline_query" in update:
coro = self._process_inline_query(update["inline_query"])
elif "callback_query" in update:
coro = self._process_callback_query(update["callback_query"])
if coro:
asyncio.ensure_future(coro)
[docs]class TgBot(Bot):
def __init__(self, *args, **kwargs):
logger.warning("TgBot is depricated, use Bot instead")
super().__init__(*args, **kwargs)
[docs]class InlineQuery:
"""
Incoming inline query
See https://core.telegram.org/bots/api#inline-mode for details
"""
def __init__(self, bot, src):
self.bot = bot
self.sender = Sender(src['from'])
self.query_id = src['id']
self.query = src['query']
[docs] def answer(self, results, **options):
return self.bot.api_call(
"answerInlineQuery",
inline_query_id=self.query_id,
results=json.dumps(results),
**options
)
[docs]class TgInlineQuery(InlineQuery):
def __init__(self, *args, **kwargs):
logger.warning("TgInlineQuery is depricated, use InlineQuery instead")
super().__init__(*args, **kwargs)
[docs]class CallbackQuery:
def __init__(self, bot, src):
self.bot = bot
self.query_id = src['id']
self.data = src['data']
[docs] def answer(self):
return self.bot.api_call(
"answerCallbackQuery",
callback_query_id=self.query_id
)