Try to improve getting forwarded message source entity

This commit is contained in:
Tulir Asokan
2022-10-14 14:53:54 +03:00
parent f63a7857a6
commit 751f1d93f3
+40 -14
View File
@@ -23,6 +23,9 @@ from telethon.errors import RPCError
from telethon.helpers import add_surrogate, del_surrogate from telethon.helpers import add_surrogate, del_surrogate
from telethon.tl.custom import Message from telethon.tl.custom import Message
from telethon.tl.types import ( from telethon.tl.types import (
Channel,
InputPeerChannelFromMessage,
InputPeerUserFromMessage,
MessageEntityBlockquote, MessageEntityBlockquote,
MessageEntityBold, MessageEntityBold,
MessageEntityBotCommand, MessageEntityBotCommand,
@@ -47,6 +50,7 @@ from telethon.tl.types import (
PeerUser, PeerUser,
SponsoredMessage, SponsoredMessage,
TypeMessageEntity, TypeMessageEntity,
User,
) )
from mautrix.types import Format, MessageType, TextMessageEventContent from mautrix.types import Format, MessageType, TextMessageEventContent
@@ -60,9 +64,33 @@ from ..util.file_transfer import UnicodeCustomEmoji, transfer_custom_emojis_to_m
log: logging.Logger = logging.getLogger("mau.fmt.tg") log: logging.Logger = logging.getLogger("mau.fmt.tg")
async def _get_fwd_entity(client: MautrixTelegramClient, evt: Message) -> Channel | User | None:
try:
return await client.get_entity(evt.fwd_from.from_id)
except (ValueError, RPCError) as e:
try:
input_peer = await client.get_input_entity(evt.peer_id)
if isinstance(evt.fwd_from.from_id, PeerUser):
return await client.get_entity(
InputPeerUserFromMessage(
peer=input_peer, msg_id=evt.id, user_id=evt.fwd_from.from_id.user_id
)
)
elif isinstance(evt.fwd_from.from_id, PeerChannel):
return await client.get_entity(
InputPeerChannelFromMessage(
peer=input_peer, msg_id=evt.id, channel_id=evt.fwd_from.from_id.channel_id
)
)
except (ValueError, RPCError) as e:
pass
return None
async def _add_forward_header( async def _add_forward_header(
client: MautrixTelegramClient, content: TextMessageEventContent, fwd_from: MessageFwdHeader client: MautrixTelegramClient, content: TextMessageEventContent, evt: Message
) -> None: ) -> None:
fwd_from = evt.fwd_from
fwd_from_html, fwd_from_text = None, None fwd_from_html, fwd_from_text = None, None
if isinstance(fwd_from.from_id, PeerUser): if isinstance(fwd_from.from_id, PeerUser):
user = await u.User.get_by_tgid(TelegramID(fwd_from.from_id.user_id)) user = await u.User.get_by_tgid(TelegramID(fwd_from.from_id.user_id))
@@ -81,12 +109,11 @@ async def _add_forward_header(
) )
if not fwd_from_text: if not fwd_from_text:
try: user = await _get_fwd_entity(client, evt)
user = await client.get_entity(fwd_from.from_id) if user:
if user: fwd_from_text, _ = pu.Puppet.get_displayname(user, False)
fwd_from_text, _ = pu.Puppet.get_displayname(user, False) fwd_from_html = f"<b>{escape(fwd_from_text)}</b>"
fwd_from_html = f"<b>{escape(fwd_from_text)}</b>" else:
except (ValueError, RPCError):
fwd_from_text = fwd_from_html = "unknown user" fwd_from_text = fwd_from_html = "unknown user"
elif isinstance(fwd_from.from_id, (PeerChannel, PeerChat)): elif isinstance(fwd_from.from_id, (PeerChannel, PeerChat)):
from_id = ( from_id = (
@@ -104,12 +131,11 @@ async def _add_forward_header(
else: else:
fwd_from_html = f"channel <b>{escape(fwd_from_text)}</b>" fwd_from_html = f"channel <b>{escape(fwd_from_text)}</b>"
else: else:
try: channel = await _get_fwd_entity(client, evt)
channel = await client.get_entity(fwd_from.from_id) if channel:
if channel: fwd_from_text = f"channel {channel.title}"
fwd_from_text = f"channel {channel.title}" fwd_from_html = f"channel <b>{escape(channel.title)}</b>"
fwd_from_html = f"channel <b>{escape(channel.title)}</b>" else:
except (ValueError, RPCError):
fwd_from_text = fwd_from_html = "unknown channel" fwd_from_text = fwd_from_html = "unknown channel"
elif fwd_from.from_name: elif fwd_from.from_name:
fwd_from_text = fwd_from.from_name fwd_from_text = fwd_from.from_name
@@ -175,7 +201,7 @@ async def telegram_to_matrix(
content.ensure_has_html() content.ensure_has_html()
if getattr(evt, "fwd_from", None): if getattr(evt, "fwd_from", None):
await _add_forward_header(client, content, evt.fwd_from) await _add_forward_header(client, content, evt)
if isinstance(evt, Message) and evt.post and evt.post_author: if isinstance(evt, Message) and evt.post and evt.post_author:
content.ensure_has_html() content.ensure_has_html()