tests: added testcases for Message.from_file() and Message.to_file()

chatmastermind/message.py

@@ -170,11 +170,11 @@ class Message():
     and a file path.
     """
     question: Question
-    answer: Optional[Answer]
-    tags: Optional[set[Tag]]
-    ai: Optional[str]
-    model: Optional[str]
-    file_path: Optional[pathlib.Path]
+    answer: Optional[Answer] = None
+    tags: Optional[set[Tag]] = None
+    ai: Optional[str] = None
+    model: Optional[str] = None
+    file_path: Optional[pathlib.Path] = None
     # class variables
     file_suffixes: ClassVar[list[str]] = ['.txt', '.yaml']
     tags_yaml_key: ClassVar[str] = 'tags'
@@ -213,12 +213,31 @@ class Message():
         return tags
 
     @classmethod
-    def from_file(cls: Type[MessageInst], file_path: pathlib.Path,  # noqa: 11
+    def from_file(cls: Type[MessageInst], file_path: pathlib.Path,
                   tags_or: Optional[set[Tag]] = None,
                   tags_and: Optional[set[Tag]] = None,
                   tags_not: Optional[set[Tag]] = None) -> Optional[MessageInst]:
         """
-        Create a Message from the given file. Expects the following file structures:
+        Create a Message from the given file. Returns 'None' if the message does
+        not fulfill the tag requirements.
+        """
+        if not file_path.exists():
+            raise MessageError(f"Message file '{file_path}' does not exist")
+        if file_path.suffix not in cls.file_suffixes:
+            raise MessageError(f"File type '{file_path.suffix}' is not supported")
+
+        if file_path.suffix == '.txt':
+            return cls.__from_file_txt(file_path, tags_or, tags_and, tags_not)
+        else:
+            return cls.__from_file_yaml(file_path, tags_or, tags_and, tags_not)
+
+    @classmethod
+    def __from_file_txt(cls: Type[MessageInst], file_path: pathlib.Path,  # noqa: 11
+                        tags_or: Optional[set[Tag]] = None,
+                        tags_and: Optional[set[Tag]] = None,
+                        tags_not: Optional[set[Tag]] = None) -> Optional[MessageInst]:
+        """
+        Create a Message from the given TXT file. Expects the following file structures:
         For '.txt':
             * TagLine [Optional]
             * AI [Optional]
@@ -226,90 +245,80 @@ class Message():
             * Question.txt_header
             * Question
             * Answer.txt_header [Optional]
-            # Answer [Optional]
-        For '.yaml':
-            * Question.yaml_key: single or multiline string
-            * Answer.yaml_key: single or multiline string [Optional]
-            * Message.tags_yaml_key: list of strings [Optional]
-            * Message.ai_yaml_key: str [Optional]
-            * Message.model_yaml_key: str [Optional]
+            * Answer [Optional]
+
         Returns 'None' if the message does not fulfill the tag requirements.
         """
-        if not file_path.exists():
-            raise MessageError(f"Message file '{file_path}' does not exist")
-        if file_path.suffix not in cls.file_suffixes:
-            raise MessageError(f"File type '{file_path.suffix}' is not supported")
-
         tags: set[Tag] = set()
         question: Question
         answer: Optional[Answer] = None
         ai: Optional[str] = None
         model: Optional[str] = None
-        # TXT
-        if file_path.suffix == '.txt':
-            with open(file_path, "r") as fd:
-                # TagLine (Optional)
-                try:
-                    pos = fd.tell()
-                    tags = TagLine(fd.readline()).tags()
-                except TagError:
-                    fd.seek(pos)
-                if tags_or or tags_and or tags_not:
-                    # match with an empty set if the file has no tags
-                    if not match_tags(tags, tags_or, tags_and, tags_not):
-                        return None
-                # AILine (Optional)
-                try:
-                    pos = fd.tell()
-                    ai = AILine(fd.readline()).ai()
-                except TagError:
-                    fd.seek(pos)
-                # ModelLine (Optional)
-                try:
-                    pos = fd.tell()
-                    model = ModelLine(fd.readline()).model()
-                except TagError:
-                    fd.seek(pos)
-                # Question and Answer
-                text = fd.read().strip().split('\n')
-                question_idx = text.index(Question.txt_header) + 1
-                try:
-                    answer_idx = text.index(Answer.txt_header)
-                    question = Question.from_list(text[question_idx:answer_idx])
-                    answer = Answer.from_list(text[answer_idx + 1:])
-                except ValueError:
-                    question = Question.from_list(text[question_idx:])
-                return cls(question, answer, tags, ai, model, file_path)
-        # YAML
-        else:
-            with open(file_path, "r") as fd:
-                data = yaml.load(fd, Loader=yaml.FullLoader)
-                if tags_or or tags_and or tags_not:
-                    if Message.tags_yaml_key in data:
-                        tags = set([Tag(tag) for tag in data[Message.tags_yaml_key]])
-                    # match with an empty set if the file has no tags
-                    if not match_tags(tags, tags_or, tags_and, tags_not):
-                        return None
-                data[cls.file_yaml_key] = file_path
-                return cls.from_dict(data)
+        with open(file_path, "r") as fd:
+            # TagLine (Optional)
+            try:
+                pos = fd.tell()
+                tags = TagLine(fd.readline()).tags()
+            except TagError:
+                fd.seek(pos)
+            if tags_or or tags_and or tags_not:
+                # match with an empty set if the file has no tags
+                if not match_tags(tags, tags_or, tags_and, tags_not):
+                    return None
+            # AILine (Optional)
+            try:
+                pos = fd.tell()
+                ai = AILine(fd.readline()).ai()
+            except TagError:
+                fd.seek(pos)
+            # ModelLine (Optional)
+            try:
+                pos = fd.tell()
+                model = ModelLine(fd.readline()).model()
+            except TagError:
+                fd.seek(pos)
+            # Question and Answer
+            text = fd.read().strip().split('\n')
+            question_idx = text.index(Question.txt_header) + 1
+            try:
+                answer_idx = text.index(Answer.txt_header)
+                question = Question.from_list(text[question_idx:answer_idx])
+                answer = Answer.from_list(text[answer_idx + 1:])
+            except ValueError:
+                question = Question.from_list(text[question_idx:])
+            return cls(question, answer, tags, ai, model, file_path)
+
+    @classmethod
+    def __from_file_yaml(cls: Type[MessageInst], file_path: pathlib.Path,
+                         tags_or: Optional[set[Tag]] = None,
+                         tags_and: Optional[set[Tag]] = None,
+                         tags_not: Optional[set[Tag]] = None) -> Optional[MessageInst]:
+        """
+        Create a Message from the given YAML file. Expects the following file structures:
+            * Question.yaml_key: single or multiline string
+            * Answer.yaml_key: single or multiline string [Optional]
+            * Message.tags_yaml_key: list of strings [Optional]
+            * Message.ai_yaml_key: str [Optional]
+            * Message.model_yaml_key: str [Optional]
+
+        Returns 'None' if the message does not fulfill the tag requirements.
+        """
+        tags: set[Tag] = set()
+        with open(file_path, "r") as fd:
+            data = yaml.load(fd, Loader=yaml.FullLoader)
+            if tags_or or tags_and or tags_not:
+                if Message.tags_yaml_key in data:
+                    tags = set([Tag(tag) for tag in data[Message.tags_yaml_key]])
+                # match with an empty set if the file has no tags
+                if not match_tags(tags, tags_or, tags_and, tags_not):
+                    return None
+            data[cls.file_yaml_key] = file_path
+            return cls.from_dict(data)
 
     def to_file(self, file_path: Optional[pathlib.Path]) -> None:  # noqa: 11
         """
-        Write a Message to the given file. Creates the following file structures:
-        For '.txt':
-            * TagLine
-            * AI [Optional]
-            * Model [Optional]
-            * Question.txt_header
-            * Question
-            * Answer.txt_header
-            * Answer
-        For '.yaml':
-            * Question.yaml_key: single or multiline string
-            * Answer.yaml_key: single or multiline string
-            * Message.tags_yaml_key: list of strings
-            * Message.ai_yaml_key: str [Optional]
-            * Message.model_yaml_key: str [Optional]
+        Write a Message to the given file. Type is determined based on the suffix.
+        Currently supported suffixes: ['.txt', '.yaml']
         """
         if file_path:
             self.file_path = file_path
@@ -319,28 +328,54 @@ class Message():
             raise MessageError(f"File type '{self.file_path.suffix}' is not supported")
         # TXT
         if self.file_path.suffix == '.txt':
-            with open(self.file_path, "w") as fd:
-                msg_tags = self.tags or set()
-                fd.write(f'{TagLine.from_set(msg_tags)}\n')
-                if self.ai:
-                    fd.write(f'{AILine.from_ai(self.ai)}\n')
-                if self.model:
-                    fd.write(f'{ModelLine.from_model(self.model)}\n')
-                fd.write(f'{Question.txt_header}\n{self.question}\n')
-                fd.write(f'{Answer.txt_header}\n{self.answer}\n')
-        # YAML
+            return self.__to_file_txt(self.file_path)
         elif self.file_path.suffix == '.yaml':
-            with open(self.file_path, "w") as fd:
-                data: YamlDict = {Question.yaml_key: str(self.question)}
-                if self.answer:
-                    data[Answer.yaml_key] = str(self.answer)
-                if self.ai:
-                    data[self.ai_yaml_key] = self.ai
-                if self.model:
-                    data[self.model_yaml_key] = self.model
-                if self.tags:
-                    data[self.tags_yaml_key] = sorted([str(tag) for tag in self.tags])
-                yaml.dump(data, fd, sort_keys=False)
+            return self.__to_file_yaml(self.file_path)
+
+    def __to_file_txt(self, file_path: pathlib.Path) -> None:
+        """
+        Write a Message to the given file in TXT format.
+        Creates the following file structures:
+            * TagLine
+            * AI [Optional]
+            * Model [Optional]
+            * Question.txt_header
+            * Question
+            * Answer.txt_header
+            * Answer
+        """
+        with open(file_path, "w") as fd:
+            if self.tags:
+                fd.write(f'{TagLine.from_set(self.tags)}\n')
+            if self.ai:
+                fd.write(f'{AILine.from_ai(self.ai)}\n')
+            if self.model:
+                fd.write(f'{ModelLine.from_model(self.model)}\n')
+            fd.write(f'{Question.txt_header}\n{self.question}\n')
+            if self.answer:
+                fd.write(f'{Answer.txt_header}\n{self.answer}\n')
+
+    def __to_file_yaml(self, file_path: pathlib.Path) -> None:
+        """
+        Write a Message to the given file in YAML format.
+        Creates the following file structures:
+            * Question.yaml_key: single or multiline string
+            * Answer.yaml_key: single or multiline string
+            * Message.tags_yaml_key: list of strings
+            * Message.ai_yaml_key: str [Optional]
+            * Message.model_yaml_key: str [Optional]
+        """
+        with open(file_path, "w") as fd:
+            data: YamlDict = {Question.yaml_key: str(self.question)}
+            if self.answer:
+                data[Answer.yaml_key] = str(self.answer)
+            if self.ai:
+                data[self.ai_yaml_key] = self.ai
+            if self.model:
+                data[self.model_yaml_key] = self.model
+            if self.tags:
+                data[self.tags_yaml_key] = sorted([str(tag) for tag in self.tags])
+            yaml.dump(data, fd, sort_keys=False)
 
     def as_dict(self) -> dict[str, Any]:
         return asdict(self)

tests/test_message.py

@@ -86,19 +86,21 @@ class MessageToFileTxtTestCase(CmmTestCase):
     def setUp(self) -> None:
         self.file = tempfile.NamedTemporaryFile(delete=False, suffix='.txt')
         self.file_path = pathlib.Path(self.file.name)
-        self.message = Message(Question('This is a question.'),
-                               Answer('This is an answer.'),
-                               {Tag('tag1'), Tag('tag2')},
-                               ai='ChatGPT',
-                               model='gpt-3.5-turbo',
-                               file_path=self.file_path)
+        self.message_complete = Message(Question('This is a question.'),
+                                        Answer('This is an answer.'),
+                                        {Tag('tag1'), Tag('tag2')},
+                                        ai='ChatGPT',
+                                        model='gpt-3.5-turbo',
+                                        file_path=self.file_path)
+        self.message_min = Message(Question('This is a question.'),
+                                   file_path=self.file_path)
 
     def tearDown(self) -> None:
         self.file.close()
         self.file_path.unlink()
 
-    def test_to_file_txt(self) -> None:
-        self.message.to_file(self.file_path)
+    def test_to_file_txt_complete(self) -> None:
+        self.message_complete.to_file(self.file_path)
 
         with open(self.file_path, "r") as fd:
             content = fd.read()
@@ -109,13 +111,23 @@ class MessageToFileTxtTestCase(CmmTestCase):
 This is a question.
 {Answer.txt_header}
 This is an answer.
+"""
+        self.assertEqual(content, expected_content)
+
+    def test_to_file_txt_min(self) -> None:
+        self.message_min.to_file(self.file_path)
+
+        with open(self.file_path, "r") as fd:
+            content = fd.read()
+        expected_content = f"""{Question.txt_header}
+This is a question.
 """
         self.assertEqual(content, expected_content)
 
     def test_to_file_unsupported_file_type(self) -> None:
         unsupported_file_path = pathlib.Path("example.doc")
         with self.assertRaises(MessageError) as cm:
-            self.message.to_file(unsupported_file_path)
+            self.message_complete.to_file(unsupported_file_path)
         self.assertEqual(str(cm.exception), "File type '.doc' is not supported")
 
     def test_to_file_no_file_path(self) -> None:
@@ -124,36 +136,38 @@ This is an answer.
         """
         with self.assertRaises(MessageError) as cm:
             # clear the internal file_path
-            self.message.file_path = None
-            self.message.to_file(None)
+            self.message_complete.file_path = None
+            self.message_complete.to_file(None)
         self.assertEqual(str(cm.exception), "Got no valid path to write message")
         # reset the internal file_path
-        self.message.file_path = self.file_path
+        self.message_complete.file_path = self.file_path
 
 
 class MessageToFileYamlTestCase(CmmTestCase):
     def setUp(self) -> None:
         self.file = tempfile.NamedTemporaryFile(delete=False, suffix='.yaml')
         self.file_path = pathlib.Path(self.file.name)
-        self.message = Message(Question('This is a question.'),
-                               Answer('This is an answer.'),
-                               {Tag('tag1'), Tag('tag2')},
-                               ai='ChatGPT',
-                               model='gpt-3.5-turbo',
-                               file_path=self.file_path)
+        self.message_complete = Message(Question('This is a question.'),
+                                        Answer('This is an answer.'),
+                                        {Tag('tag1'), Tag('tag2')},
+                                        ai='ChatGPT',
+                                        model='gpt-3.5-turbo',
+                                        file_path=self.file_path)
         self.message_multiline = Message(Question('This is a\nmultiline question.'),
                                          Answer('This is a\nmultiline answer.'),
                                          {Tag('tag1'), Tag('tag2')},
                                          ai='ChatGPT',
                                          model='gpt-3.5-turbo',
                                          file_path=self.file_path)
+        self.message_min = Message(Question('This is a question.'),
+                                   file_path=self.file_path)
 
     def tearDown(self) -> None:
         self.file.close()
         self.file_path.unlink()
 
-    def test_to_file_yaml(self) -> None:
-        self.message.to_file(self.file_path)
+    def test_to_file_yaml_complete(self) -> None:
+        self.message_complete.to_file(self.file_path)
 
         with open(self.file_path, "r") as fd:
             content = fd.read()
@@ -186,6 +200,14 @@ class MessageToFileYamlTestCase(CmmTestCase):
 """
         self.assertEqual(content, expected_content)
 
+    def test_to_file_yaml_min(self) -> None:
+        self.message_min.to_file(self.file_path)
+
+        with open(self.file_path, "r") as fd:
+            content = fd.read()
+        expected_content = f"{Question.yaml_key}: This is a question.\n"
+        self.assertEqual(content, expected_content)
+
 
 class MessageFromFileTxtTestCase(CmmTestCase):
     def setUp(self) -> None:
@@ -382,6 +404,12 @@ This is an answer.
   - tag2
 """)
 
+    def tearDown(self) -> None:
+        self.file_txt.close()
+        self.file_path_txt.unlink()
+        self.file_yaml.close()
+        self.file_path_yaml.unlink()
+
     def test_tags_from_file_txt(self) -> None:
         tags = Message.tags_from_file(self.file_path_txt)
         self.assertSetEqual(tags, {Tag('tag1'), Tag('tag2')})