Skip to content

Commit

Permalink
fix review comments
Browse files Browse the repository at this point in the history
  • Loading branch information
armpico committed Sep 11, 2022
1 parent 015c77e commit 033e461
Show file tree
Hide file tree
Showing 8 changed files with 223 additions and 180 deletions.
116 changes: 72 additions & 44 deletions picobot/handlers.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@
import os
import shlex
from functools import wraps
from pathlib import Path

from slugify import slugify
from telegram import Bot, Message, Update
Expand All @@ -16,8 +17,7 @@
from .video_editor import sticker_from_video, VideoTooLongError
from .repository.repo import repository

IMG_DIR = ROOT_DIR / 'images'
VID_DIR = ROOT_DIR / 'videos'
MEDIA_DIR = ROOT_DIR / 'media'
IMG_PREFIX = 'img'
VIDEO_PREFIX = 'vid'
AVATAR_PREFIX = 'avatar'
Expand Down Expand Up @@ -53,12 +53,12 @@ def build_pack_name(title: str, bot: Bot) -> str:
return f'{slug}_by_{bot.username}'


def start(update: Update, context: CallbackContext):
def start(update: Update, context: CallbackContext) -> None:
"""Send a message when the command /start is issued."""
update.message.reply_text(responses.GREETING)


def create_pack(update: Update, context: CallbackContext):
def create_pack(update: Update, context: CallbackContext) -> None:
bot = context.bot
user = update.message.from_user

Expand All @@ -70,28 +70,35 @@ def create_pack(update: Update, context: CallbackContext):

title = splittext[1]
name = build_pack_name(title, bot)
png_sticker = open(IMG_DIR / 'caravela.png', 'rb')
png_sticker = open(MEDIA_DIR / 'caravela.png', 'rb')
emoji = splittext[2] if len(splittext) > 2 else DEFAULT_EMOJI

# Create Pack
try:
bot.create_new_sticker_set(
user_id=user.id, name=name, title=title, png_sticker=png_sticker, emojis=emoji,
user_id=user.id,
name=name,
title=title,
png_sticker=png_sticker,
emojis=emoji,
)
sticker = bot.get_sticker_set(name).stickers[0]
update.message.reply_sticker(sticker)
repository().add_pack_to_user(user, name)
except Exception as exc:
logger.error(
"Exception on Create Pack. User %s (id %d) Pack %s", user.first_name, user.id, name,
"Exception on Create Pack. User %s (id %d) Pack %s",
user.first_name,
user.id,
name,
)

logger.error(exc)
update.message.reply_text(responses.ERROR_MSG)
png_sticker.close()


def create_video_pack(update: Update, context: CallbackContext):
def create_video_pack(update: Update, context: CallbackContext) -> None:
bot = context.bot
user = update.message.from_user

Expand All @@ -103,27 +110,34 @@ def create_video_pack(update: Update, context: CallbackContext):

title = splittext[1]
name = build_pack_name(title, bot)
with open(VID_DIR / 'caravela.webm', 'rb') as webm_sticker:
with open(MEDIA_DIR / 'caravela.webm', 'rb') as webm_sticker:
emoji = splittext[2] if len(splittext) > 2 else DEFAULT_EMOJI

# Create Video Pack
try:
bot.create_new_sticker_set(
user_id=user.id, name=name, title=title, webm_sticker=webm_sticker, emojis=emoji,
user_id=user.id,
name=name,
title=title,
webm_sticker=webm_sticker,
emojis=emoji,
)
sticker = bot.get_sticker_set(name).stickers[0]
update.message.reply_sticker(sticker)
repository().add_pack_to_user(user, name)
except Exception as exc:
logger.error(
"Exception on Create Pack. User %s (id %d) Pack %s", user.first_name, user.id, name,
"Exception on Create Pack. User %s (id %d) Pack %s",
user.first_name,
user.id,
name,
)

logger.error(exc)
update.message.reply_text(responses.ERROR_MSG)


def add_sticker(update: Update, context: CallbackContext):
def add_sticker(update: Update, context: CallbackContext) -> None:
bot = context.bot
msg = update.message

Expand Down Expand Up @@ -167,27 +181,27 @@ def add_sticker(update: Update, context: CallbackContext):
return
elif msg_type == MsgType.VIDEO:
media = msg.video[-1]
if add_video(bot, msg, media, user_id, pack_name, emoji):
if add_video(bot, msg, media, user_id, pack_name, emoji, circle=False):
return
elif msg_type == MsgType.REP_VIDEO:
media = msg.reply_to_message.video
if add_video(bot, msg, media, user_id, pack_name, emoji, False):
if add_video(bot, msg, media, user_id, pack_name, emoji, circle=False):
return
elif msg_type == MsgType.VIDEO_NOTE:
media = msg.video_note[-1]
if add_video(bot, msg, media, user_id, pack_name, emoji, True):
if add_video(bot, msg, media, user_id, pack_name, emoji, circle=True):
return
elif msg_type == MsgType.REP_VIDEO_NOTE:
media = msg.reply_to_message.video_note
if add_video(bot, msg, media, user_id, pack_name, emoji, True):
if add_video(bot, msg, media, user_id, pack_name, emoji, circle=True):
return
elif msg_type == MsgType.DOCUMENT_VIDEO:
media = msg.document[-1]
if add_video(bot, msg, media, user_id, pack_name, emoji, False):
if add_video(bot, msg, media, user_id, pack_name, emoji, circle=False):
return
elif msg_type == MsgType.REP_DOCUMENT_VIDEO:
media = msg.reply_to_message.document
if add_video(bot, msg, media, user_id, pack_name, emoji, False):
if add_video(bot, msg, media, user_id, pack_name, emoji, circle=False):
return
elif msg_type == MsgType.DOCUMENT:
if add_document(bot, msg, user_id, pack_name, emoji, False):
Expand Down Expand Up @@ -219,14 +233,13 @@ def add_text(bot: Bot, msg: Message, user_id: int, pack_name: str, emoji: str):
other_user_id = msg.reply_to_message.from_user.id
msg_time = msg.reply_to_message.date.strftime('%H:%M')
photos = bot.get_user_profile_photos(other_user_id, limit=1).photos
avatar_path = ''
avatar_path = Path('')
try:
photo = photos[0][0]
avatar_path = IMG_DIR / f'{AVATAR_PREFIX}{other_user_id}.jpg'
avatar_path = MEDIA_DIR / f'{AVATAR_PREFIX}{other_user_id}.jpg'
bot.get_file(photo.file_id).download(custom_path=avatar_path)
except Exception:
msg.reply_text(responses.ERROR_DOWNLOAD_PHOTO)
avatar_path = ''

text = msg.reply_to_message.text
# save as png
Expand All @@ -238,23 +251,24 @@ def add_text(bot: Bot, msg: Message, user_id: int, pack_name: str, emoji: str):
)
sticker = bot.get_sticker_set(pack_name).stickers[-1]
msg.reply_sticker(sticker)
img_path.unlink(missing_ok=True)
except Exception as exc:
if isinstance(exc, telegram.error.BadRequest):
exception_msg = exc.message.lower()
if exception_msg in responses.TELEGRAM_ERROR_CODES:
msg.reply_text(responses.TELEGRAM_ERROR_CODES[exception_msg])
return True
logger.error(
"Exception on add_text. User %s (id %d) Pack %s", username, user_id, pack_name,
"Exception on add_text. User %s (id %d) Pack %s",
username,
user_id,
pack_name,
)
logger.error(exc)
return False
finally:
if os.path.exists(img_path):
os.remove(img_path)
if os.path.exists(avatar_path):
os.remove(avatar_path)

img_path.unlink(missing_ok=True)
avatar_path.unlink(missing_ok=True)
return True


Expand All @@ -272,25 +286,29 @@ def add_photo(bot: Bot, msg: Message, user_id: int, pack_name: str, emoji: str,
photo = msg.reply_to_message.photo[-1]
else:
photo = msg.photo[-1]
img_path = IMG_DIR / f'{IMG_PREFIX}{user_id}.jpg'
img_path = MEDIA_DIR / f'{IMG_PREFIX}{user_id}.jpg'
try:
bot.get_file(photo.file_id).download(custom_path=img_path)
# resize and save as png
img_path = sticker_from_image(img_path)
with open(img_path, 'rb') as png_sticker:
png_path = sticker_from_image(img_path)
with open(png_path, 'rb') as png_sticker:
bot.add_sticker_to_set(
user_id=user_id, name=pack_name, png_sticker=png_sticker, emojis=emoji
)
sticker = bot.get_sticker_set(pack_name).stickers[-1]
msg.reply_sticker(sticker)
img_path.unlink(missing_ok=True)
png_path.unlink(missing_ok=True)
except Exception as exc:
if isinstance(exc, telegram.error.BadRequest):
exception_msg = exc.message.lower()
if exception_msg in responses.TELEGRAM_ERROR_CODES:
msg.reply_text(responses.TELEGRAM_ERROR_CODES[exception_msg])
return True
logger.error(
"Exception on add_photo. User id %d Pack %s", user_id, pack_name,
"Exception on add_photo. User id %d Pack %s",
user_id,
pack_name,
)
logger.error(exc)
return False
Expand All @@ -299,25 +317,27 @@ def add_photo(bot: Bot, msg: Message, user_id: int, pack_name: str, emoji: str,


def add_video(
bot: Bot, msg: Message, video, user_id: int, pack_name: str, emoji: str, circle: bool
bot: Bot, msg: Message, video, user_id: int, pack_name: str, emoji: str, *, circle: bool
):
video_path = VID_DIR / f'{VIDEO_PREFIX}{user_id}.mp4'
video_path = MEDIA_DIR / f'{VIDEO_PREFIX}{user_id}.mp4'
try:
if bot.get_file(video.file_id).file_size > 2172859:
msg.reply_text(responses.FILE_TOO_LARGE)
raise FileTooLargeError()("File size exceeds limits")
raise FileTooLargeError("File size exceeds limits")
bot.get_file(video.file_id).download(custom_path=video_path)
# resize and save as webm
if not circle:
video_path = sticker_from_video(video_path)
webm_path = sticker_from_video(video_path)
else:
video_path = sticker_from_video(video_path, IMG_DIR / "circle_mask.png")
with open(video_path, 'rb') as webm_sticker:
webm_path = sticker_from_video(video_path, MEDIA_DIR / "circle_mask.png")
with open(webm_path, 'rb') as webm_sticker:
bot.add_sticker_to_set(
user_id=user_id, name=pack_name, webm_sticker=webm_sticker, emojis=emoji
)
sticker = bot.get_sticker_set(pack_name).stickers[-1]
msg.reply_sticker(sticker)
video_path.unlink(missing_ok=True)
webm_path.unlink(missing_ok=True)
except Exception as exc:
if isinstance(exc, VideoTooLongError):
msg.reply_text(responses.VIDEO_TOO_LONG)
Expand All @@ -328,7 +348,9 @@ def add_video(
msg.reply_text(responses.TELEGRAM_ERROR_CODES[exception_msg])
return True
logger.error(
"Exception on add_video. User id %d Pack %s", user_id, pack_name,
"Exception on add_video. User id %d Pack %s",
user_id,
pack_name,
)
logger.error(exc)
return False
Expand Down Expand Up @@ -361,26 +383,30 @@ def add_document(bot: Bot, msg: Message, user_id: int, pack_name: str, emoji: st
def insert_sticker_in_pack(bot: Bot, msg: Message, user_id: int, pack_name: str, emoji: str):
sticker_id = msg.reply_to_message.sticker.file_id

img_path = IMG_DIR / f'{IMG_PREFIX}{user_id}.jpg'
img_path = MEDIA_DIR / f'{IMG_PREFIX}{user_id}.jpg'
try:
sticker_file = bot.get_file(sticker_id)
sticker_file.download(custom_path=str(img_path))
# resize and save as png
img_path = sticker_from_image(img_path)
with open(img_path, 'rb') as png_sticker:
png_path = sticker_from_image(img_path)
with open(png_path, 'rb') as png_sticker:
bot.add_sticker_to_set(
user_id=user_id, name=pack_name, png_sticker=png_sticker, emojis=emoji
)
sticker = bot.get_sticker_set(pack_name).stickers[-1]
msg.reply_sticker(sticker)
img_path.unlink(missing_ok=True)
png_path.unlink(missing_ok=True)
except Exception as exc:
if isinstance(exc, telegram.error.BadRequest):
exception_msg = exc.message.lower()
if exception_msg in responses.TELEGRAM_ERROR_CODES:
msg.reply_text(responses.TELEGRAM_ERROR_CODES[exception_msg])
return True
logger.error(
"Exception inserting sticker in pack. User id %d Pack %s", user_id, pack_name,
"Exception inserting sticker in pack. User id %d Pack %s",
user_id,
pack_name,
)
logger.error(exc)
return False
Expand All @@ -390,7 +416,7 @@ def insert_sticker_in_pack(bot: Bot, msg: Message, user_id: int, pack_name: str,
def insert_video_sticker_in_pack(bot: Bot, msg: Message, user_id: int, pack_name: str, emoji: str):
sticker_id = msg.reply_to_message.sticker.file_id

video_path = VID_DIR / f'{VIDEO_PREFIX}{user_id}.webm'
video_path = MEDIA_DIR / f'{VIDEO_PREFIX}{user_id}.webm'
try:
sticker_file = bot.get_file(sticker_id)
sticker_file.download(custom_path=str(video_path))
Expand All @@ -407,7 +433,9 @@ def insert_video_sticker_in_pack(bot: Bot, msg: Message, user_id: int, pack_name
msg.reply_text(responses.TELEGRAM_ERROR_CODES[exception_msg])
return True
logger.error(
"Exception inserting sticker in pack. User id %d Pack %s", user_id, pack_name,
"Exception inserting sticker in pack. User id %d Pack %s",
user_id,
pack_name,
)
logger.error(exc)
return False
Expand Down
File renamed without changes
Binary file added picobot/media/caravela.webm
Binary file not shown.
File renamed without changes
4 changes: 2 additions & 2 deletions picobot/painter.py
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,7 @@
Box,
)

IMG_DIR = ROOT_DIR / 'images'
MEDIA_DIR = ROOT_DIR / 'media'
FONT_DIR = ROOT_DIR / 'fonts'
IMG_PREFIX = 'img'
AVATAR_MASK_NAME = 'avatar_mask.png'
Expand Down Expand Up @@ -236,7 +236,7 @@ def sticker_from_text(user_id: int, username: str, text: str, avatar_path: str,
draw_time(dr, text=msg_time, points=points_balloon)

img = resize_to_sticker_limits(img)
img_path = IMG_DIR / f'{IMG_PREFIX}{user_id}.png'
img_path = MEDIA_DIR / f'{IMG_PREFIX}{user_id}.png'
img.save(img_path)
img.close()
return img_path
Expand Down
Loading

0 comments on commit 033e461

Please sign in to comment.