iartamonov/mautrix-telegram
1
0
Fork 0
Code Issues Pull Requests 1 Actions Packages Projects Releases Wiki Activity

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

Browse Source 
[db updated]
This commit is contained in:
Tulir Asokan
2018-03-09 16:54:17 +02:00
parent a469e6ed10
commit b1c85d5cda
7 changed files with 153 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files
alembic/versions/cfc972368e50_add_metadata_to_telegramfile.py
+35
View File
@@ -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
+7 -2
View File
@@ -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):
mautrix_telegram/portal.py
+11 -15
View File
@@ -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
mautrix_telegram/util/file_transfer.py
+89 -8
View File
@@ -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
-1
View File
@@ -5,5 +5,4 @@ python-magic
SQLAlchemy
alembic
Markdown
Pillow
future-fstrings
requirements/optional.txt
+2
View File
@@ -1,2 +1,4 @@
lxml
cryptg
Pillow
moviepy
setup.py
+9 -5
View File
@@ -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",