Nikita Kalanakov 2 år sedan
förälder
incheckning
106fc6816d
2 ändrade filer med 68 tillägg och 2 borttagningar
  1. 10 2
      commands.py
  2. 58 0
      utils.py

+ 10 - 2
commands.py

@@ -21,7 +21,15 @@ from actions import (
   add_or_update_birthday,
   get_birthdays
 )
-from utils import WORDS_TABLE, make_temporary_filename, parse_kind, get_user_name, get_word_for, calculate_age
+from utils import (
+  WORDS_TABLE,
+  make_temporary_filename,
+  parse_kind,
+  get_user_name,
+  get_word_for,
+  calculate_age,
+  unparse
+)
 
 class Handler:
   def __init__(self, handler, is_restricted=False):
@@ -183,7 +191,7 @@ async def make_message_shot(bot, message):
 
   data += pack('I', (message.sender.id if message.sender else message.peer_id.channel_id) % 7)
 
-  text = bytes(message.text, encoding='UTF-8')
+  text = bytes(unparse(message.raw_text, [entity for entity, _ in message.get_entities_text()]), encoding='UTF-8')
   data += pack('I', len(text))
   data += text
 

+ 58 - 0
utils.py

@@ -4,6 +4,7 @@ from datetime import datetime
 from collections import namedtuple
 
 from telethon.utils import get_display_name
+from telethon.tl.types import MessageEntityBold, MessageEntityItalic, MessageEntityStrike, MessageEntityCode
 
 Command = namedtuple(
   'Command', 
@@ -130,3 +131,60 @@ def calculate_age(date):
     age_now=age_now,
     date_string=date.strftime('%d.%m.%Y')
   )
+
+DELIMITERS = {
+  MessageEntityBold:   '**',
+  MessageEntityItalic: '__',
+  MessageEntityStrike: '~~',
+  MessageEntityCode:   '`',
+}
+
+class LookupTable:
+  def __init__(self):
+    self._start = {}
+    self._end = {}
+
+  def insert(self, entity):
+    delimiter = DELIMITERS.get(type(entity))
+    if not delimiter:
+      return
+
+    start = entity.offset
+    end = entity.offset + entity.length
+
+    if start in self._start:
+      self._start[start].append(delimiter)
+    else:
+      self._start[start] = [delimiter]
+
+    if end in self._end:
+      self._end[end].insert(0, delimiter)
+    else:
+      self._end[end] = [delimiter]
+
+  def _lookup(self, position):
+    if position in self._start:
+      return ''.join(self._start[position])
+    elif position in self._end:
+      return ''.join(self._end[position]) 
+
+  def process(self, text):
+    text = list(text) + ['']
+    result = ''
+
+    for position, character in zip(range(len(text)), text):
+      delimiter = self._lookup(position)
+      if delimiter:
+        result += delimiter
+
+      result += character
+
+    return result
+
+def unparse(text, entities):
+  table = LookupTable()
+
+  for entity in entities:
+    table.insert(entity)
+
+  return table.process(text)