Add moviepy as optional dep for HQ thumbnails, make Pillow optional

[db updated]

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')

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)
+    user = Column(Integer, ForeignKey("user.tgid"), primary_key=True)
-    contact = Column("contact", Integer, ForeignKey("puppet.id"), 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):

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):
+        if file.thumbnail:
-            thumbnail = await util.transfer_file_to_matrix(self.db, source.client, intent,
+            info["thumbnail_url"] = file.thumbnail.mxc
-                                                           document.thumb.location)
             info["thumbnail_info"] = {
-                "mimetype": thumbnail.mime_type,
+                "mimetype": file.thumbnail.mime_type,
-                "h": document.thumb.h,
+                "h": file.thumbnail.height or document.thumb.h,
-                "w": document.thumb.w,
+                "w": file.thumbnail.width or document.thumb.w,
-                "size": (len(document.thumb.bytes)
+                "size": file.thumbnail.size,
-                         if isinstance(document.thumb, PhotoCachedSize)
-                         else document.thumb.size)
             }
-            info["thumbnail_url"] = thumbnail.mxc
         if height and width:
             info["h"] = height
             info["w"] = width

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()

requirements/base.txt

@@ -5,5 +5,4 @@ python-magic
 SQLAlchemy
 alembic
 Markdown
-Pillow
 future-fstrings

requirements/optional.txt

@@ -1,2 +1,4 @@
 lxml
 cryptg
+Pillow
+moviepy

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={
+    extras_require=extras,
-        "highlight_edits": ["lxml>=4.1.1,<5"],
-        "fast_crypto": ["cryptg>=0.1,<0.2"],
-    },
 
     classifiers=[
         "Development Status :: 4 - Beta",