From 63040b368895fb8065c0c03a15b0f40beb561339 Mon Sep 17 00:00:00 2001 From: juk0de Date: Mon, 4 Sep 2023 08:49:43 +0200 Subject: [PATCH] message / chat: output improvements --- chatmastermind/chat.py | 16 ++++------------ chatmastermind/message.py | 24 ++++++++++++++++++++++++ tests/test_chat.py | 16 ++++++++++++---- tests/test_message.py | 24 ++++++++++++++++++++++++ 4 files changed, 64 insertions(+), 16 deletions(-) diff --git a/chatmastermind/chat.py b/chatmastermind/chat.py index 7e6df8f..c631dab 100644 --- a/chatmastermind/chat.py +++ b/chatmastermind/chat.py @@ -7,7 +7,7 @@ from pprint import PrettyPrinter from pydoc import pager from dataclasses import dataclass from typing import TypeVar, Type, Optional, ClassVar, Any, Callable -from .message import Question, Answer, Message, MessageFilter, MessageError, source_code, message_in +from .message import Message, MessageFilter, MessageError, message_in from .tags import Tag ChatInst = TypeVar('ChatInst', bound='Chat') @@ -170,18 +170,10 @@ class Chat: output: list[str] = [] for message in self.messages: if source_code_only: - output.extend(source_code(message.question, include_delims=True)) + output.append(message.to_str(source_code_only=True)) continue - output.append('-' * terminal_width()) - if with_tags: - output.append(message.tags_str()) - if with_files: - output.append('FILE: ' + str(message.file_path)) - output.append(Question.txt_header) - output.append(message.question) - if message.answer: - output.append(Answer.txt_header) - output.append(message.answer) + output.append(message.to_str(with_tags, with_files)) + output.append('\n' + ('-' * terminal_width()) + '\n') if paged: print_paged('\n'.join(output)) else: diff --git a/chatmastermind/message.py b/chatmastermind/message.py index 0fb949c..35de3b9 100644 --- a/chatmastermind/message.py +++ b/chatmastermind/message.py @@ -392,6 +392,30 @@ class Message(): data[cls.file_yaml_key] = file_path return cls.from_dict(data) + def to_str(self, with_tags: bool = False, with_file: bool = False, source_code_only: bool = False) -> str: + """ + Return the current Message as a string. + """ + output: list[str] = [] + if source_code_only: + # use the source code from answer only + if self.answer: + output.extend(self.answer.source_code(include_delims=True)) + return '\n'.join(output) if len(output) > 0 else '' + if with_tags: + output.append(self.tags_str()) + if with_file: + output.append('FILE: ' + str(self.file_path)) + output.append(Question.txt_header) + output.append(self.question) + if self.answer: + output.append(Answer.txt_header) + output.append(self.answer) + return '\n'.join(output) + + def __str__(self) -> str: + return self.to_str(False, False, False) + def to_file(self, file_path: Optional[pathlib.Path]=None) -> None: # noqa: 11 """ Write a Message to the given file. Type is determined based on the suffix. diff --git a/tests/test_chat.py b/tests/test_chat.py index a1c020e..f8302eb 100644 --- a/tests/test_chat.py +++ b/tests/test_chat.py @@ -66,16 +66,20 @@ class TestChat(CmmTestCase): def test_print(self, mock_stdout: StringIO) -> None: self.chat.add_messages([self.message1, self.message2]) self.chat.print(paged=False) - expected_output = f"""{'-'*terminal_width()} -{Question.txt_header} + expected_output = f"""{Question.txt_header} Question 1 {Answer.txt_header} Answer 1 + {'-'*terminal_width()} + {Question.txt_header} Question 2 {Answer.txt_header} Answer 2 + +{'-'*terminal_width()} + """ self.assertEqual(mock_stdout.getvalue(), expected_output) @@ -83,20 +87,24 @@ Answer 2 def test_print_with_tags_and_file(self, mock_stdout: StringIO) -> None: self.chat.add_messages([self.message1, self.message2]) self.chat.print(paged=False, with_tags=True, with_files=True) - expected_output = f"""{'-'*terminal_width()} -{TagLine.prefix} atag1 btag2 + expected_output = f"""{TagLine.prefix} atag1 btag2 FILE: 0001.txt {Question.txt_header} Question 1 {Answer.txt_header} Answer 1 + {'-'*terminal_width()} + {TagLine.prefix} btag2 FILE: 0002.txt {Question.txt_header} Question 2 {Answer.txt_header} Answer 2 + +{'-'*terminal_width()} + """ self.assertEqual(mock_stdout.getvalue(), expected_output) diff --git a/tests/test_message.py b/tests/test_message.py index e860538..a49c893 100644 --- a/tests/test_message.py +++ b/tests/test_message.py @@ -804,3 +804,27 @@ class MessageRenameTagsTestCase(CmmTestCase): self.message.rename_tags({(Tag('atag1'), Tag('atag2')), (Tag('btag2'), Tag('btag3'))}) self.assertIsNotNone(self.message.tags) self.assertSetEqual(self.message.tags, {Tag('atag2'), Tag('btag3')}) # type: ignore [arg-type] + + +class MessageToStrTestCase(CmmTestCase): + def setUp(self) -> None: + self.message = Message(Question('This is a question.'), + Answer('This is an answer.'), + tags={Tag('atag1'), Tag('btag2')}, + file_path=pathlib.Path('/tmp/foo/bla')) + + def test_to_str(self) -> None: + expected_output = f"""{Question.txt_header} +This is a question. +{Answer.txt_header} +This is an answer.""" + self.assertEqual(self.message.to_str(), expected_output) + + def test_to_str_with_tags_and_file(self) -> None: + expected_output = f"""{TagLine.prefix} atag1 btag2 +FILE: /tmp/foo/bla +{Question.txt_header} +This is a question. +{Answer.txt_header} +This is an answer.""" + self.assertEqual(self.message.to_str(with_tags=True, with_file=True), expected_output)