chatmastermind/main.py

@@ -2,17 +2,16 @@
 # -*- coding: utf-8 -*-
 # vim: set fileencoding=utf-8 :
 
-import yaml
 import sys
 import argcomplete
 import argparse
 from pathlib import Path
-from .utils import terminal_width, print_tag_args, print_chat_hist, display_source_code, ChatType
-from .storage import save_answers, create_chat_hist, read_file, dump_data
+from .utils import terminal_width, print_tag_args, print_chat_hist, ChatType
+from .storage import save_answers, create_chat_hist
 from .api_client import ai, openai_api_key, print_models
 from .configuration import Config
 from .chat import ChatDB
-from .message import Message, MessageFilter
+from .message import Message, MessageFilter, MessageError
 from itertools import zip_longest
 from typing import Any
 
@@ -128,13 +127,12 @@ def print_cmd(args: argparse.Namespace, config: Config) -> None:
     Handler for the 'print' command.
     """
     fname = Path(args.file)
-    if fname.suffix == '.yaml':
-        with open(args.file, 'r') as f:
-            data = yaml.load(f, Loader=yaml.FullLoader)
-    elif fname.suffix == '.txt':
-        data = read_file(fname)
-    else:
-        print(f"Unknown file type: {args.file}")
+    try:
+        message = Message.from_file(fname)
+        if message:
+            print(message.to_str(source_code_only=args.source_code_only))
+    except MessageError:
+        print(f"File is not a valid message: {args.file}")
         sys.exit(1)
     if args.source_code_only:
         display_source_code(data['answer'])
@@ -227,14 +225,22 @@ def create_parser() -> argparse.ArgumentParser:
 
     # 'print' command parser
     print_cmd_parser = cmdparser.add_parser('print',
-                                            help="Print files.",
+                                            help="Print message files.",
                                             aliases=['p'])
     print_cmd_parser.set_defaults(func=print_cmd)
     print_cmd_parser.add_argument('-f', '--file', help='File to print', required=True)
+<<<<<<< HEAD
     print_cmd_modes = print_cmd_parser.add_mutually_exclusive_group()
     print_cmd_modes.add_argument('-q', '--question', help='Print only question', action='store_true')
     print_cmd_modes.add_argument('-a', '--answer', help='Print only answer', action='store_true')
     print_cmd_modes.add_argument('-S', '--only-source-code', help='Print only source code', action='store_true')
+||||||| parent of bf1cbff (cmm: the 'print' command now uses 'Message.from_file()')
+    print_cmd_parser.add_argument('-S', '--source-code-only', help='Print only source code',
+                                  action='store_true')
+=======
+    print_cmd_parser.add_argument('-S', '--source-code-only', help='Print source code only (from the answer, if available)',
+                                  action='store_true')
+>>>>>>> bf1cbff (cmm: the 'print' command now uses 'Message.from_file()')
 
     argcomplete.autocomplete(parser)
     return parser