message / chat: output improvements

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())
+            output.append(message.to_str(with_tags, with_files))
-            if with_tags:
+            output.append('\n' + ('-' * terminal_width()) + '\n')
-                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)
         if paged:
             print_paged('\n'.join(output))
         else:

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.

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()}
+        expected_output = f"""{Question.txt_header}
-{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()}
+        expected_output = f"""{TagLine.prefix} atag1 btag2
-{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)
 

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)