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

Blacken and isort code

Browse Source
This commit is contained in:
Tulir Asokan
2021-12-21 01:36:24 +02:00
parent f2af17d359
commit 6d25e9687e
55 changed files with 3752 additions and 2018 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mautrix_telegram/util/file_transfer.py
+149 -68
View File
@@ -1,5 +1,5 @@
# mautrix-telegram - A Matrix-Telegram puppeting bridge
# Copyright (C) 2019 Tulir Asokan
# Copyright (C) 2021 Tulir Asokan
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published by
@@ -13,27 +13,40 @@
#
# You should have received a copy of the GNU Affero General Public License
# along with this program. If not, see <https://www.gnu.org/licenses/>.
from typing import Optional, Tuple, Union, Dict
from __future__ import annotations
from typing import Optional, Union
from io import BytesIO
import time
import logging
import asyncio
import tempfile
import magic
from asyncpg import UniqueViolationError
from sqlite3 import IntegrityError
import asyncio
import logging
import tempfile
import time
from telethon.tl.types import (Document, InputFileLocation, InputDocumentFileLocation,
TypePhotoSize, PhotoSize, PhotoCachedSize, InputPhotoFileLocation,
InputPeerPhotoFileLocation)
from telethon.errors import (AuthBytesInvalidError, AuthKeyInvalidError, LocationInvalidError,
SecurityError, FileIdInvalidError)
from asyncpg import UniqueViolationError
from telethon.errors import (
AuthBytesInvalidError,
AuthKeyInvalidError,
FileIdInvalidError,
LocationInvalidError,
SecurityError,
)
from telethon.tl.types import (
Document,
InputDocumentFileLocation,
InputFileLocation,
InputPeerPhotoFileLocation,
InputPhotoFileLocation,
PhotoCachedSize,
PhotoSize,
TypePhotoSize,
)
import magic
from mautrix.appservice import IntentAPI
from ..tgclient import MautrixTelegramClient
from ..db import TelegramFile as DBTelegramFile
from ..tgclient import MautrixTelegramClient
from ..util import sane_mimetypes
from .parallel_file_transfer import parallel_transfer_to_matrix
from .tgs_converter import convert_tgs_to
@@ -55,13 +68,21 @@ except ImportError:
log: logging.Logger = logging.getLogger("mau.util")
TypeLocation = Union[Document, InputDocumentFileLocation, InputPeerPhotoFileLocation,
InputFileLocation, InputPhotoFileLocation]
TypeLocation = Union[
Document,
InputDocumentFileLocation,
InputPeerPhotoFileLocation,
InputFileLocation,
InputPhotoFileLocation,
]
def convert_image(file: bytes, source_mime: str = "image/webp", target_type: str = "png",
thumbnail_to: Optional[Tuple[int, int]] = None
) -> Tuple[str, bytes, Optional[int], Optional[int]]:
def convert_image(
file: bytes,
source_mime: str = "image/webp",
target_type: str = "png",
thumbnail_to: tuple[int, int] | None = None,
) -> tuple[str, bytes, int | None, int | None]:
if not Image:
return source_mime, file, None, None
try:
@@ -77,8 +98,12 @@ def convert_image(file: bytes, source_mime: str = "image/webp", target_type: str
return source_mime, file, None, None
def _read_video_thumbnail(data: bytes, video_ext: str = "mp4", frame_ext: str = "png",
max_size: Tuple[int, int] = (1024, 720)) -> Tuple[bytes, int, int]:
def _read_video_thumbnail(
data: bytes,
video_ext: str = "mp4",
frame_ext: str = "png",
max_size: tuple[int, int] = (1024, 720),
) -> tuple[bytes, int, int]:
with tempfile.NamedTemporaryFile(prefix="mxtg_video_", suffix=f".{video_ext}") as file:
# We don't have any way to read the video from memory, so save it to disk.
file.write(data)
@@ -109,11 +134,17 @@ def _location_to_id(location: TypeLocation) -> str:
return str(location.photo_id)
async def transfer_thumbnail_to_matrix(client: MautrixTelegramClient, intent: IntentAPI,
thumbnail_loc: TypeLocation, mime_type: str, encrypt: bool,
video: Optional[bytes], custom_data: Optional[bytes] = None,
width: Optional[int] = None, height: [int] = None
) -> Optional[DBTelegramFile]:
async def transfer_thumbnail_to_matrix(
client: MautrixTelegramClient,
intent: IntentAPI,
thumbnail_loc: TypeLocation,
mime_type: str,
encrypt: bool,
video: bytes | None,
custom_data: bytes | None = None,
width: int | None = None,
height: int | None = None,
) -> DBTelegramFile | None:
if not Image or not VideoFileClip:
return None
@@ -151,28 +182,45 @@ async def transfer_thumbnail_to_matrix(client: MautrixTelegramClient, intent: In
if decryption_info:
decryption_info.url = content_uri
db_file = DBTelegramFile(id=loc_id, mxc=content_uri, mime_type=mime_type,
was_converted=False, timestamp=int(time.time()), size=len(file),
width=width, height=height, decryption_info=decryption_info)
db_file = DBTelegramFile(
id=loc_id,
mxc=content_uri,
mime_type=mime_type,
was_converted=False,
timestamp=int(time.time()),
size=len(file),
width=width,
height=height,
decryption_info=decryption_info,
)
try:
await db_file.insert()
except (UniqueViolationError, IntegrityError) as e:
log.exception(f"{e.__class__.__name__} while saving transferred file thumbnail data. "
"This was probably caused by two simultaneous transfers of the same file, "
"and might (but probably won't) cause problems with thumbnails or something.")
log.exception(
f"{e.__class__.__name__} while saving transferred file thumbnail data. "
"This was probably caused by two simultaneous transfers of the same file, "
"and might (but probably won't) cause problems with thumbnails or something."
)
return db_file
transfer_locks: Dict[str, asyncio.Lock] = {}
transfer_locks: dict[str, asyncio.Lock] = {}
TypeThumbnail = Optional[Union[TypeLocation, TypePhotoSize]]
async def transfer_file_to_matrix(client: MautrixTelegramClient, intent: IntentAPI,
location: TypeLocation, thumbnail: TypeThumbnail = None, *,
is_sticker: bool = False, tgs_convert: Optional[dict] = None,
filename: Optional[str] = None, encrypt: bool = False,
parallel_id: Optional[int] = None) -> Optional[DBTelegramFile]:
async def transfer_file_to_matrix(
client: MautrixTelegramClient,
intent: IntentAPI,
location: TypeLocation,
thumbnail: TypeThumbnail = None,
*,
is_sticker: bool = False,
tgs_convert: dict | None = None,
filename: str | None = None,
encrypt: bool = False,
parallel_id: int | None = None,
) -> DBTelegramFile | None:
location_id = _location_to_id(location)
if not location_id:
return None
@@ -187,17 +235,32 @@ async def transfer_file_to_matrix(client: MautrixTelegramClient, intent: IntentA
lock = asyncio.Lock()
transfer_locks[location_id] = lock
async with lock:
return await _unlocked_transfer_file_to_matrix(client, intent, location_id, location,
thumbnail, is_sticker, tgs_convert,
filename, encrypt, parallel_id)
return await _unlocked_transfer_file_to_matrix(
client,
intent,
location_id,
location,
thumbnail,
is_sticker,
tgs_convert,
filename,
encrypt,
parallel_id,
)
async def _unlocked_transfer_file_to_matrix(client: MautrixTelegramClient, intent: IntentAPI,
loc_id: str, location: TypeLocation,
thumbnail: TypeThumbnail, is_sticker: bool,
tgs_convert: Optional[dict], filename: Optional[str],
encrypt: bool, parallel_id: Optional[int]
) -> Optional[DBTelegramFile]:
async def _unlocked_transfer_file_to_matrix(
client: MautrixTelegramClient,
intent: IntentAPI,
loc_id: str,
location: TypeLocation,
thumbnail: TypeThumbnail,
is_sticker: bool,
tgs_convert: dict | None,
filename: str | None,
encrypt: bool,
parallel_id: int | None,
) -> DBTelegramFile | None:
db_file = await DBTelegramFile.get(loc_id)
if db_file:
return db_file
@@ -205,8 +268,9 @@ async def _unlocked_transfer_file_to_matrix(client: MautrixTelegramClient, inten
converted_anim = None
if parallel_id and isinstance(location, Document) and (not is_sticker or not tgs_convert):
db_file = await parallel_transfer_to_matrix(client, intent, loc_id, location, filename,
encrypt, parallel_id)
db_file = await parallel_transfer_to_matrix(
client, intent, loc_id, location, filename, encrypt, parallel_id
)
mime_type = location.mime_type
file = None
else:
@@ -223,12 +287,13 @@ async def _unlocked_transfer_file_to_matrix(client: MautrixTelegramClient, inten
image_converted = False
# A weird bug in alpine/magic makes it return application/octet-stream for gzips...
is_tgs = (mime_type == "application/gzip"
or (mime_type == "application/octet-stream"
and magic.from_buffer(file).startswith("gzip")))
is_tgs = mime_type == "application/gzip" or (
mime_type == "application/octet-stream" and magic.from_buffer(file).startswith("gzip")
)
if is_sticker and tgs_convert and is_tgs:
converted_anim = await convert_tgs_to(file, tgs_convert["target"],
**tgs_convert["args"])
converted_anim = await convert_tgs_to(
file, tgs_convert["target"], **tgs_convert["args"]
)
mime_type = converted_anim.mime
file = converted_anim.data
width, height = converted_anim.width, converted_anim.height
@@ -244,29 +309,45 @@ async def _unlocked_transfer_file_to_matrix(client: MautrixTelegramClient, inten
if decryption_info:
decryption_info.url = content_uri
db_file = DBTelegramFile(id=loc_id, mxc=content_uri, decryption_info=decryption_info,
mime_type=mime_type, was_converted=image_converted,
timestamp=int(time.time()), size=len(file),
width=width, height=height)
db_file = DBTelegramFile(
id=loc_id,
mxc=content_uri,
decryption_info=decryption_info,
mime_type=mime_type,
was_converted=image_converted,
timestamp=int(time.time()),
size=len(file),
width=width,
height=height,
)
if thumbnail and (mime_type.startswith("video/") or mime_type == "image/gif"):
if isinstance(thumbnail, (PhotoSize, PhotoCachedSize)):
thumbnail = thumbnail.location
try:
db_file.thumbnail = await transfer_thumbnail_to_matrix(client, intent, thumbnail,
video=file, mime_type=mime_type,
encrypt=encrypt)
db_file.thumbnail = await transfer_thumbnail_to_matrix(
client, intent, thumbnail, video=file, mime_type=mime_type, encrypt=encrypt
)
except FileIdInvalidError:
log.warning(f"Failed to transfer thumbnail for {thumbnail!s}", exc_info=True)
elif converted_anim and converted_anim.thumbnail_data:
db_file.thumbnail = await transfer_thumbnail_to_matrix(
client, intent, location, video=None, encrypt=encrypt,
custom_data=converted_anim.thumbnail_data, mime_type=converted_anim.thumbnail_mime,
width=converted_anim.width, height=converted_anim.height)
client,
intent,
location,
video=None,
encrypt=encrypt,
custom_data=converted_anim.thumbnail_data,
mime_type=converted_anim.thumbnail_mime,
width=converted_anim.width,
height=converted_anim.height,
)
try:
await db_file.insert()
except (UniqueViolationError, IntegrityError) as e:
log.exception(f"{e.__class__.__name__} while saving transferred file data. "
"This was probably caused by two simultaneous transfers of the same file, "
"and should not cause any problems.")
log.exception(
f"{e.__class__.__name__} while saving transferred file data. "
"This was probably caused by two simultaneous transfers of the same file, "
"and should not cause any problems."
)
return db_file