hist_cmd: added module 'test_hist_cmd.py'

chatmastermind/commands/hist.py

@@ -1,13 +1,53 @@
+import sys
 import argparse
 from pathlib import Path
 from ..configuration import Config
 from ..chat import ChatDB
-from ..message import MessageFilter
+from ..message import MessageFilter, Message
 
 
-def hist_cmd(args: argparse.Namespace, config: Config) -> None:
+msg_suffix = Message.file_suffix_write  # currently '.msg'
+
+
+def convert_messages(args: argparse.Namespace, config: Config) -> None:
     """
-    Handler for the 'hist' command.
+    Convert messages to a new format. Also used to change old suffixes
+    ('.txt', '.yaml') to the latest default message file suffix ('.msg').
+    """
+    chat = ChatDB.from_dir(Path(config.cache),
+                           Path(config.db))
+    # read all known message files
+    msgs = chat.msg_gather(loc='disk', glob='*.*')
+    # make a set of all message IDs
+    msg_ids = set([m.msg_id() for m in msgs])
+    # set requested format and write all messages
+    chat.set_msg_format(args.convert)
+    # delete the current suffix
+    # -> a new one will automatically be created
+    for m in msgs:
+        if m.file_path:
+            m.file_path = m.file_path.with_suffix('')
+    chat.msg_write(msgs)
+    # read all messages with the current default suffix
+    msgs = chat.msg_gather(loc='disk', glob='*{msg_suffix}')
+    # make sure we converted all of the original messages
+    for mid in msg_ids:
+        if not any(mid == m.msg_id() for m in msgs):
+            print(f"Message '{mid}' has not been found after conversion. Aborting.")
+            sys.exit(1)
+    # delete messages with old suffixes
+    msgs = chat.msg_gather(loc='disk', glob='*.*')
+    for m in msgs:
+        if m.file_path and m.file_path.suffix != msg_suffix:
+            m.rm_file()
+    print(f"Successfully converted {len(msg_ids)} messages.")
+    if len(msgs):
+        print(f"Deleted {len(msgs)} messages with deprecated suffixes.")
+
+
+def print_chat(args: argparse.Namespace, config: Config) -> None:
+    """
+    Print the DB chat history.
     """
 
     mfilter = MessageFilter(tags_or=args.or_tags,
@@ -21,3 +61,13 @@ def hist_cmd(args: argparse.Namespace, config: Config) -> None:
     chat.print(args.source_code_only,
                args.with_tags,
                args.with_files)
+
+
+def hist_cmd(args: argparse.Namespace, config: Config) -> None:
+    """
+    Handler for the 'hist' command.
+    """
+    if args.print:
+        print_chat(args, config)
+    elif args.convert:
+        convert_messages(args, config)

chatmastermind/main.py

@@ -73,17 +73,20 @@ def create_parser() -> argparse.ArgumentParser:
 
     # 'hist' command parser
     hist_cmd_parser = cmdparser.add_parser('hist', parents=[tag_parser],
-                                           help="Print chat history.",
+                                           help="Print and manage chat history.",
                                            aliases=['h'])
     hist_cmd_parser.set_defaults(func=hist_cmd)
+    hist_group = hist_cmd_parser.add_mutually_exclusive_group(required=True)
+    hist_group.add_argument('-p', '--print', help='Print the DB chat history', action='store_true')
+    hist_group.add_argument('-c', '--convert', help='Convert all message files to the given format [txt|yaml]', metavar='FORMAT')
     hist_cmd_parser.add_argument('-w', '--with-tags', help="Print chat history with tags.",
                                  action='store_true')
     hist_cmd_parser.add_argument('-W', '--with-files', help="Print chat history with filenames.",
                                  action='store_true')
     hist_cmd_parser.add_argument('-S', '--source-code-only', help='Only print embedded source code',
                                  action='store_true')
-    hist_cmd_parser.add_argument('-A', '--answer', help='Search for answer substring', metavar='SUBSTRING')
-    hist_cmd_parser.add_argument('-Q', '--question', help='Search for question substring', metavar='SUBSTRING')
+    hist_cmd_parser.add_argument('-A', '--answer', help='Print only answers with given substring', metavar='SUBSTRING')
+    hist_cmd_parser.add_argument('-Q', '--question', help='Print only questions with given substring', metavar='SUBSTRING')
 
     # 'tags' command parser
     tags_cmd_parser = cmdparser.add_parser('tags',

tests/test_hist_cmd.py

@@ -0,0 +1,57 @@
+import unittest
+import argparse
+import tempfile
+import yaml
+from pathlib import Path
+from chatmastermind.message import Message, Question
+from chatmastermind.chat import ChatDB
+from chatmastermind.configuration import Config
+from chatmastermind.commands.hist import convert_messages
+
+
+msg_suffix = Message.file_suffix_write
+
+
+class TestConvertMessages(unittest.TestCase):
+    def setUp(self) -> None:
+        self.db_dir = tempfile.TemporaryDirectory()
+        self.cache_dir = tempfile.TemporaryDirectory()
+        self.db_path = Path(self.db_dir.name)
+        self.cache_path = Path(self.cache_dir.name)
+        self.args = argparse.Namespace()
+        self.config = Config()
+        self.config.cache = self.cache_dir.name
+        self.config.db = self.db_dir.name
+        self.args.convert = 'yaml'
+        # Prepare some messages
+        self.chat = ChatDB.from_dir(Path(self.cache_path),
+                                    Path(self.db_path))
+        self.messages = [Message(Question(f'Question {i}')) for i in range(1, 6)]
+        self.chat.db_write(self.messages[0:2])
+        self.chat.cache_write(self.messages[3:])
+        # Change some of the suffixes
+        assert self.messages[0].file_path
+        assert self.messages[1].file_path
+        self.messages[0].file_path.rename(self.messages[0].file_path.with_suffix('.txt'))
+        self.messages[1].file_path.rename(self.messages[1].file_path.with_suffix('.yaml'))
+
+    def tearDown(self) -> None:
+        self.db_dir.cleanup()
+        self.cache_dir.cleanup()
+
+    def test_convert_messages(self) -> None:
+        convert_messages(self.args, self.config)
+        msgs = self.chat.msg_gather(loc='disk', glob='*.*')
+        # Check if the number of messages is the same as before
+        self.assertEqual(len(msgs), len(self.messages))
+        # Check if all messages have the requested suffix
+        for msg in msgs:
+            assert msg.file_path
+            self.assertEqual(msg.file_path.suffix, msg_suffix)
+        # Check if the message IDs are correctly maintained
+        for m_new, m_old in zip(msgs, self.messages):
+            self.assertEqual(m_new.msg_id(), m_old.msg_id())
+        # check if all messages have the new format
+        for m in msgs:
+            with open(str(m.file_path), "r") as fd:
+                yaml.load(fd, Loader=yaml.FullLoader)