From 18e184d1f009b731c7501c510a117e6ed40ef8c5 Mon Sep 17 00:00:00 2001 From: miruka Date: Wed, 6 Nov 2019 09:43:05 -0400 Subject: [PATCH] Cache local media on upload, fix media local echo --- src/python/matrix_client.py | 38 ++++++++++++++++++++++++++----------- src/python/media_cache.py | 11 +++++++++++ src/python/models/items.py | 20 +++++++++++++------ 3 files changed, 52 insertions(+), 17 deletions(-) diff --git a/src/python/matrix_client.py b/src/python/matrix_client.py index 31be6522..a7db2fbe 100644 --- a/src/python/matrix_client.py +++ b/src/python/matrix_client.py @@ -237,8 +237,13 @@ class MatrixClient(nio.AsyncClient): path, upload_item, encrypt=encrypt, ) + await self.media_cache.create_media(url, path.read_bytes()) + kind = (mime or "").split("/")[0] + thumb_url: str = "" + thumb_info: Dict[str, Any] = {} + content: dict = { "body": path.name, "info": { @@ -266,13 +271,20 @@ class MatrixClient(nio.AsyncClient): ) try: - thumb_url, thumb_info, thumb_crypt_dict = \ + thumb_data, thumb_url, thumb_info, thumb_crypt_dict = \ await self.upload_thumbnail( path, upload_item, is_svg=is_svg, encrypt=encrypt, ) except (UneededThumbnail, UnthumbnailableError): pass else: + await self.media_cache.create_thumbnail( + thumb_url, + thumb_data, + content["info"]["w"], + content["info"]["h"], + ) + if encrypt: content["info"]["thumbnail_file"] = { "url": thumb_url, @@ -323,14 +335,17 @@ class MatrixClient(nio.AsyncClient): await self._local_echo( room_id, uuid, event_type, - inline_content = path.name, - media_url = url, - media_title = path.name, - media_width = content["info"].get("w", 0), - media_height = content["info"].get("h", 0), - media_duration = content["info"].get("duration", 0), - media_size = content["info"]["size"], - media_mime = content["info"]["mimetype"], + inline_content = path.name, + media_url = url, + media_title = path.name, + media_width = content["info"].get("w", 0), + media_height = content["info"].get("h", 0), + media_duration = content["info"].get("duration", 0), + media_size = content["info"]["size"], + media_mime = content["info"]["mimetype"], + thumbnail_url = thumb_url, + thumbnail_width = thumb_info.get("w", 0), + thumbnail_height = thumb_info.get("h", 0), ) await self._send_message(room_id, uuid, content) @@ -352,7 +367,7 @@ class MatrixClient(nio.AsyncClient): sender_name = our_info.display_name, sender_avatar = our_info.avatar_url, is_local_echo = True, - local_event_type = event_type.__name__, + local_event_type = event_type, **event_fields, ) @@ -457,7 +472,7 @@ class MatrixClient(nio.AsyncClient): item: Optional[Upload] = None, is_svg: bool = False, encrypt: bool = False, - ) -> Tuple[str, Dict[str, Any], CryptDict]: + ) -> Tuple[bytes, str, Dict[str, Any], CryptDict]: png_modes = ("1", "L", "P", "RGBA") @@ -511,6 +526,7 @@ class MatrixClient(nio.AsyncClient): item.status = UploadStatus.UploadingThumbnail return ( + data, await self.upload(data, upload_mime, Path(path).name), { "w": thumb.width, diff --git a/src/python/media_cache.py b/src/python/media_cache.py index 3b194d9a..ed485eb9 100644 --- a/src/python/media_cache.py +++ b/src/python/media_cache.py @@ -219,10 +219,21 @@ class MediaCache: self.downloads_dir.mkdir(parents=True, exist_ok=True) + async def create_media(self, mxc: str, data: bytes) -> None: + await Media(self, mxc, data, {}).create() + + async def get_media(self, mxc: str, crypt_dict: CryptDict = None) -> str: return str(await Media(self, mxc, None, crypt_dict).get()) + async def create_thumbnail( + self, mxc: str, data: bytes, width: int, height: int, + ) -> None: + thumb = Thumbnail(self, mxc, data, {}, (round(width), round(height))) + await thumb.create() + + async def get_thumbnail( self, mxc: str, width: int, height: int, crypt_dict: CryptDict = None, ) -> str: diff --git a/src/python/models/items.py b/src/python/models/items.py index 92e50f50..102f7882 100644 --- a/src/python/models/items.py +++ b/src/python/models/items.py @@ -2,7 +2,7 @@ import re from dataclasses import dataclass, field from datetime import datetime from pathlib import Path -from typing import Any, Dict, List, Optional, Tuple +from typing import Any, Dict, List, Optional, Tuple, Type from uuid import uuid4 import lxml # nosec @@ -159,8 +159,8 @@ class Event(ModelItem): target_name: str = "" target_avatar: str = "" - is_local_echo: bool = False - local_event_type: str = "" + is_local_echo: bool = False + local_event_type: Optional[Type[nio.Event]] = None media_url: str = "" media_title: str = "" @@ -187,12 +187,20 @@ class Event(ModelItem): @property def event_type(self) -> str: - return self.local_event_type or type(self.source).__name__ + if self.local_event_type: + return self.local_event_type.__name__ + + return type(self.source).__name__ @property def links(self) -> List[str]: - if isinstance(self.source, - (nio.RoomMessageMedia, nio.RoomEncryptedMedia)): + local_type = self.local_event_type + media_classes = (nio.RoomMessageMedia, nio.RoomEncryptedMedia) + + if local_type and issubclass(local_type, media_classes): + return [self.media_url] + + if isinstance(self.source, media_classes): return [self.media_url] if not self.content.strip():