From 1cd52acf2d012f095e493ceb714b4ed7316c350e Mon Sep 17 00:00:00 2001 From: juk0de Date: Sun, 24 Sep 2023 18:20:38 +0200 Subject: [PATCH] message: introduced file suffix '.msg' - '.msg' suffix is always used for writing - 'Message.to_file()' will set the file suffix if the given file_path has none - added 'mformat' argument to 'Message.to_file()' for choosing the file format - '.txt' and '.yaml' suffixes are only supported for reading --- chatmastermind/message.py | 75 +++++++++++++++++++++------------------ 1 file changed, 41 insertions(+), 34 deletions(-) diff --git a/chatmastermind/message.py b/chatmastermind/message.py index 443455e..04def6d 100644 --- a/chatmastermind/message.py +++ b/chatmastermind/message.py @@ -5,7 +5,8 @@ import pathlib import yaml import tempfile import shutil -from typing import Type, TypeVar, ClassVar, Optional, Any, Union, Final, Literal, Iterable +from typing import Type, TypeVar, ClassVar, Optional, Any, Union, Final, Literal, Iterable, Tuple +from typing import get_args as typing_get_args from dataclasses import dataclass, asdict, field from .tags import Tag, TagLine, TagError, match_tags, rename_tags @@ -15,6 +16,9 @@ MessageInst = TypeVar('MessageInst', bound='Message') AILineInst = TypeVar('AILineInst', bound='AILine') ModelLineInst = TypeVar('ModelLineInst', bound='ModelLine') YamlDict = dict[str, Union[QuestionInst, AnswerInst, set[Tag]]] +MessageFormat = Literal['txt', 'yaml'] +message_valid_formats: Final[Tuple[MessageFormat, ...]] = typing_get_args(MessageFormat) +message_default_format: Final[MessageFormat] = 'txt' class MessageError(Exception): @@ -92,7 +96,7 @@ class MessageFilter: class AILine(str): """ - A line that represents the AI name in a '.txt' file.. + A line that represents the AI name in the 'txt' format. """ prefix: Final[str] = 'AI:' @@ -112,7 +116,7 @@ class AILine(str): class ModelLine(str): """ - A line that represents the model name in a '.txt' file.. + A line that represents the model name in the 'txt' format. """ prefix: Final[str] = 'MODEL:' @@ -216,7 +220,9 @@ class Message(): model: Optional[str] = field(default=None, compare=False) file_path: Optional[pathlib.Path] = field(default=None, compare=False) # class variables - file_suffixes: ClassVar[list[str]] = ['.txt', '.yaml'] + file_suffixes_read: ClassVar[list[str]] = ['.msg', '.txt', '.yaml'] + file_suffix_write: ClassVar[str] = '.msg' + default_format: ClassVar[MessageFormat] = message_default_format tags_yaml_key: ClassVar[str] = 'tags' file_yaml_key: ClassVar[str] = 'file_path' ai_yaml_key: ClassVar[str] = 'ai' @@ -276,24 +282,16 @@ class Message(): tags: set[Tag] = set() if not file_path.exists(): raise MessageError(f"Message file '{file_path}' does not exist") - if file_path.suffix not in cls.file_suffixes: + if file_path.suffix not in cls.file_suffixes_read: raise MessageError(f"File type '{file_path.suffix}' is not supported") - # for TXT, it's enough to read the TagLine - if file_path.suffix == '.txt': - with open(file_path, "r") as fd: - try: - tags = TagLine(fd.readline()).tags(prefix, contain) - except TagError: - pass # message without tags - else: # '.yaml' - try: - message = cls.from_file(file_path) - if message: - msg_tags = message.filter_tags(prefix=prefix, contain=contain) - except MessageError as e: - print(f"Error processing message in '{file_path}': {str(e)}") - if msg_tags: - tags = msg_tags + try: + message = cls.from_file(file_path) + if message: + msg_tags = message.filter_tags(prefix=prefix, contain=contain) + except MessageError as e: + print(f"Error processing message in '{file_path}': {str(e)}") + if msg_tags: + tags = msg_tags return tags @classmethod @@ -328,15 +326,16 @@ class Message(): """ if not file_path.exists(): raise MessageError(f"Message file '{file_path}' does not exist") - if file_path.suffix not in cls.file_suffixes: + if file_path.suffix not in cls.file_suffixes_read: raise MessageError(f"File type '{file_path.suffix}' is not supported") - - if file_path.suffix == '.txt': + # try TXT first + try: message = cls.__from_file_txt(file_path, mfilter.tags_or if mfilter else None, mfilter.tags_and if mfilter else None, mfilter.tags_not if mfilter else None) - else: + # then YAML + except MessageError: message = cls.__from_file_yaml(file_path) if message and (mfilter is None or message.match(mfilter)): return message @@ -442,21 +441,29 @@ class Message(): output.append(self.answer) return '\n'.join(output) - def to_file(self, file_path: Optional[pathlib.Path]=None) -> None: # noqa: 11 + def to_file(self, file_path: Optional[pathlib.Path]=None, mformat: MessageFormat = message_default_format) -> None: # noqa: 11 """ - Write a Message to the given file. Type is determined based on the suffix. - Currently supported suffixes: ['.txt', '.yaml'] + Write a Message to the given file. Supported message file formats are 'txt' and 'yaml'. + Suffix is always '.msg'. """ if file_path: self.file_path = file_path if not self.file_path: raise MessageError("Got no valid path to write message") - if self.file_path.suffix not in self.file_suffixes: - raise MessageError(f"File type '{self.file_path.suffix}' is not supported") + if mformat not in message_valid_formats: + raise MessageError(f"File format '{mformat}' is not supported") + # check for valid suffix + # -> add one if it's empty + # -> refuse old or otherwise unsupported suffixes + if not self.file_path.suffix: + self.file_path = self.file_path.with_suffix(self.file_suffix_write) + elif self.file_path.suffix != self.file_suffix_write: + raise MessageError(f"File suffix '{self.file_path.suffix}' is not supported") # TXT - if self.file_path.suffix == '.txt': + if mformat == 'txt': return self.__to_file_txt(self.file_path) - elif self.file_path.suffix == '.yaml': + # YAML + elif mformat == 'yaml': return self.__to_file_yaml(self.file_path) def __to_file_txt(self, file_path: pathlib.Path) -> None: @@ -468,8 +475,8 @@ class Message(): * Model [Optional] * Question.txt_header * Question - * Answer.txt_header - * Answer + * Answer.txt_header [Optional] + * Answer [Optional] """ with tempfile.NamedTemporaryFile(dir=file_path.parent, prefix=file_path.name, mode="w", delete=False) as temp_fd: temp_file_path = pathlib.Path(temp_fd.name)