Migrate formatter and utils to mautrix-python
This commit is contained in:
Tulir Asokan
@@ -13,145 +13,84 @@
|
||||
#
|
||||
# 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 Callable, List, Optional, Sequence, Type, Union
|
||||
from typing import Optional, Union, Any, List, Type, Dict
|
||||
from enum import Enum
|
||||
|
||||
from telethon.tl.types import (MessageEntityMentionName as MentionName,
|
||||
MessageEntityTextUrl as TextURL, MessageEntityPre as Pre,
|
||||
TypeMessageEntity, InputMessageEntityMentionName as InputMentionName)
|
||||
from telethon.tl.types import (MessageEntityMention as Mention, MessageEntityBotCommand as Command,
|
||||
MessageEntityMentionName as MentionName, MessageEntityUrl as URL,
|
||||
MessageEntityEmail as Email, MessageEntityTextUrl as TextURL,
|
||||
MessageEntityBold as Bold, MessageEntityItalic as Italic,
|
||||
MessageEntityCode as Code, MessageEntityPre as Pre,
|
||||
MessageEntityStrike as Strike, MessageEntityUnderline as Underline,
|
||||
MessageEntityBlockquote as Blockquote, TypeMessageEntity,
|
||||
InputMessageEntityMentionName as InputMentionName)
|
||||
|
||||
from mautrix.util.formatter import EntityString, SemiAbstractEntity
|
||||
|
||||
|
||||
class Entity:
|
||||
@staticmethod
|
||||
def copy(entity: TypeMessageEntity) -> Optional[TypeMessageEntity]:
|
||||
if not entity:
|
||||
return None
|
||||
kwargs = {
|
||||
"offset": entity.offset,
|
||||
"length": entity.length,
|
||||
}
|
||||
if isinstance(entity, Pre):
|
||||
kwargs["language"] = entity.language
|
||||
elif isinstance(entity, TextURL):
|
||||
kwargs["url"] = entity.url
|
||||
elif isinstance(entity, (MentionName, InputMentionName)):
|
||||
kwargs["user_id"] = entity.user_id
|
||||
return entity.__class__(**kwargs)
|
||||
class TelegramEntityType(Enum):
|
||||
"""EntityType is a Matrix formatting entity type."""
|
||||
BOLD = Bold
|
||||
ITALIC = Italic
|
||||
STRIKETHROUGH = Strike
|
||||
UNDERLINE = Underline
|
||||
URL = URL
|
||||
INLINE_URL = TextURL
|
||||
EMAIL = Email
|
||||
PREFORMATTED = Pre
|
||||
INLINE_CODE = Code
|
||||
BLOCKQUOTE = Blockquote
|
||||
MENTION = Mention
|
||||
MENTION_NAME = MentionName
|
||||
COMMAND = Command
|
||||
|
||||
@classmethod
|
||||
def adjust(cls, entity: Union[TypeMessageEntity, List[TypeMessageEntity]],
|
||||
func: Callable[[TypeMessageEntity], None]
|
||||
) -> Union[Optional[TypeMessageEntity], List[TypeMessageEntity]]:
|
||||
if isinstance(entity, list):
|
||||
return [Entity.adjust(element, func) for element in entity if entity]
|
||||
elif not entity:
|
||||
return None
|
||||
entity = cls.copy(entity)
|
||||
func(entity)
|
||||
if entity.offset < 0:
|
||||
entity.length += entity.offset
|
||||
entity.offset = 0
|
||||
return entity
|
||||
USER_MENTION = 1
|
||||
ROOM_MENTION = 2
|
||||
HEADER = 3
|
||||
|
||||
|
||||
def offset_diff(amount: int) -> Callable[[TypeMessageEntity], None]:
|
||||
def func(entity: TypeMessageEntity) -> None:
|
||||
entity.offset += amount
|
||||
class TelegramEntity(SemiAbstractEntity):
|
||||
internal: TypeMessageEntity
|
||||
|
||||
return func
|
||||
def __init__(self, type: Union[TelegramEntityType, Type[TypeMessageEntity]],
|
||||
offset: int, length: int, extra_info: Dict[str, Any]) -> None:
|
||||
if isinstance(type, TelegramEntityType):
|
||||
if isinstance(type.value, int):
|
||||
raise ValueError(f"Can't create Entity with non-Telegram EntityType {type}")
|
||||
type = type.value
|
||||
self.internal = type(offset=offset, length=length, **extra_info)
|
||||
|
||||
def copy(self) -> Optional['TelegramEntity']:
|
||||
extra_info = {}
|
||||
if isinstance(self.internal, Pre):
|
||||
extra_info["language"] = self.internal.language
|
||||
elif isinstance(self.internal, TextURL):
|
||||
extra_info["url"] = self.internal.url
|
||||
elif isinstance(self.internal, (MentionName, InputMentionName)):
|
||||
extra_info["user_id"] = self.internal.user_id
|
||||
return TelegramEntity(type(self.internal), offset=self.internal.offset,
|
||||
length=self.internal.length, extra_info=extra_info)
|
||||
|
||||
@property
|
||||
def offset(self) -> int:
|
||||
return self.internal.offset
|
||||
|
||||
@offset.setter
|
||||
def offset(self, value: int) -> None:
|
||||
self.internal.offset = value
|
||||
|
||||
@property
|
||||
def length(self) -> int:
|
||||
return self.internal.length
|
||||
|
||||
@length.setter
|
||||
def length(self, value: int) -> None:
|
||||
self.internal.length = value
|
||||
|
||||
|
||||
def offset_length_multiply(amount: int) -> Callable[[TypeMessageEntity], None]:
|
||||
def func(entity: TypeMessageEntity) -> None:
|
||||
entity.offset *= amount
|
||||
entity.length *= amount
|
||||
class TelegramMessage(EntityString[TelegramEntity, TelegramEntityType]):
|
||||
entity_class = TelegramEntity
|
||||
|
||||
return func
|
||||
|
||||
|
||||
class TelegramMessage:
|
||||
def __init__(self, text: str = "", entities: Optional[List[TypeMessageEntity]] = None) -> None:
|
||||
self.text = text # type: str
|
||||
self.entities = entities or [] # type: List[TypeMessageEntity]
|
||||
|
||||
def offset_entities(self, offset: int) -> 'TelegramMessage':
|
||||
def apply_offset(entity: TypeMessageEntity, inner_offset: int
|
||||
) -> Optional[TypeMessageEntity]:
|
||||
entity = Entity.copy(entity)
|
||||
entity.offset += inner_offset
|
||||
if entity.offset < 0:
|
||||
entity.offset = 0
|
||||
elif entity.offset > len(self.text):
|
||||
return None
|
||||
elif entity.offset + entity.length > len(self.text):
|
||||
entity.length = len(self.text) - entity.offset
|
||||
return entity
|
||||
|
||||
self.entities = [apply_offset(entity, offset) for entity in self.entities if entity]
|
||||
self.entities = [x for x in self.entities if x is not None]
|
||||
return self
|
||||
|
||||
def append(self, *args: Union[str, 'TelegramMessage']) -> 'TelegramMessage':
|
||||
for msg in args:
|
||||
if isinstance(msg, str):
|
||||
msg = TelegramMessage(text=msg)
|
||||
self.entities += Entity.adjust(msg.entities, offset_diff(len(self.text)))
|
||||
self.text += msg.text
|
||||
return self
|
||||
|
||||
def prepend(self, *args: Union[str, 'TelegramMessage']) -> 'TelegramMessage':
|
||||
for msg in args:
|
||||
if isinstance(msg, str):
|
||||
msg = TelegramMessage(text=msg)
|
||||
self.entities = msg.entities + Entity.adjust(self.entities, offset_diff(len(msg.text)))
|
||||
self.text = msg.text + self.text
|
||||
return self
|
||||
|
||||
def format(self, entity_type: Type[TypeMessageEntity], offset: int = None, length: int = None,
|
||||
**kwargs) -> 'TelegramMessage':
|
||||
self.entities.append(entity_type(offset=offset or 0,
|
||||
length=length if length is not None else len(self.text),
|
||||
**kwargs))
|
||||
return self
|
||||
|
||||
def concat(self, *args: Union[str, 'TelegramMessage']) -> 'TelegramMessage':
|
||||
return TelegramMessage().append(self, *args)
|
||||
|
||||
def trim(self) -> 'TelegramMessage':
|
||||
orig_len = len(self.text)
|
||||
self.text = self.text.lstrip()
|
||||
diff = orig_len - len(self.text)
|
||||
self.text = self.text.rstrip()
|
||||
self.offset_entities(-diff)
|
||||
return self
|
||||
|
||||
def split(self, separator, max_items: int = 0) -> List['TelegramMessage']:
|
||||
text_parts = self.text.split(separator, max_items - 1)
|
||||
output = [] # type: List[TelegramMessage]
|
||||
|
||||
offset = 0
|
||||
for part in text_parts:
|
||||
msg = TelegramMessage(part)
|
||||
for entity in self.entities:
|
||||
start_in_range = len(part) > entity.offset - offset >= 0
|
||||
end_in_range = len(part) >= entity.offset - offset + entity.length > 0
|
||||
if start_in_range and end_in_range:
|
||||
msg.entities.append(Entity.adjust(entity, offset_diff(-offset)))
|
||||
output.append(msg)
|
||||
|
||||
offset += len(part)
|
||||
offset += len(separator)
|
||||
|
||||
return output
|
||||
|
||||
@staticmethod
|
||||
def join(items: Sequence[Union[str, 'TelegramMessage']],
|
||||
separator: str = " ") -> 'TelegramMessage':
|
||||
main = TelegramMessage()
|
||||
for msg in items:
|
||||
if isinstance(msg, str):
|
||||
msg = TelegramMessage(text=msg)
|
||||
main.entities += Entity.adjust(msg.entities, offset_diff(len(main.text)))
|
||||
main.text += msg.text + separator
|
||||
if len(separator) > 0:
|
||||
main.text = main.text[:-len(separator)]
|
||||
return main
|
||||
@property
|
||||
def telegram_entities(self) -> List[TypeMessageEntity]:
|
||||
return [entity.internal for entity in self.entities]
|
||||
|
||||
Reference in New Issue
Block a user