async def send(cls,
bot: Bot,
trakt: TraktClient,
storage: Storage,
user_id,
se: ShowEpisode,
watched: bool = None):
text = rendering.render_html(
'calendar_notification',
show_episode=se,
)
creds, user_pref = await asyncio.gather(
storage.get_creds(user_id),
storage.get_pref(user=user_id),
)
on_watch = user_pref.get('on_watch', 'hide')
sess = trakt.auth(creds.access_token)
if watched is None:
watched = await sess.watched(se.episode.id)
if watched and on_watch == 'delete':
return
keyboard_markup = await cls.markup(se,
watched=watched,
hide=on_watch == 'hide')
await bot.send_message(user_id,
text,
reply_markup=keyboard_markup,
disable_web_page_preview=watched)
async def cancel_handler(message: Message):
user_id = message.from_user.id
storage = Storage.get_current()
user_data = await storage.get_data(user=user_id)
de = user_data.get('deleted_episode')
if de:
asyncio.create_task(message.bot.send_chat_action(user_id, 'typing'))
# todo: use lock
logger.debug(f"restoring deleted notification")
sess = await trakt_session(user_id, storage)
del user_data['deleted_episode']
se, rfh_data, _ = await asyncio.gather(
sess.search_by_episode_id(de, extended=True),
sess.remove_from_history(de),
storage.set_data(user=user_id, data=user_data),
)
logger.debug((se, rfh_data))
await CalendarNotification.send(
message.bot,
TraktClient.get_current(),
storage,
user_id,
se,
watched=False,
)
else:
await asyncio.gather(
storage.finish(user=user_id),
message.answer("Whatever it was, it was canceled."),
)
async def process_auth_flow(message: Message):
storage = Storage.get_current()
trakt = TraktClient.get_current()
user_id = message.from_user.id
flow = trakt.device_auth_flow()
data = await flow.__anext__()
code = data['user_code']
url = data['verification_url']
msg_text = render_html('auth_message', url=url, code=code)
reply_kwargs = dict(disable_web_page_preview=True)
reply = await message.answer(msg_text, **reply_kwargs)
async for ok, data in flow:
state = await storage.get_state(user=user_id)
if state != 'auth':
await reply.edit_text("❌ canceled", **reply_kwargs)
return
if ok:
await asyncio.gather(
storage.save_creds(user_id, data),
reply.edit_text("✅ successfully authenticated"),
)
return data['access_token']
else:
text = render_html('auth_message', url=url, code=code, time_left=data)
await reply.edit_text(text, **reply_kwargs)
else:
await reply.edit_text("❌ failed to authenticated in time")
async def test_cached_call_func(store: Storage):
def func(a):
return a + 1
key = store.make_func_key(func, 4)
assert await store.get_cache(key) is None
assert await store.cached_call(func, 4) == 5
assert await store.get_cache(key) == b'5'
assert await store.cached_call(func, 4) == 5
async def test_cached_call_coro(store: Storage):
async def coro(a):
return a + 1
key = store.make_func_key(coro, 4)
assert await store.get_cache(key) is None
assert await store.cached_call(coro, 4) == 5
assert await store.get_cache(key) == b'5'
assert await store.cached_call(coro, 4) == 5
async def store():
store = Storage(uri=REDIS_URL, db=1)
try:
yield store
finally:
conn = await store.redis()
await conn.flushdb()
await store.close()
await store.wait_closed()
async def calendar_notification_watch_handler(query: CallbackQuery,
callback_data: dict):
user_id = query.from_user.id
episode_id = callback_data['id']
prev_watched = callback_data.get('watched') == '1'
store = Storage.get_current()
sess, user_pref = await asyncio.gather(trakt_session(user_id),
store.get_pref(user=user_id))
on_watch = user_pref.get('on_watch', 'hide')
watched = await sess.watched(episode_id)
# if message was created more than 48 hours ago then it cannot be deleted
now = datetime.now()
delta = now - query.message.date
if on_watch == 'delete' and delta >= timedelta(hours=48):
warn = await query.message.reply(
"quick note: bot cannot delete messages older then 48 hours",
disable_notification=True)
asyncio.create_task(postponed_delete(warn, delay=5))
on_watch = 'hide'
# delete message if it is marked as watched
if on_watch == 'delete':
watched = await toggle_watched_status(sess, episode_id, watched)
if watched:
logger.debug("episode is watched and on_watch=delete")
await asyncio.gather(
store.update_data(user=user_id,
data={'deleted_episode': episode_id}),
query.message.delete(),
query.answer("added to history"),
)
return
# sync with current watch status
else:
if watched is prev_watched:
watched = await toggle_watched_status(sess, episode_id, watched)
else:
msg = 'watched' if watched else 'unwatched'
logger.debug(f"user already marked this as {msg}")
se = await sess.search_by_episode_id(episode_id, extended=True)
logger.debug(se)
# update keyboard
hide = on_watch == 'hide'
markup = await CalendarNotification.markup(se, watched, hide=hide)
await asyncio.gather(
query.message.edit_text(query.message.html_text,
reply_markup=markup,
disable_web_page_preview=hide and watched),
query.answer(
"added to history" if watched else "removed from history"),
)
async def on_watch_behavior_handler(message: Message, command_args,
command_args_error):
if command_args_error:
await message.answer(command_args_error)
return
b = command_args.behavior
storage = Storage.get_current()
if b is None:
user_pref = await storage.get_pref(user=message.from_user.id)
b = user_pref.get('on_watch', 'hide')
await message.answer(f"current 'on watch' behavior is {b!r}")
else:
await asyncio.gather(
storage.update_pref(user=message.from_user.id, on_watch=b),
message.answer(f"new 'on watch' behavior was set: {b}"),
)
async def on_startup(dispatcher: Dispatcher, **kwargs):
logger.debug('setting up services')
dispatcher.storage = Storage(REDIS_URL)
queue = await arq.create_pool(get_redis_settings())
dispatcher.context_vars.update({
'storage': dispatcher.storage,
'trakt': TraktClient(),
'queue': (worker_queue_var, queue),
})
# setup handlers
from traktogram.handlers import auth_router, cmd_router, notification_router, error_router
for router in (auth_router, cmd_router, notification_router, error_router):
dispatcher.add_router(router)
await dispatcher.storage.redis() # test connection
async def logout_handler(message: Message):
storage = Storage.get_current()
trakt = TraktClient.get_current()
queue = worker_queue_var.get()
user_id = message.from_user.id
action = asyncio.create_task(message.bot.send_chat_action(message.chat.id, 'typing'))
creds = await storage.get_creds(user_id)
if creds:
sess = trakt.auth(creds.access_token)
tasks = [
action,
sess.revoke_token(),
storage.finish(user=user_id),
storage.remove_creds(message.from_user.id),
message.answer("Successfully logged out."),
]
if keys := await get_tasks_keys(queue, user_id):
tasks.append(queue.delete(*keys))
await asyncio.gather(*tasks)
async def auth_handler(message: Message):
storage = Storage.get_current()
trakt = TraktClient.get_current()
queue = worker_queue_var.get()
user_id = message.from_user.id
if await storage.has_creds(user_id):
logger.debug("user already authenticated")
text = "You are already authenticated. Do you want to /logout?"
await message.answer(text)
return
await storage.set_state(user=user_id, state='auth')
try:
if access_token := await process_auth_flow(message):
sess = trakt.auth(access_token)
service = NotificationScheduler(queue)
await service.schedule(sess, user_id)
finally:
await storage.finish(user=user_id)
async def on_startup(ctx: dict):
NotificationScheduler.send_single_task_name = send_calendar_notifications.__name__
NotificationScheduler.send_multi_task_name = send_calendar_multi_notifications.__name__
ctx['trakt'] = TraktClient()
ctx['storage'] = Storage(REDIS_URL)
ctx['bot'] = Bot(BOT_TOKEN, parse_mode='html')