Improve handling of printing source code only output.

chatmastermind/main.py

@@ -6,15 +6,26 @@ import yaml
 import sys
 import argcomplete
 import argparse
-from .utils import terminal_width, pp, process_tags, display_chat
-from .storage import save_answers, create_chat, get_tags
+import pathlib
+from .utils import terminal_width, process_tags, display_chat, display_source_code
+from .storage import save_answers, create_chat, get_tags, read_file, dump_data
 from .api_client import ai, openai_api_key
 
 
 def run_print_command(args: argparse.Namespace, config: dict) -> None:
-    with open(args.print, 'r') as f:
-        data = yaml.load(f, Loader=yaml.FullLoader)
-    pp(data)
+    fname = pathlib.Path(args.print)
+    if fname.suffix == '.yaml':
+        with open(args.print, '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.print}")
+        sys.exit(1)
+    if args.only_source_code:
+        display_source_code(data['answer'])
+    else:
+        print(dump_data(data).strip())
 
 
 def process_and_display_chat(args: argparse.Namespace,
@@ -74,7 +85,7 @@ def create_parser() -> argparse.ArgumentParser:
     parser = argparse.ArgumentParser(
         description="ChatMastermind is a Python application that automates conversation with AI")
     group = parser.add_mutually_exclusive_group(required=True)
-    group.add_argument('-p', '--print', help='YAML file to print')
+    group.add_argument('-p', '--print', help='File to print')
     group.add_argument('-q', '--question', nargs='*', help='Question to ask')
     group.add_argument('-D', '--chat-dump', help="Print chat as Python structure", action='store_true')
     group.add_argument('-d', '--chat', help="Print chat as readable text", action='store_true')