chat_message.py

# -*- coding: utf-8 -*-

import yaml


# Only export the class
__all__ = ['Message']

START_MARK = '{'
END_MARK = '}'

ITEM_TAG = 'item'
ORDER_TAG = 'order'
# hyperlink是一个单词，既然前面是tag，后面应该是tags？？
HYPERLINK_TAGS = [ORDER_TAG, ITEM_TAG]


def escape(message):
    return message.replace(START_MARK, START_MARK * 2)\
        .replace(END_MARK, END_MARK * 2)


class Message():
    def __init__(self):
        self.content = ''

    def to_yaml(self):
        return yaml.dump(self.__dict__)

    @staticmethod
    def from_yaml(dumped):
        message = Message()
        for k, v in yaml.safe_load(dumped).items():
            setattr(message, k, v)
        return message

    def add_plain_message(self, message):
        self.content += escape(message)

    def __add_hyperlink_message(self, message, hyperlink_type):
        if hyperlink_type not in HYPERLINK_TAGS:
            return

        self.content += START_MARK
        self.add_plain_message(hyperlink_type)
        self.content += END_MARK
        self.__dict__.setdefault(hyperlink_type, []).append(message)

    def add_item_message(self, message):
        self.__add_hyperlink_message(message, ITEM_TAG)

    def add_order_message(self, message):
        self.__add_hyperlink_message(message, ORDER_TAG)


def test_escape():
    expected = 'plus a weird me{{s}}}}s{{{{a}}ge.'
    assert escape('plus a weird me{s}}s{{a}ge.') == expected


def test_message():
    message = Message()
    keyboard_inputs = [
        'this is a message.',
        'item: here is an item.',
        'plus a weird me{s}}s{{a}ge.',
        'item: here is another item.',
        'order: here is an order.',
        'exit',
    ]
    for keyboard_input in keyboard_inputs:
        if keyboard_input == "exit":
            break
        elif 'item' in keyboard_input:
            message.add_item_message(keyboard_input)
        elif 'order' in keyboard_input:
            message.add_order_message(keyboard_input)
        else:
            message.add_plain_message(keyboard_input)

    expected_dump = '''
content: this is a message.{item}plus a weird me{{s}}}}s{{{{a}}ge.{item}{order}
item: ['item: here is an item.', 'item: here is another item.']
order: ['order: here is an order.']
'''.strip()
    assert message.to_yaml().strip() == expected_dump

    expected_dict = {
        'content': 'this is a message.{item}plus a weird '
                   'me{{s}}}}s{{{{a}}ge.{item}{order}',
        'item': ['item: here is an item.', 'item: here is another item.'],
        'order': ['order: here is an order.']}
    message = Message.from_yaml(expected_dump)
    assert message.__dict__ == expected_dict