From 9e5cb8414050da32f477ded3a15ae262e0ce81b9 Mon Sep 17 00:00:00 2001 From: Tulir Asokan Date: Fri, 23 Feb 2018 21:24:18 +0200 Subject: [PATCH] Refactor more code --- mautrix_telegram/context.py | 1 - mautrix_telegram/portal.py | 36 ++++++++++++++++---------- mautrix_telegram/util/file_transfer.py | 10 +++++-- 3 files changed, 30 insertions(+), 17 deletions(-) diff --git a/mautrix_telegram/context.py b/mautrix_telegram/context.py index d53f7dc9..6960321c 100644 --- a/mautrix_telegram/context.py +++ b/mautrix_telegram/context.py @@ -31,4 +31,3 @@ class Context: yield self.config yield self.loop yield self.bot - # yield self.mx diff --git a/mautrix_telegram/portal.py b/mautrix_telegram/portal.py index 6997bf22..c43070c1 100644 --- a/mautrix_telegram/portal.py +++ b/mautrix_telegram/portal.py @@ -332,16 +332,19 @@ class Portal: user.register_portal(self) await self.main_intent.invite(self.mxid, user.mxid) - async def delete_telegram_user(self, user_id, kick_message=None): + async def delete_telegram_user(self, user_id, sender): puppet = p.Puppet.get(user_id) user = u.User.get_by_tgid(user_id) - if kick_message: + kick_message = (f"Kicked by {sender.displayname}" + if sender and sender.tgid != user.tgid + else "Left Telegram chat") + if sender and sender.tgid != user.tgid: await self.main_intent.kick(self.mxid, puppet.mxid, kick_message) else: await puppet.intent.leave_room(self.mxid) if user: user.unregister_portal(self) - await self.main_intent.kick(self.mxid, user.mxid, kick_message or "Left Telegram chat") + await self.main_intent.kick(self.mxid, user.mxid, kick_message) async def update_info(self, user, entity=None): if self.peer_type == "user": @@ -367,7 +370,7 @@ class Portal: if changed: self.save() - async def update_username(self, username): + async def update_username(self, username, save=False): if self.username != username: if self.username: await self.main_intent.remove_room_alias(self._get_alias_localpart()) @@ -377,20 +380,27 @@ class Portal: await self.main_intent.set_join_rule(self.mxid, "public") else: await self.main_intent.set_join_rule(self.mxid, "invite") + + if save: + self.save() return True return False - async def update_about(self, about): + async def update_about(self, about, save=False): if self.about != about: self.about = about await self.main_intent.set_room_topic(self.mxid, self.about) + if save: + self.save() return True return False - async def update_title(self, title): + async def update_title(self, title, save=False): if self.title != title: self.title = title await self.main_intent.set_room_name(self.mxid, self.title) + if save: + self.save() return True return False @@ -399,7 +409,7 @@ class Portal: return max(photo.sizes, key=(lambda photo2: ( len(photo2.bytes) if isinstance(photo2, PhotoCachedSize) else photo2.size))) - async def update_avatar(self, user, photo): + async def update_avatar(self, user, photo, save=False): photo_id = f"{photo.volume_id}-{photo.local_id}" if self.photo_id != photo_id: file = await util.transfer_file_to_matrix(self.db, user.client, self.main_intent, @@ -407,6 +417,8 @@ class Portal: if file: await self.main_intent.set_room_avatar(self.mxid, file.mxc) self.photo_id = photo_id + if save: + self.save() return True return False @@ -974,21 +986,17 @@ class Portal: # TODO figure out how to see changes to about text / channel username if isinstance(action, MessageActionChatEditTitle): - if await self.update_title(action.title): - self.save() + await self.update_title(action.title, save=True) elif isinstance(action, MessageActionChatEditPhoto): largest_size = self._get_largest_photo_size(action.photo) - if await self.update_avatar(source, largest_size.location): - self.save() + self.update_avatar(source, largest_size.location, save=True) elif isinstance(action, MessageActionChatAddUser): for user_id in action.users: await self.add_telegram_user(user_id, source) elif isinstance(action, MessageActionChatJoinedByLink): await self.add_telegram_user(sender.id, source) elif isinstance(action, MessageActionChatDeleteUser): - kick_message = (f"Kicked by {sender.displayname}" - if sender.id != action.user_id else None) - await self.delete_telegram_user(action.user_id, kick_message) + await self.delete_telegram_user(action.user_id, sender) elif isinstance(action, MessageActionChatMigrateTo): self.peer_type = "channel" self.migrate_and_save(action.channel_id) diff --git a/mautrix_telegram/util/file_transfer.py b/mautrix_telegram/util/file_transfer.py index 76e113b3..2fd686a0 100644 --- a/mautrix_telegram/util/file_transfer.py +++ b/mautrix_telegram/util/file_transfer.py @@ -20,6 +20,7 @@ import logging import magic from PIL import Image +from sqlalchemy.exc import IntegrityError from telethon.tl.types import (Document, FileLocation, InputFileLocation, InputDocumentFileLocation) @@ -69,7 +70,12 @@ async def transfer_file_to_matrix(db, client, intent, location): db_file = DBTelegramFile(id=id, mxc=uploaded["content_uri"], mime_type=mime_type, was_converted=image_converted, timestamp=int(time.time())) - db.add(db_file) - db.commit() + try: + db.add(db_file) + db.commit() + except IntegrityError: + log.exception("Integrity error 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