fixed some bugs

1 files changed, 13 insertions, 10 deletions

diff --git a/gentree.py b/gentree.py
index 995fa1a..fe980bd 100644
--- a/gentree.py
+++ b/gentree.py
@@ -14,11 +14,11 @@ class Message:
         self.reply_to_msgs = []
         self.previous_msg = -1
 
-def to_dict(message):
-    return {'sender': message.sender,
-            'message': message.message,
-            'reply_to_msgs': message.reply_to_msgs,
-            'previous_msg': message.previous_msg}
+    def to_dict(self):
+        return {'sender': self.sender,
+                'message': self.message,
+                'reply_to_msgs': self.reply_to_msgs,
+                'previous_msg': self.previous_msg}
 
 
 def get_log_type(log):
@@ -42,6 +42,7 @@ def parse_log_quassel(log):
     [hh:mm:ss] <sender-nick> blahblah.
     or
     [hh:mm:ss] -*- sender-nick does something.
+    [06:05:28] *** Mode #fsf +o zoe by ChanServ
     args:
       log: a string of quassel log.
     return:
@@ -54,6 +55,8 @@ def parse_log_quassel(log):
         message = Message()
         is_me = False
         [_, nick, payload] = line.split(' ', maxsplit=2)
+        if nick == '***':
+            continue
         if nick == '-*-':
             [_, _, nick, payload] = line.split(' ', maxsplit=3)
             is_me = True
@@ -61,11 +64,11 @@ def parse_log_quassel(log):
             nick = nick[1:-1] # removes <>
         # Remove special symbol for op
         if nick[0] in ['@', '%']:
-            nick = nick[:-1]
+            nick = nick[1:]
         message.sender = nick
-        last_message = last_messages.setdefault(nick, index)
-        if last_message != index:
-            message.previous_msg = last_message
+        if nick in last_messages:
+            message.previous_msg = last_messages[nick]
+        last_messages[nick] = index
         for participant, msg_id in last_messages.items():
             if payload.find(participant) != -1:
                 message.reply_to_msgs.append(msg_id)
@@ -108,7 +111,7 @@ def main():
     meeting = parse_log(log)
     # graph = gen_graph(meeting)
     pp = pprint.PrettyPrinter(indent=4)
-    pp.pprint([to_dict(msg) for msg in meeting])
+    pp.pprint([msg.to_dict() for msg in meeting])
 
 if __name__ == '__main__':
     main()