Added a check for tabs (because there is a different way to query messages for MUC and any other type of tab)

1 files changed, 23 insertions, 11 deletions

diff --git a/poezio/mam.py b/poezio/mam.py
index 28173994..b3145295 100644
--- a/poezio/mam.py
+++ b/poezio/mam.py
@@ -9,15 +9,19 @@
 import asyncio
 from datetime import datetime, timedelta, timezone
 from poezio.theming import get_theme
+from poezio import tabs
 from poezio.text_buffer import Message, TextBuffer
 
-def add_line(text_buffer: TextBuffer, text: str, str_time: str, nick: str, top: bool):
+def add_line(self, text_buffer: TextBuffer, text: str, str_time: str, nick: str, top: bool):
     """Adds a textual entry in the TextBuffer"""
 
     time = datetime.strftime(str_time, '%Y-%m-%d %H:%M:%S')
     time = datetime.strptime(time, '%Y-%m-%d %H:%M:%S')
     if '/' in nick:
-        nick = nick.split('/')[1]
+        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)
@@ -46,15 +50,23 @@ async def query(self, remote_jid, start, end, top):
     if 'urn:xmpp:mam:2' not in iq['disco_info'].get_features():
         return self.core.information("This MUC doesn't support MAM.", "Error")
     if top:
-        results = self.core.xmpp['xep_0313'].retrieve(jid=self.remote_jid,
-        iterator=True, reverse=top, end=self.end_date)
+        if isinstance(self, tabs.MucTab):
+            results = self.core.xmpp['xep_0313'].retrieve(jid=self.remote_jid,
+            iterator=True, reverse=top, end=self.end_date)
+        else:
+            results = self.core.xmpp['xep_0313'].retrieve(with_jid=self.remote_jid,
+            iterator=True, reverse=top, end=self.end_date)
     else:
-        results = self.core.xmpp['xep_0313'].retrieve(jid=self.remote_jid,
-        iterator=True, reverse=top, start=self.start_date, end=self.end_date)
+        if isinstance(self, tabs.MucTab):
+            results = self.core.xmpp['xep_0313'].retrieve(jid=self.remote_jid,
+            iterator=True, reverse=top, start=self.start_date, end=self.end_date)
+        else:
+            results = self.core.xmpp['xep_0313'].retrieve(with_jid=self.remote_jid,
+            iterator=True, reverse=top, start=self.start_date, end=self.end_date)
     msg_count = 0
     msgs = []
     timestamp = datetime.now()
-    add_line(text_buffer, 'Start of MAM query: ', timestamp, 'MAM', top)
+    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']:
@@ -62,7 +74,7 @@ async def query(self, remote_jid, start, end, top):
                 if msg_count == 10:
                     self.query_id = 0
                     timestamp = datetime.now()
-                    add_line(text_buffer, 'End of MAM query: ', timestamp, 'MAM', top)
+                    add_line(self, text_buffer, 'End of MAM query: ', timestamp, 'MAM', top)
                     self.core.refresh_window()
                     return
                 msg_count += 1
@@ -71,18 +83,18 @@ async def query(self, remote_jid, start, end, top):
                 forwarded = msg['mam_result']['forwarded']
                 timestamp = forwarded['delay']['stamp']
                 message = forwarded['stanza']
-                add_line(text_buffer, message['body'], timestamp, str(message['from']), top)
+                add_line(self, text_buffer, message['body'], timestamp, str(message['from']), top)
                 self.text_win.scroll_up(len(self.text_win.built_lines))
         else:
             for msg in rsm['mam']['results']:
                 forwarded = msg['mam_result']['forwarded']
                 timestamp = forwarded['delay']['stamp']
                 message = forwarded['stanza']
-                add_line(text_buffer, message['body'], timestamp, str(message['from']), top)
+                add_line(self, text_buffer, message['body'], timestamp, str(message['from']), top)
                 self.core.refresh_window()
     self.query_id = 0
     timestamp = datetime.now()
-    add_line(text_buffer, 'End of MAM query: ', timestamp, 'MAM', top)
+    add_line(self, text_buffer, 'End of MAM query: ', timestamp, 'MAM', top)
 
 
 def mam_scroll(self):