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] = None
+    answer: Optional[Answer]
-    tags: Optional[set[Tag]] = None
+    tags: Optional[set[Tag]]
-    ai: Optional[str] = None
+    ai: Optional[str]
-    model: Optional[str] = None
+    model: Optional[str]
-    file_path: Optional[pathlib.Path] = None
+    file_path: Optional[pathlib.Path]
     # class variables
     file_suffixes: ClassVar[list[str]] = ['.txt', '.yaml']
     tags_yaml_key: ClassVar[str] = 'tags'
@@ -213,31 +213,12 @@ class Message():
         return tags
 
     @classmethod
-    def from_file(cls: Type[MessageInst], file_path: pathlib.Path,
+    def from_file(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 file. Returns 'None' if the message does
+        Create a Message from the given file. Expects the following file structures:
-        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]
@@ -245,80 +226,90 @@ class Message():
             * Question.txt_header
             * Question
             * Answer.txt_header [Optional]
-            * Answer [Optional]
+            # Answer [Optional]
-
+        For '.yaml':
-        Returns 'None' if the message does not fulfill the tag requirements.
-        """
-        tags: set[Tag] = set()
-        question: Question
-        answer: Optional[Answer] = None
-        ai: Optional[str] = None
-        model: Optional[str] = None
-        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.
         """
+        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()
-        with open(file_path, "r") as fd:
+        question: Question
-            data = yaml.load(fd, Loader=yaml.FullLoader)
+        answer: Optional[Answer] = None
-            if tags_or or tags_and or tags_not:
+        ai: Optional[str] = None
-                if Message.tags_yaml_key in data:
+        model: Optional[str] = None
-                    tags = set([Tag(tag) for tag in data[Message.tags_yaml_key]])
+        # TXT
-                # match with an empty set if the file has no tags
+        if file_path.suffix == '.txt':
-                if not match_tags(tags, tags_or, tags_and, tags_not):
+            with open(file_path, "r") as fd:
-                    return None
+                # TagLine (Optional)
-            data[cls.file_yaml_key] = file_path
+                try:
-            return cls.from_dict(data)
+                    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)
 
     def to_file(self, file_path: Optional[pathlib.Path]) -> None:  # noqa: 11
         """
-        Write a Message to the given file. Type is determined based on the suffix.
+        Write a Message to the given file. Creates the following file structures:
-        Currently supported suffixes: ['.txt', '.yaml']
+        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]
         """
         if file_path:
             self.file_path = file_path
@@ -328,54 +319,28 @@ class Message():
             raise MessageError(f"File type '{self.file_path.suffix}' is not supported")
         # TXT
         if self.file_path.suffix == '.txt':
-            return self.__to_file_txt(self.file_path)
+            with open(self.file_path, "w") as fd:
-        elif self.file_path.suffix == '.yaml':
+                msg_tags = self.tags or set()
-            return self.__to_file_yaml(self.file_path)
+                fd.write(f'{TagLine.from_set(msg_tags)}\n')
-
+                if self.ai:
-    def __to_file_txt(self, file_path: pathlib.Path) -> None:
+                    fd.write(f'{AILine.from_ai(self.ai)}\n')
-        """
+                if self.model:
-        Write a Message to the given file in TXT format.
+                    fd.write(f'{ModelLine.from_model(self.model)}\n')
-        Creates the following file structures:
+                fd.write(f'{Question.txt_header}\n{self.question}\n')
-            * 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')
-
+        # YAML
-    def __to_file_yaml(self, file_path: pathlib.Path) -> None:
+        elif self.file_path.suffix == '.yaml':
-        """
+            with open(self.file_path, "w") as fd:
-        Write a Message to the given file in YAML format.
+                data: YamlDict = {Question.yaml_key: str(self.question)}
-        Creates the following file structures:
+                if self.answer:
-            * Question.yaml_key: single or multiline string
+                    data[Answer.yaml_key] = str(self.answer)
-            * Answer.yaml_key: single or multiline string
+                if self.ai:
-            * Message.tags_yaml_key: list of strings
+                    data[self.ai_yaml_key] = self.ai
-            * Message.ai_yaml_key: str [Optional]
+                if self.model:
-            * Message.model_yaml_key: str [Optional]
+                    data[self.model_yaml_key] = self.model
-        """
+                if self.tags:
-        with open(file_path, "w") as fd:
+                    data[self.tags_yaml_key] = sorted([str(tag) for tag in self.tags])
-            data: YamlDict = {Question.yaml_key: str(self.question)}
+                yaml.dump(data, fd, sort_keys=False)
-            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,21 +86,19 @@ 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_complete = Message(Question('This is a question.'),
+        self.message = Message(Question('This is a question.'),
-                                        Answer('This is an answer.'),
+                               Answer('This is an answer.'),
-                                        {Tag('tag1'), Tag('tag2')},
+                               {Tag('tag1'), Tag('tag2')},
-                                        ai='ChatGPT',
+                               ai='ChatGPT',
-                                        model='gpt-3.5-turbo',
+                               model='gpt-3.5-turbo',
-                                        file_path=self.file_path)
+                               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_complete(self) -> None:
+    def test_to_file_txt(self) -> None:
-        self.message_complete.to_file(self.file_path)
+        self.message.to_file(self.file_path)
 
         with open(self.file_path, "r") as fd:
             content = fd.read()
@@ -111,23 +109,13 @@ 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_complete.to_file(unsupported_file_path)
+            self.message.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:
@@ -136,38 +124,36 @@ This is a question.
         """
         with self.assertRaises(MessageError) as cm:
             # clear the internal file_path
-            self.message_complete.file_path = None
+            self.message.file_path = None
-            self.message_complete.to_file(None)
+            self.message.to_file(None)
         self.assertEqual(str(cm.exception), "Got no valid path to write message")
         # reset the internal file_path
-        self.message_complete.file_path = self.file_path
+        self.message.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_complete = Message(Question('This is a question.'),
+        self.message = Message(Question('This is a question.'),
-                                        Answer('This is an answer.'),
+                               Answer('This is an answer.'),
-                                        {Tag('tag1'), Tag('tag2')},
+                               {Tag('tag1'), Tag('tag2')},
-                                        ai='ChatGPT',
+                               ai='ChatGPT',
-                                        model='gpt-3.5-turbo',
+                               model='gpt-3.5-turbo',
-                                        file_path=self.file_path)
+                               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_complete(self) -> None:
+    def test_to_file_yaml(self) -> None:
-        self.message_complete.to_file(self.file_path)
+        self.message.to_file(self.file_path)
 
         with open(self.file_path, "r") as fd:
             content = fd.read()
@@ -200,14 +186,6 @@ 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:
@@ -404,12 +382,6 @@ 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')})