py/modsys: Add optional mutable attributes sys.ps1/ps2 and use them.

This allows customising the REPL prompt strings.

Signed-off-by: Damien George <[email protected]>

docs/library/sys.rst

@@ -132,6 +132,12 @@ Constants
    If you need to check whether your program runs on MicroPython (vs other
    Python implementation), use `sys.implementation` instead.
 
+.. data:: ps1
+          ps2
+
+   Mutable attributes holding strings, which are used for the REPL prompt.  The defaults
+   give the standard Python prompt of ``>>>`` and ``...``.
+
 .. data:: stderr
 
    Standard error `stream`.

ports/unix/main.c

@@ -193,7 +193,7 @@ STATIC int do_repl(void) {
 
     input_restart:
         vstr_reset(&line);
-        int ret = readline(&line, ">>> ");
+        int ret = readline(&line, mp_repl_get_ps1());
         mp_parse_input_kind_t parse_input_kind = MP_PARSE_SINGLE_INPUT;
 
         if (ret == CHAR_CTRL_C) {
@@ -240,7 +240,7 @@ STATIC int do_repl(void) {
             // got a line with non-zero length, see if it needs continuing
             while (mp_repl_continue_with_input(vstr_null_terminated_str(&line))) {
                 vstr_add_byte(&line, '\n');
-                ret = readline(&line, "... ");
+                ret = readline(&line, mp_repl_get_ps2());
                 if (ret == CHAR_CTRL_C) {
                     // cancel everything
                     printf("\n");
@@ -265,13 +265,13 @@ STATIC int do_repl(void) {
     // use simple readline
 
     for (;;) {
-        char *line = prompt(">>> ");
+        char *line = prompt((char *)mp_repl_get_ps1());
         if (line == NULL) {
             // EOF
             return 0;
         }
         while (mp_repl_continue_with_input(line)) {
-            char *line2 = prompt("... ");
+            char *line2 = prompt((char *)mp_repl_get_ps2());
             if (line2 == NULL) {
                 break;
             }

ports/unix/variants/coverage/mpconfigvariant.h

@@ -45,6 +45,7 @@
 #define MICROPY_PY_BUILTINS_HELP       (1)
 #define MICROPY_PY_BUILTINS_HELP_MODULES (1)
 #define MICROPY_PY_SYS_GETSIZEOF       (1)
+#define MICROPY_PY_SYS_PS1_PS2         (1)
 #define MICROPY_PY_SYS_TRACEBACKLIMIT  (1)
 #define MICROPY_PY_MATH_CONSTANTS      (1)
 #define MICROPY_PY_MATH_FACTORIAL      (1)

py/modsys.c

@@ -185,6 +185,10 @@ MP_DEFINE_CONST_FUN_OBJ_1(mp_sys_settrace_obj, mp_sys_settrace);
 
 #if MICROPY_PY_SYS_ATTR_DELEGATION
 STATIC const uint16_t sys_mutable_keys[] = {
+    #if MICROPY_PY_SYS_PS1_PS2
+    MP_QSTR_ps1,
+    MP_QSTR_ps2,
+    #endif
     #if MICROPY_PY_SYS_TRACEBACKLIMIT
     MP_QSTR_tracebacklimit,
     #endif

py/mpconfig.h

@@ -1356,6 +1356,11 @@ typedef double mp_float_t;
 #define MICROPY_PY_SYS_ATEXIT (0)
 #endif
 
+// Whether to provide sys.{ps1,ps2} mutable attributes, to control REPL prompts
+#ifndef MICROPY_PY_SYS_PS1_PS2
+#define MICROPY_PY_SYS_PS1_PS2 (MICROPY_CONFIG_ROM_LEVEL_AT_LEAST_EXTRA_FEATURES)
+#endif
+
 // Whether to provide "sys.settrace" function
 #ifndef MICROPY_PY_SYS_SETTRACE
 #define MICROPY_PY_SYS_SETTRACE (0)
@@ -1385,7 +1390,7 @@ typedef double mp_float_t;
 // Whether the sys module supports attribute delegation
 // This is enabled automatically when needed by other features
 #ifndef MICROPY_PY_SYS_ATTR_DELEGATION
-#define MICROPY_PY_SYS_ATTR_DELEGATION (MICROPY_PY_SYS_TRACEBACKLIMIT)
+#define MICROPY_PY_SYS_ATTR_DELEGATION (MICROPY_PY_SYS_PS1_PS2 || MICROPY_PY_SYS_TRACEBACKLIMIT)
 #endif
 
 // Whether to provide "uerrno" module

py/mpstate.h

@@ -41,6 +41,10 @@
 // variable, but in the future it is hoped that the state can become local.
 
 enum {
+    #if MICROPY_PY_SYS_PS1_PS2
+    MP_SYS_MUTABLE_PS1,
+    MP_SYS_MUTABLE_PS2,
+    #endif
     #if MICROPY_PY_SYS_TRACEBACKLIMIT
     MP_SYS_MUTABLE_TRACEBACKLIMIT,
     #endif

py/qstrdefs.h

@@ -39,6 +39,10 @@ Q()
 Q(*)
 Q(_)
 Q(/)
+#if MICROPY_PY_SYS_PS1_PS2
+Q(>>> )
+Q(... )
+#endif
 #if MICROPY_PY_BUILTINS_STR_OP_MODULO
 Q(%#o)
 Q(%#x)

py/repl.c

@@ -33,6 +33,16 @@
 
 #if MICROPY_HELPER_REPL
 
+#if MICROPY_PY_SYS_PS1_PS2
+const char *mp_repl_get_psx(unsigned int entry) {
+    if (mp_obj_is_str(MP_STATE_VM(sys_mutable)[entry])) {
+        return mp_obj_str_get_str(MP_STATE_VM(sys_mutable)[entry]);
+    } else {
+        return "";
+    }
+}
+#endif
+
 STATIC bool str_startswith_word(const char *str, const char *head) {
     size_t i;
     for (i = 0; str[i] && head[i]; i++) {

py/repl.h

@@ -31,8 +31,34 @@
 #include "py/mpprint.h"
 
 #if MICROPY_HELPER_REPL
+
+#if MICROPY_PY_SYS_PS1_PS2
+
+const char *mp_repl_get_psx(unsigned int entry);
+
+static inline const char *mp_repl_get_ps1(void) {
+    return mp_repl_get_psx(MP_SYS_MUTABLE_PS1);
+}
+
+static inline const char *mp_repl_get_ps2(void) {
+    return mp_repl_get_psx(MP_SYS_MUTABLE_PS2);
+}
+
+#else
+
+static inline const char *mp_repl_get_ps1(void) {
+    return ">>> ";
+}
+
+static inline const char *mp_repl_get_ps2(void) {
+    return "... ";
+}
+
+#endif
+
 bool mp_repl_continue_with_input(const char *input);
 size_t mp_repl_autocomplete(const char *str, size_t len, const mp_print_t *print, const char **compl_str);
+
 #endif
 
 #endif // MICROPY_INCLUDED_PY_REPL_H

py/runtime.c

@@ -135,6 +135,11 @@ void mp_init(void) {
     MP_STATE_VM(sys_exitfunc) = mp_const_none;
     #endif
 
+    #if MICROPY_PY_SYS_PS1_PS2
+    MP_STATE_VM(sys_mutable[MP_SYS_MUTABLE_PS1]) = MP_OBJ_NEW_QSTR(MP_QSTR__gt__gt__gt__space_);
+    MP_STATE_VM(sys_mutable[MP_SYS_MUTABLE_PS2]) = MP_OBJ_NEW_QSTR(MP_QSTR__dot__dot__dot__space_);
+    #endif
+
     #if MICROPY_PY_SYS_SETTRACE
     MP_STATE_THREAD(prof_trace_callback) = MP_OBJ_NULL;
     MP_STATE_THREAD(prof_callback_is_executing) = false;

shared/runtime/pyexec.c

@@ -433,7 +433,7 @@ STATIC int pyexec_friendly_repl_process_char(int c) {
 
         vstr_add_byte(MP_STATE_VM(repl_line), '\n');
         repl.cont_line = true;
-        readline_note_newline("... ");
+        readline_note_newline(mp_repl_get_ps2());
         return 0;
 
     } else {
@@ -454,7 +454,7 @@ STATIC int pyexec_friendly_repl_process_char(int c) {
 
         if (mp_repl_continue_with_input(vstr_null_terminated_str(MP_STATE_VM(repl_line)))) {
             vstr_add_byte(MP_STATE_VM(repl_line), '\n');
-            readline_note_newline("... ");
+            readline_note_newline(mp_repl_get_ps2());
             return 0;
         }
 
@@ -468,7 +468,7 @@ STATIC int pyexec_friendly_repl_process_char(int c) {
         vstr_reset(MP_STATE_VM(repl_line));
         repl.cont_line = false;
         repl.paste_mode = false;
-        readline_init(MP_STATE_VM(repl_line), ">>> ");
+        readline_init(MP_STATE_VM(repl_line), mp_repl_get_ps1());
         return 0;
     }
 }
@@ -598,7 +598,7 @@ int pyexec_friendly_repl(void) {
         }
 
         vstr_reset(&line);
-        int ret = readline(&line, ">>> ");
+        int ret = readline(&line, mp_repl_get_ps1());
         mp_parse_input_kind_t parse_input_kind = MP_PARSE_SINGLE_INPUT;
 
         if (ret == CHAR_CTRL_A) {
@@ -651,7 +651,7 @@ int pyexec_friendly_repl(void) {
             // got a line with non-zero length, see if it needs continuing
             while (mp_repl_continue_with_input(vstr_null_terminated_str(&line))) {
                 vstr_add_byte(&line, '\n');
-                ret = readline(&line, "... ");
+                ret = readline(&line, mp_repl_get_ps2());
                 if (ret == CHAR_CTRL_C) {
                     // cancel everything
                     mp_hal_stdout_tx_str("\r\n");

tests/cmdline/repl_sys_ps1_ps2.py

@@ -0,0 +1,6 @@
+# test changing ps1/ps2
+import usys
+usys.ps1 = "PS1"
+usys.ps2 = "PS2"
+(1 +
+2)

tests/cmdline/repl_sys_ps1_ps2.py.exp

@@ -0,0 +1,10 @@
+MicroPython \.\+ version
+Use \.\+
+>>> # test changing ps1/ps2
+>>> import usys
+>>> usys.ps1 = "PS1"
+PS1usys.ps2 = "PS2"
+PS1(1 +
+PS22)
+3
+PS1

tests/run-tests.py

@@ -433,6 +433,7 @@ def run_tests(pyb, tests, args, result_dir, num_threads=1):
 
     if not has_coverage:
         skip_tests.add("cmdline/cmd_parsetree.py")
+        skip_tests.add("cmdline/repl_sys_ps1_ps2.py")
 
     # Some tests shouldn't be run on a PC
     if args.target == "unix":

tests/unix/extra_coverage.py.exp

@@ -45,8 +45,8 @@ utime           utimeq
 argv            atexit          byteorder       exc_info
 exit            getsizeof       implementation  maxsize
 modules         path            platform        print_exception
-stderr          stdin           stdout          tracebacklimit
-version         version_info
+ps1             ps2             stderr          stdin
+stdout          tracebacklimit  version         version_info
 ementation
 # attrtuple
 (start=1, stop=2, step=3)