def add_wmk(update, context):
if not check_user_data(update, context, WMK_ID):
return ConversationHandler.END
_ = set_lang(update, context)
update.effective_message.reply_text(
_("Adding the watermark onto your PDF file"),
reply_markup=ReplyKeyboardRemove())
# Setup temporary files
temp_files = [tempfile.NamedTemporaryFile() for _ in range(2)]
src_fn, wmk_fn = [x.name for x in temp_files]
user_data = context.user_data
src_file_id = user_data[WMK_ID]
wmk_file_id = update.effective_message.document.file_id
src_reader = open_pdf(update, context, src_file_id, src_fn)
if src_reader is not None:
wmk_reader = open_pdf(update, context, wmk_file_id, wmk_fn,
_("watermark"))
if wmk_reader is not None:
# Add watermark
pdf_writer = PdfFileWriter()
for page in src_reader.pages:
page.mergePage(wmk_reader.getPage(0))
pdf_writer.addPage(page)
# Send result file
write_send_pdf(update, context, pdf_writer, "file.pdf",
"watermarked")
# Clean up memory and files
if user_data[WMK_ID] == src_file_id:
del user_data[WMK_ID]
for tf in temp_files:
tf.close()
return ConversationHandler.END
def get_pdf_preview(update, context):
result = check_back_user_data(update, context)
if result is not None:
return result
_ = set_lang(update, context)
update.effective_message.reply_text(
_("Extracting a preview for your PDF file"),
reply_markup=ReplyKeyboardRemove())
with tempfile.NamedTemporaryFile() as tf1:
user_data = context.user_data
file_id, file_name = user_data[PDF_INFO]
pdf_reader = open_pdf(update, context, file_id, tf1.name)
if pdf_reader:
# Get first page of PDF file
pdf_writer = PdfFileWriter()
pdf_writer.addPage(pdf_reader.getPage(0))
with tempfile.NamedTemporaryFile() as tf2:
# Write cover preview PDF file
with open(tf2.name, "wb") as f:
pdf_writer.write(f)
with tempfile.TemporaryDirectory() as dir_name:
# Convert cover preview to JPEG
out_fn = os.path.join(
dir_name,
f"Preview_{os.path.splitext(file_name)[0]}.png")
imgs = pdf2image.convert_from_path(tf2.name, fmt="png")
imgs[0].save(out_fn)
# Send result file
send_result_file(update, context, out_fn, "preview")
# Clean up memory and files
if user_data[PDF_INFO] == file_id:
del user_data[PDF_INFO]
return ConversationHandler.END
def split_pdf(update: Update, context: CallbackContext) -> int:
result = check_back_user_data(update, context)
if result is not None:
return result
_ = set_lang(update, context)
message = update.effective_message
split_range = message.text
if not PageRange.valid(split_range):
message.reply_text(
_(
"The range is invalid. Try again",
reply_markup=get_back_markup(update, context),
))
return WAIT_SPLIT_RANGE
message.reply_text(_("Splitting your PDF file"),
reply_markup=ReplyKeyboardRemove())
with tempfile.NamedTemporaryFile() as tf:
user_data = context.user_data
file_id, file_name = user_data[PDF_INFO]
pdf_reader = open_pdf(update, context, file_id, tf.name)
if pdf_reader is not None:
merger = PdfFileMerger()
merger.append(pdf_reader, pages=PageRange(split_range))
write_send_pdf(update, context, merger, file_name, "split")
# Clean up memory
if user_data[PDF_INFO] == file_id:
del user_data[PDF_INFO]
return ConversationHandler.END