path: root/gentree.py

from enum import Enum
import sys
import pprint
import textwrap

class LogType(Enum):
    QUASSEL = 1
    ERC = 2


class Message:
    def __init__(self):
        self.id = -1
        self.sender = ""
        self.message = ""
        self.reply_to_msgs = []
        self.previous_msg = -1

    def to_dict(self):
        return {'id': self.id,
                'sender': self.sender,
                'message': self.message,
                'reply_to_msgs': self.reply_to_msgs,
                'previous_msg': self.previous_msg}


def get_log_type(log):
    """determine the type of irc log.
    args:
      log: a string of irc log.
    returns:
      a LogType.
    """
    if log[0] == '[':
        return LogType.QUASSEL
    elif log[0] == '<':
        return LogType.ERC
    else:
        raise ValueError('Unknown irc log type.')
    

def parse_log_quassel(log):
    """parse a quassel log.
    A quassel message is in the form of
    [hh:mm:ss] <sender-nick> blahblah.
    or
    [hh:mm:ss] -*- sender-nick does something.
    [hh:mm:ss] *** Mode #abc +o def by ChanServ
    args:
      log: a string of quassel log.
    return:
      a list of Messages
    """
    last_messages = dict()
    index = 0
    messages = []
    for line in log.splitlines():
        message = Message()
        line = line.lstrip()
        [_, nick, payload] = line.split(' ', maxsplit=2)
        if nick == '***':
            continue
        if nick == '-*-':
            [_, _, nick, payload] = line.split(' ', maxsplit=3)
        else:
            nick = nick[1:-1] # removes <>
        message.message = line
        message.id = index
        # Remove special symbol for op
        if nick[0] in ['@', '%']:
            nick = nick[1:]
        message.sender = nick
        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)
        messages.append(message)
        index += 1
    return messages

def parse_log_erc(log):
    """parse an erc log.
    An erc message is in the form of
    <nick> blahblah
           blahblah blah  [hh:mm]
    ** nick does something
    **** old-nick is now known as new-nick
    And the timestamp is optional
    args:
      log: a string of quassel log.
    return:
      a list of Messages
    """
    last_messages = dict()
    index = 0
    messages = []
    log_lines = log.splitlines()
    log_lines.append('')
    nick = ''
    for line, next_line in zip(log_lines, log_lines[1:]):
        line = line.lstrip()
        if line[0] == '<':
            [nick, payload] = line.split(' ', maxsplit=1)
        elif line[:3] = '** ':
            [_, nick, payload] = line.split(' ', maxsplit=2)
        elif line[:5] = '**** ':
            continue
        else:
            payload += f'\n{line}'
        if nick == '***':
            continue
        if nick == '-*-':
            [_, _, nick, payload] = line.split(' ', maxsplit=3)
        else:
            nick = nick[1:-1] # removes <>
        message.message = line
        message.id = index
        # Remove special symbol for op
        if nick[0] in ['@', '%']:
            nick = nick[1:]
        message.sender = nick
        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)
        messages.append(message)
        index += 1
    return messages

def parse_log(log):
    """parse log.
    args:
      log: a string of irc log.
    returns:
      a list of Messages
    """
    log_type = get_log_type(log)
    if log_type == LogType.QUASSEL:
        return parse_log_quassel(log)
    elif log_type == LogType.ERC:
        return parse_log_erc(log)
    else:
        raise ValueError('Unknown irc log type.')

def gen_graph(meeting):
    dot = """
digraph meeting {
  graph [
    width = 20
  ];

  node [
    shape = box,
    margin = 0,
    pad = 0
  ];
"""
    for msg in meeting:
        quoted_message = textwrap.fill(msg.message.replace('"', '\\"'))
        dot += f'\nm{str(msg.id)} [label="{quoted_message}"];'
        if msg.id > 0:
            dot += f'\nm{str(msg.id - 1)} -> m{str(msg.id)} [style="invis"];'
    for msg in meeting:
        if msg.previous_msg > -1:
            dot += f'\nm{str(msg.previous_msg)} -> m{str(msg.id)} [style="dashed"];'
        for id in msg.reply_to_msgs:
            dot += f'\nm{str(id)} -> m{str(msg.id)};'
    dot += '\n}'
    return dot

def main():
    log = open(sys.argv[1]).read()
    meeting = parse_log(log)
    graph = gen_graph(meeting)
    open(sys.argv[2], 'w').write(graph)
    # pp = pprint.PrettyPrinter(indent=4)
    # pp.pprint([msg.to_dict() for msg in meeting])

if __name__ == '__main__':
    main()