change input/output prompts and color scheme

petercorke · petercorke · commit ec48bc1c4323 · 2021-01-24T09:55:21.000+10:00
diff --git a/examples/rtb.py b/examples/rtb.py
@@ -31,13 +31,23 @@
     help='specify script to run')
 parser.add_argument('--backend', '-b', default=None,
     help='specify graphics frontend')
+parser.add_argument('--color', '-c', default='neutral',
+    help='specify terminal color scheme (neutral, lightbg, nocolor, linux), linux is for dark mode')
+parser.add_argument('--confirmexit', '-x', default=False,
+    help='confirm exit')
+parser.add_argument('--prompt', '-p', default='>>> ',
+    help='input prompt')
+parser.add_argument('--resultprefix', '-r', default=None,
+    help='execution result prefix, include {} for execution count number')
+parser.add_argument('--showassign', '-a', default=False,
+    help='display the result of assignments')
 args = parser.parse_args()
 
 if args.backend is not None:
     print(f"Using matplotlub backend {args.backend}")
     plt.use(args.backend)
 
-# load some models
+# load some robot models
 puma = models.DH.Puma560()
 panda = models.DH.Panda()
 
@@ -74,15 +84,26 @@
 
 class MyPrompt(Prompts):
     def in_prompt_tokens(self, cli=None):
-        return [(Token.Prompt, '>>> ')]
+        return [(Token.Prompt, args.prompt)]
+    def out_prompt_tokens(self, cli=None):
+        if args.resultprefix is None:
+            # traditional behaviour
+            return [
+                (Token.OutPrompt, 'Out['),
+                (Token.OutPromptNum, str(self.shell.execution_count)),
+                (Token.OutPrompt, ']: '),
+        ]
+        else:
+            return [(Token.Prompt, args.resultprefix.format(self.shell.execution_count))]
 
 # set configuration options, there are lots, see
 # https://ipython.readthedocs.io/en/stable/config/options/terminal.html
 c = Config()
-c.InteractiveShellEmbed.colors = "Linux"
-c.InteractiveShell.colors = 'Neutral'
-c.InteractiveShell.confirm_exit = False
+c.InteractiveShellEmbed.colors = args.color
+c.InteractiveShell.confirm_exit = args.confirmexit
 # c.InteractiveShell.prompts_class = ClassicPrompts
 c.InteractiveShell.prompts_class = MyPrompt
+if showassign:
+    c.InteractiveShell.ast_node_interactivity = 'last_expr_or_assign'
 
 IPython.embed(config=c)
