Aligned the timestamp of MAM start/end messages, added info message if no more messages are left.

1 files changed, 14 insertions, 13 deletions

diff --git a/poezio/mam.py b/poezio/mam.py
index effad4c2..163e6439 100644
--- a/poezio/mam.py
+++ b/poezio/mam.py
@@ -18,14 +18,14 @@ def add_line(self, text_buffer: TextBuffer, text: str, str_time: str, nick: str,
 
     time = datetime.strftime(str_time, '%Y-%m-%d %H:%M:%S')
     time = datetime.strptime(time, '%Y-%m-%d %H:%M:%S')
+    time = time.replace(tzinfo=timezone.utc).astimezone(tz=None)
+    time = time.replace(tzinfo=None)
     if '/' in nick:
         if isinstance(self, tabs.MucTab):
             nick = nick.split('/')[1]
         else:
             nick = nick.split('/')[0]
         color = get_theme().COLOR_OWN_NICK
-        time = time.replace(tzinfo=timezone.utc).astimezone(tz=None)
-        time = time.replace(tzinfo=None)
     else:
         color = get_theme().COLOR_ME_MESSAGE
     text_buffer.add_message(
@@ -59,30 +59,28 @@ async def query(self, remote_jid, start, end, top):
                 results = self.core.xmpp['xep_0313'].retrieve(jid=self.remote_jid,
                 iterator=True, reverse=top, end=self.end_date)
             except (IqError, IqTimeout):
-                return self.information('Failed to retrieve messages', 'Error')
+                return self.core.information('Failed to retrieve messages', 'Error')
         else:
             try:
                 results = self.core.xmpp['xep_0313'].retrieve(with_jid=self.remote_jid,
                 iterator=True, reverse=top, end=self.end_date)
             except (IqError, IqTimeout):
-                return self.information('Failed to retrieve messages', 'Error')
+                return self.core.information('Failed to retrieve messages', 'Error')
     else:
         if isinstance(self, tabs.MucTab):
             try:
                 results = self.core.xmpp['xep_0313'].retrieve(jid=self.remote_jid,
                 iterator=True, reverse=top, start=self.start_date, end=self.end_date)
             except (IqError, IqTimeout):
-                return self.information('Failed to retrieve messages', 'Error')
+                return self.core.information('Failed to retrieve messages', 'Error')
         else:
             try:
                 results = self.core.xmpp['xep_0313'].retrieve(with_jid=self.remote_jid,
                 iterator=True, reverse=top, start=self.start_date, end=self.end_date)
             except (IqError, IqTimeout):
-                return self.information('Failed to retrieve messages', 'Error')
+                return self.core.information('Failed to retrieve messages', 'Error')
     msg_count = 0
     msgs = []
-    timestamp = datetime.now()
-    add_line(self, text_buffer, 'Start of MAM query: ', timestamp, 'MAM', top)
     async for rsm in results:
         if top:
             for msg in rsm['mam']['results']:
@@ -90,17 +88,21 @@ async def query(self, remote_jid, start, end, top):
                     msgs.append(msg)
                 if msg_count == 10:
                     self.query_id = 0
-                    timestamp = datetime.now()
-                    add_line(self, text_buffer, 'End of MAM query: ', timestamp, 'MAM', top)
                     self.core.refresh_window()
                     return
                 msg_count += 1
             msgs.reverse()
             for msg in msgs:
+                if msg is msgs[0]:
+                    timestamp = msg['mam_result']['forwarded']['delay']['stamp']
+                    add_line(self, text_buffer, 'Start of MAM query: ', timestamp, 'MAM', top)
                 forwarded = msg['mam_result']['forwarded']
                 timestamp = forwarded['delay']['stamp']
                 message = forwarded['stanza']
                 add_line(self, text_buffer, message['body'], timestamp, str(message['from']), top)
+                if msg is msgs[len(msgs)-1]:
+                    timestamp = msg['mam_result']['forwarded']['delay']['stamp']
+                    add_line(self, text_buffer, 'End of MAM query: ', timestamp, 'MAM', top)
                 self.text_win.scroll_up(len(self.text_win.built_lines))
         else:
             for msg in rsm['mam']['results']:
@@ -109,10 +111,9 @@ async def query(self, remote_jid, start, end, top):
                 message = forwarded['stanza']
                 add_line(self, text_buffer, message['body'], timestamp, str(message['from']), top)
                 self.core.refresh_window()
+    if len(msgs) == 0:
+        return self.core.information('No more messages left to retrieve', 'Info')
     self.query_id = 0
-    timestamp = datetime.now()
-    add_line(self, text_buffer, 'End of MAM query: ', timestamp, 'MAM', top)
-
 
 def mam_scroll(self):
     remote_jid = self.jid