From b1c85d5cda805605de3def712604bfa6cb6c571d Mon Sep 17 00:00:00 2001 From: Tulir Asokan Date: Fri, 9 Mar 2018 16:54:17 +0200 Subject: [PATCH] Add moviepy as optional dep for HQ thumbnails, make Pillow optional [db updated] --- ...c972368e50_add_metadata_to_telegramfile.py | 35 +++++++ mautrix_telegram/db.py | 9 +- mautrix_telegram/portal.py | 26 +++-- mautrix_telegram/util/file_transfer.py | 97 +++++++++++++++++-- requirements/base.txt | 1 - requirements/optional.txt | 2 + setup.py | 14 ++- 7 files changed, 153 insertions(+), 31 deletions(-) create mode 100644 alembic/versions/cfc972368e50_add_metadata_to_telegramfile.py diff --git a/alembic/versions/cfc972368e50_add_metadata_to_telegramfile.py b/alembic/versions/cfc972368e50_add_metadata_to_telegramfile.py new file mode 100644 index 00000000..205cbf3a --- /dev/null +++ b/alembic/versions/cfc972368e50_add_metadata_to_telegramfile.py @@ -0,0 +1,35 @@ +"""Add metadata to TelegramFile + +Revision ID: cfc972368e50 +Revises: 501dad2868bc +Create Date: 2018-03-09 16:07:01.236712 + +""" +from alembic import op +import sqlalchemy as sa + +# revision identifiers, used by Alembic. +revision = 'cfc972368e50' +down_revision = '501dad2868bc' +branch_labels = None +depends_on = None + + +def upgrade(): + with op.batch_alter_table("telegram_file") as batch_op: + batch_op.add_column(sa.Column('size', sa.Integer(), nullable=True)) + batch_op.add_column(sa.Column('width', sa.Integer(), nullable=True)) + batch_op.add_column(sa.Column('height', sa.Integer(), nullable=True)) + batch_op.add_column(sa.Column('thumbnail', sa.String(), nullable=True)) + batch_op.create_foreign_key(constraint_name="fk_file_thumbnail", + referent_table="telegram_file", + local_cols=['thumbnail'], + remote_cols=['id']) + + +def downgrade(): + with op.batch_alter_table("telegram_file") as batch_op: + batch_op.drop_column('size') + batch_op.drop_column('width') + batch_op.drop_column('height') + batch_op.drop_column('thumbnail') diff --git a/mautrix_telegram/db.py b/mautrix_telegram/db.py index ba82001b..a057e15f 100644 --- a/mautrix_telegram/db.py +++ b/mautrix_telegram/db.py @@ -81,8 +81,8 @@ class Contact(Base): query = None __tablename__ = "contact" - user = Column("user", Integer, ForeignKey("user.tgid"), primary_key=True) - contact = Column("contact", Integer, ForeignKey("puppet.id"), primary_key=True) + user = Column(Integer, ForeignKey("user.tgid"), primary_key=True) + contact = Column(Integer, ForeignKey("puppet.id"), primary_key=True) class Puppet(Base): @@ -112,6 +112,11 @@ class TelegramFile(Base): mime_type = Column(String) was_converted = Column(Boolean) timestamp = Column(BigInteger) + size = Column(Integer, nullable=True) + width = Column(Integer, nullable=True) + height = Column(Integer, nullable=True) + thumbnail_id = Column("thumbnail", String, ForeignKey("telegram_file.id"), nullable=True) + thumbnail = relationship("TelegramFile", uselist=False) def init(db_session): diff --git a/mautrix_telegram/portal.py b/mautrix_telegram/portal.py index 8a40ae0f..a60da924 100644 --- a/mautrix_telegram/portal.py +++ b/mautrix_telegram/portal.py @@ -33,7 +33,6 @@ from mautrix_appservice import MatrixRequestError, IntentError from .db import Portal as DBPortal, Message as DBMessage from . import puppet as p, user as u, formatter, util -from .formatter.util import trim_reply_fallback_html, trim_reply_fallback_text mimetypes.init() @@ -858,11 +857,12 @@ class Portal: async def handle_telegram_document(self, source, intent, evt: Message, relates_to=None): document = evt.media.document - file = await util.transfer_file_to_matrix(self.db, source.client, intent, document) + file = await util.transfer_file_to_matrix(self.db, source.client, intent, document, + document.thumb) if not file: return None name = evt.message - width, height = 0, 0 + width, height = file.width, file.height for attr in document.attributes: if isinstance(attr, DocumentAttributeFilename): name = name or attr.file_name @@ -871,25 +871,21 @@ class Portal: file.mime_type = mime_from_name or file.mime_type elif isinstance(attr, DocumentAttributeSticker): name = f"Sticker for {attr.alt}" - elif isinstance(attr, DocumentAttributeVideo): + elif isinstance(attr, DocumentAttributeVideo) and (not width or not height): width, height = attr.w, attr.h mime_type = document.mime_type or file.mime_type info = { - "size": document.size, + "size": file.size, "mimetype": mime_type, } - if document.thumb and not isinstance(document.thumb, PhotoSizeEmpty): - thumbnail = await util.transfer_file_to_matrix(self.db, source.client, intent, - document.thumb.location) + if file.thumbnail: + info["thumbnail_url"] = file.thumbnail.mxc info["thumbnail_info"] = { - "mimetype": thumbnail.mime_type, - "h": document.thumb.h, - "w": document.thumb.w, - "size": (len(document.thumb.bytes) - if isinstance(document.thumb, PhotoCachedSize) - else document.thumb.size) + "mimetype": file.thumbnail.mime_type, + "h": file.thumbnail.height or document.thumb.h, + "w": file.thumbnail.width or document.thumb.w, + "size": file.thumbnail.size, } - info["thumbnail_url"] = thumbnail.mxc if height and width: info["h"] = height info["w"] = width diff --git a/mautrix_telegram/util/file_transfer.py b/mautrix_telegram/util/file_transfer.py index b64f6346..d8542247 100644 --- a/mautrix_telegram/util/file_transfer.py +++ b/mautrix_telegram/util/file_transfer.py @@ -19,11 +19,22 @@ import time import logging import magic -from PIL import Image from sqlalchemy.exc import IntegrityError, InvalidRequestError +try: + from PIL import Image +except ImportError: + Image = None +try: + from moviepy.editor import VideoFileClip + import random + import string + import os + import mimetypes +except ImportError: + VideoFileClip = random = string = os = mimetypes = None from telethon_aio.tl.types import (Document, FileLocation, InputFileLocation, - InputDocumentFileLocation, PhotoCachedSize) + InputDocumentFileLocation, PhotoSize, PhotoCachedSize) from telethon_aio.errors import LocationInvalidError from ..db import TelegramFile as DBTelegramFile @@ -32,24 +43,86 @@ log = logging.getLogger("mau.util") def _convert_webp(file, to="png"): + if not Image: + return "image/webp", file try: image = Image.open(BytesIO(file)).convert("RGBA") new_file = BytesIO() image.save(new_file, to) - return f"image/{to}", new_file.getvalue() + w, h = image.size + return f"image/{to}", new_file.getvalue(), w, h except Exception: log.exception(f"Failed to convert webp to {to}") return "image/webp", file -async def transfer_file_to_matrix(db, client, intent, location): +def _temp_file_name(ext): + return ("/tmp/mxtg-video-" + + "".join(random.choice(string.ascii_uppercase + string.digits) for _ in range(10)) + + ext) + + +def _read_video_thumbnail(data, video_ext="mp4", frame_ext="png", max_size=(1024, 720)): + # We don't have any way to read the video from memory, so save it to disk. + temp_file = _temp_file_name(video_ext) + with open(temp_file, "wb") as file: + file.write(data) + + # Read temp file and get frame + clip = VideoFileClip(temp_file) + frame = clip.get_frame(0) + + # Convert to png and save to BytesIO + image = Image.fromarray(frame).convert("RGBA") + thumbnail_file = BytesIO() + if max_size: + image.thumbnail(max_size, Image.ANTIALIAS) + image.save(thumbnail_file, frame_ext) + + os.remove(temp_file) + + w, h = image.size + return thumbnail_file.getvalue(), w, h + + +def _location_to_id(location): if isinstance(location, (Document, InputDocumentFileLocation)): - id = f"{location.id}-{location.version}" + return f"{location.id}-{location.version}" elif isinstance(location, (FileLocation, InputFileLocation)): - id = f"{location.volume_id}-{location.local_id}" + return f"{location.volume_id}-{location.local_id}" else: return None + +async def transfer_thumbnail_to_matrix(client, intent, thumbnail_loc, video, mime): + if not Image or not VideoFileClip: + return None + + id = _location_to_id(thumbnail_loc) + if not id: + return None + + video_ext = mimetypes.guess_extension(mime) + if VideoFileClip and video_ext: + file, width, height = _read_video_thumbnail(video, video_ext, frame_ext="png") + mime_type = "image/png" + else: + file = await client.download_file_bytes(thumbnail_loc) + width, height = None, None + mime_type = magic.from_buffer(file, mime=True) + + uploaded = await intent.upload_file(file, mime_type) + + return DBTelegramFile(id=id, mxc=uploaded["content_uri"], mime_type=mime_type, + was_converted=False, timestamp=int(time.time()), size=len(file), + width=width, height=height) + + +async def transfer_file_to_matrix(db, client, intent, location, thumbnail=None): + id = _location_to_id(location) + if not id: + return None + db_file = DBTelegramFile.query.get(id) if db_file: return db_file @@ -58,18 +131,26 @@ async def transfer_file_to_matrix(db, client, intent, location): file = await client.download_file_bytes(location) except LocationInvalidError: return None + width, height = None, None mime_type = magic.from_buffer(file, mime=True) image_converted = False if mime_type == "image/webp": - mime_type, file = _convert_webp(file, to="png") + mime_type, file, width, height = _convert_webp(file, to="png") image_converted = True uploaded = await intent.upload_file(file, mime_type) db_file = DBTelegramFile(id=id, mxc=uploaded["content_uri"], mime_type=mime_type, was_converted=image_converted, - timestamp=int(time.time())) + timestamp=int(time.time()), size=len(file), + width=width, height=height) + if thumbnail: + if isinstance(thumbnail, (PhotoSize, PhotoCachedSize)): + thumbnail = thumbnail.location + db_file.thumbnail = await transfer_thumbnail_to_matrix(client, intent, thumbnail, file, + mime_type) + try: db.add(db_file) db.commit() diff --git a/requirements/base.txt b/requirements/base.txt index b7b8cc90..6e955265 100644 --- a/requirements/base.txt +++ b/requirements/base.txt @@ -5,5 +5,4 @@ python-magic SQLAlchemy alembic Markdown -Pillow future-fstrings diff --git a/requirements/optional.txt b/requirements/optional.txt index 10087cb8..a4400877 100644 --- a/requirements/optional.txt +++ b/requirements/optional.txt @@ -1,2 +1,4 @@ lxml cryptg +Pillow +moviepy diff --git a/setup.py b/setup.py index a9292e9e..5166748e 100644 --- a/setup.py +++ b/setup.py @@ -3,6 +3,14 @@ import sys import glob import mautrix_telegram +extras = { + "highlight_edits": ["lxml>=4.1.1,<5"], + "fast_crypto": ["cryptg>=0.1,<0.2"], + "webp_convert": ["Pillow>=5.0.0,<6"], + "hq_thumbnails": ["moviepy>=0.2,<0.3"], +} +extras["all"] = [deps[0] for deps in extras.values()] + setuptools.setup( name="mautrix-telegram", version=mautrix_telegram.__version__, @@ -23,7 +31,6 @@ setuptools.setup( "alembic>=0.9.8,<0.10", "Markdown>=2.6.11,<3", "ruamel.yaml>=0.15.35,<0.16", - "Pillow>=5.0.0,<6", "future-fstrings>=0.4.2", "python-magic>=0.4.15,<0.5", "telethon-aio>=0.18,<0.19" if sys.version_info >= (3, 6) else "telethon-aio-git", @@ -31,10 +38,7 @@ setuptools.setup( dependency_links=[ "https://github.com/tulir/telethon-asyncio/tarball/9b389cfb4b6d3876e9661c23507f17e96897e4b0#egg=telethon-aio-git-0.18.0+1" ], - extras_require={ - "highlight_edits": ["lxml>=4.1.1,<5"], - "fast_crypto": ["cryptg>=0.1,<0.2"], - }, + extras_require=extras, classifiers=[ "Development Status :: 4 - Beta",