added new module 'openai.py'

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:

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:
+            output.extend(self.question.source_code(include_delims=True))
+            if self.answer:
+                output.extend(self.answer.source_code(include_delims=True))
+            return '\n'.join(output)
+        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()}
-{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)
 
@@ -388,4 +396,25 @@ class TestChatDB(CmmTestCase):
         with self.assertRaises(ChatError):
             chat_db.add_to_cache([Message(Question("?"), file_path=pathlib.Path("foo"))])
 
-    # TODO: add testcase for "ChatDB.write_messages()"
+    def test_chat_db_write_messages(self) -> None:
+        # create a new ChatDB instance
+        chat_db = ChatDB.from_dir(pathlib.Path(self.cache_path.name),
+                                  pathlib.Path(self.db_path.name))
+
+        db_dir_files = self.message_list(self.db_path)
+        self.assertEqual(len(db_dir_files), 4)
+        cache_dir_files = self.message_list(self.cache_path)
+        self.assertEqual(len(cache_dir_files), 0)
+
+        # try to write a message without a valid file_path
+        message = Message(question=Question("Question 1"),
+                          answer=Answer("Answer 1"))
+        with self.assertRaises(ChatError):
+            chat_db.write_messages([message])
+
+        # write a message with a valid file_path
+        message.file_path = pathlib.Path(self.cache_path.name) / '123456.txt'
+        chat_db.write_messages([message])
+        cache_dir_files = self.message_list(self.cache_path)
+        self.assertEqual(len(cache_dir_files), 1)
+        self.assertIn(pathlib.Path(self.cache_path.name, '123456.txt'), cache_dir_files)

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)