chat: new possibilites for adding messages and better tests

chatmastermind/chat.py

@@ -55,6 +55,16 @@ def read_dir(dir_path: pathlib.Path,
     return messages
 
 
+def make_file_path(dir_path: pathlib.Path,
+                   file_suffix: str,
+                   next_fid: Callable[[], int]) -> pathlib.Path:
+    """
+    Create a file_path for the given directory using the
+    given file_suffix and ID generator function.
+    """
+    return dir_path / f"{next_fid():04d}{file_suffix}"
+
+
 def write_dir(dir_path: pathlib.Path,
               messages: list[Message],
               file_suffix: str,
@@ -73,9 +83,7 @@ def write_dir(dir_path: pathlib.Path,
         file_path = message.file_path
         # message has no file_path: create one
         if not file_path:
-            fid = next_fid()
-            fname = f"{fid:04d}{file_suffix}"
-            file_path = dir_path / fname
+            file_path = make_file_path(dir_path, file_suffix, next_fid)
         # file_path does not point to given directory: modify it
         elif not file_path.parent.samefile(dir_path):
             file_path = dir_path / file_path.name
@@ -124,11 +132,11 @@ class Chat:
         """
         self.messages = []
 
-    def add_msgs(self, msgs: list[Message]) -> None:
+    def add_messages(self, messages: list[Message]) -> None:
         """
         Add new messages and sort them if possible.
         """
-        self.messages += msgs
+        self.messages += messages
         self.sort()
 
     def tags(self, prefix: Optional[str] = None, contain: Optional[str] = None) -> set[Tag]:
@@ -279,25 +287,25 @@ class ChatDB(Chat):
         self.messages += new_messages
         self.sort()
 
-    def write_db(self, msgs: Optional[list[Message]] = None) -> None:
+    def write_db(self, messages: Optional[list[Message]] = None) -> None:
         """
         Write messages to the DB directory. If a message has no file_path, a new one
         will be created. If message.file_path exists, it will be modified to point
         to the DB directory.
         """
         write_dir(self.db_path,
-                  msgs if msgs else self.messages,
+                  messages if messages else self.messages,
                   self.file_suffix,
                   self.get_next_fid)
 
-    def write_cache(self, msgs: Optional[list[Message]] = None) -> None:
+    def write_cache(self, messages: Optional[list[Message]] = None) -> None:
         """
         Write messages to the cache directory. If a message has no file_path, a new one
         will be created. If message.file_path exists, it will be modified to point to
         the cache directory.
         """
         write_dir(self.cache_path,
-                  msgs if msgs else self.messages,
+                  messages if messages else self.messages,
                   self.file_suffix,
                   self.get_next_fid)
 
@@ -309,3 +317,52 @@ class ChatDB(Chat):
         clear_dir(self.cache_path, self.glob)
         # only keep messages from DB dir (or those that have not yet been written)
         self.messages = [m for m in self.messages if not m.file_path or m.file_path.parent.samefile(self.db_path)]
+
+    def add_to_db(self, messages: list[Message], write: bool = True) -> None:
+        """
+        Add the given new messages and set the file_path to the DB directory.
+        Only accepts messages without a file_path.
+        """
+        if any(m.file_path is not None for m in messages):
+            raise ChatError("Can't add new messages with existing file_path")
+        if write:
+            write_dir(self.db_path,
+                      messages,
+                      self.file_suffix,
+                      self.get_next_fid)
+        else:
+            for m in messages:
+                m.file_path = make_file_path(self.db_path, self.default_file_suffix, self.get_next_fid)
+        self.messages += messages
+        self.sort()
+
+    def add_to_cache(self, messages: list[Message], write: bool = True) -> None:
+        """
+        Add the given new messages and set the file_path to the cache directory.
+        Only accepts messages without a file_path.
+        """
+        if any(m.file_path is not None for m in messages):
+            raise ChatError("Can't add new messages with existing file_path")
+        if write:
+            write_dir(self.cache_path,
+                      messages,
+                      self.file_suffix,
+                      self.get_next_fid)
+        else:
+            for m in messages:
+                m.file_path = make_file_path(self.cache_path, self.default_file_suffix, self.get_next_fid)
+        self.messages += messages
+        self.sort()
+
+    def write_messages(self, messages: Optional[list[Message]] = None) -> None:
+        """
+        Write either the given messages or the internal ones to their current file_path.
+        If messages are given, they all must have a valid file_path. When writing the
+        internal messages, the ones with a valid file_path are written, the others
+        are ignored.
+        """
+        if messages and any(m.file_path is None for m in messages):
+            raise ChatError("Can't write files without a valid file_path")
+        msgs = iter(messages if messages else self.messages)
+        while (m := next(msgs, None)):
+            m.to_file()