Fix extracting source file with type specification.

chatmastermind/main.py

@@ -125,6 +125,10 @@ def print_cmd(args: argparse.Namespace, config: Config) -> None:
         sys.exit(1)
     if args.only_source_code:
         display_source_code(data['answer'])
+    elif args.answer:
+        print(data['answer'].strip())
+    elif args.question:
+        print(data['question'].strip())
     else:
         print(dump_data(data).strip())
 
@@ -213,8 +217,10 @@ def create_parser() -> argparse.ArgumentParser:
                                             aliases=['p'])
     print_cmd_parser.set_defaults(func=print_cmd)
     print_cmd_parser.add_argument('-f', '--file', help='File to print', required=True)
-    print_cmd_parser.add_argument('-S', '--only-source-code', help='Print only source code',
+    print_cmd_modes = print_cmd_parser.add_mutually_exclusive_group()
-                                  action='store_true')
+    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')
 
     argcomplete.autocomplete(parser)
     return parser

chatmastermind/utils.py

@@ -55,9 +55,10 @@ def message_to_chat(message: dict[str, str],
 def display_source_code(content: str) -> None:
     try:
         content_start = content.index('```')
+        content_start = content.index('\n', content_start) + 1
         content_end = content.rindex('```')
-        if content_start + 3 < content_end:
+        if content_start < content_end:
-            print(content[content_start + 3:content_end].strip())
+            print(content[content_start:content_end].strip())
     except ValueError:
         pass