Add ability to add output prefixes to the REPL output and use that to…

… implement an IPython mode (JuliaLang#46474)
green-nsk · Sep 15, 2022 · e118845 · e118845
1 parent 6557542
commit e118845
Show file tree

Hide file tree

Showing 5 changed files with 140 additions and 6 deletions.
diff --git a/NEWS.md b/NEWS.md
@@ -129,6 +129,9 @@ Standard library changes
   via the `REPL.activate(::Module)` function or via typing the module in the REPL and pressing
   the keybinding Alt-m ([#33872]).
 
+* An "IPython mode" which mimics the behaviour of the prompts and storing the evaluated result in `Out` can be
+  activated with `REPL.ipython_mode()`. See the manual for how to enable this at startup.
+
 #### SparseArrays
 
 #### Test

diff --git a/stdlib/REPL/docs/src/index.md b/stdlib/REPL/docs/src/index.md
@@ -616,6 +616,20 @@ julia> REPL.activate(CustomMod)
   var       8 bytes Int64
 ```
 
+## IPython mode
+
+It is possible to get an interface which is similar to the IPython REPL with numbered input prompts and output prefixes. This is done by calling `REPL.ipython_mode()`. If you want to have this enabled on startup, add
+```julia
+atreplinit() do repl
+    if !isdefined(repl, :interface)
+        repl.interface = REPL.setup_interface(repl)
+    end
+    REPL.ipython_mode(repl)
+end
+```
+
+to your `startup.jl` file.
+
 ## TerminalMenus
 
 TerminalMenus is a submodule of the Julia REPL and enables small, low-profile interactive menus in the terminal.

diff --git a/stdlib/REPL/src/LineEdit.jl b/stdlib/REPL/src/LineEdit.jl
@@ -49,6 +49,9 @@ mutable struct Prompt <: TextInterface
     prompt_prefix::Union{String,Function}
     # Same as prefix except after the prompt
     prompt_suffix::Union{String,Function}
+    output_prefix::Union{String,Function}
+    output_prefix_prefix::Union{String,Function}
+    output_prefix_suffix::Union{String,Function}
     keymap_dict::Dict{Char,Any}
     repl::Union{AbstractREPL,Nothing}
     complete::CompletionProvider
@@ -1452,7 +1455,6 @@ default_completion_cb(::IOBuffer) = []
 default_enter_cb(_) = true
 
 write_prompt(terminal::AbstractTerminal, s::PromptState, color::Bool) = write_prompt(terminal, s.p, color)
-
 function write_prompt(terminal::AbstractTerminal, p::Prompt, color::Bool)
     prefix = prompt_string(p.prompt_prefix)
     suffix = prompt_string(p.prompt_suffix)
@@ -1464,6 +1466,17 @@ function write_prompt(terminal::AbstractTerminal, p::Prompt, color::Bool)
     return width
 end
 
+function write_output_prefix(io::IO, p::Prompt, color::Bool)
+    prefix = prompt_string(p.output_prefix_prefix)
+    suffix = prompt_string(p.output_prefix_suffix)
+    print(io, prefix)
+    color && write(io, Base.text_colors[:bold])
+    width = write_prompt(io, p.output_prefix, color)
+    color && write(io, Base.text_colors[:normal])
+    print(io, suffix)
+    return width
+end
+
 # On Windows, when launching external processes, we cannot control what assumption they make on the
 # console mode. We thus forcibly reset the console mode at the start of the prompt to ensure they do
 # not leave the console mode in a corrupt state.
@@ -2591,6 +2604,9 @@ function Prompt(prompt
     ;
     prompt_prefix = "",
     prompt_suffix = "",
+    output_prefix = "",
+    output_prefix_prefix = "",
+    output_prefix_suffix = "",
     keymap_dict = default_keymap_dict,
     repl = nothing,
     complete = EmptyCompletionProvider(),
@@ -2599,8 +2615,8 @@ function Prompt(prompt
     hist = EmptyHistoryProvider(),
     sticky = false)
 
-    return Prompt(prompt, prompt_prefix, prompt_suffix, keymap_dict, repl,
-        complete, on_enter, on_done, hist, sticky)
+    return Prompt(prompt, prompt_prefix, prompt_suffix, output_prefix, output_prefix_prefix, output_prefix_suffix,
+                   keymap_dict, repl, complete, on_enter, on_done, hist, sticky)
 end
 
 run_interface(::Prompt) = nothing

diff --git a/stdlib/REPL/src/REPL.jl b/stdlib/REPL/src/REPL.jl
@@ -260,6 +260,11 @@ function display(d::REPLDisplay, mime::MIME"text/plain", x)
     x = Ref{Any}(x)
     with_repl_linfo(d.repl) do io
         io = IOContext(io, :limit => true, :module => active_module(d)::Module)
+        if d.repl isa LineEditREPL
+            mistate = d.repl.mistate
+            mode = LineEdit.mode(mistate)
+            LineEdit.write_output_prefix(io, mode, get(io, :color, false))
+        end
         get(io, :color, false) && write(io, answer_color(d.repl))
         if isdefined(d.repl, :options) && isdefined(d.repl.options, :iocontext)
             # this can override the :limit property set initially
@@ -354,8 +359,7 @@ end
 
     consumer is an optional function that takes a REPLBackend as an argument
 """
-function run_repl(repl::AbstractREPL, @nospecialize(consumer = x -> nothing); backend_on_current_task::Bool = true)
-    backend = REPLBackend()
+function run_repl(repl::AbstractREPL, @nospecialize(consumer = x -> nothing); backend_on_current_task::Bool = true, backend = REPLBackend())
     backend_ref = REPLBackendRef(backend)
     cleanup = @task try
             destroy(backend_ref, t)
@@ -1062,7 +1066,7 @@ function setup_interface(
 
     shell_prompt_len = length(SHELL_PROMPT)
     help_prompt_len = length(HELP_PROMPT)
-    jl_prompt_regex = r"^(?:\(.+\) )?julia> "
+    jl_prompt_regex = r"^In \[[0-9]+\]: |^(?:\(.+\) )?julia> "
     pkg_prompt_regex = r"^(?:\(.+\) )?pkg> "
 
     # Canonicalize user keymap input
@@ -1388,4 +1392,62 @@ function run_frontend(repl::StreamREPL, backend::REPLBackendRef)
     nothing
 end
 
+module IPython
+
+using ..REPL
+
+__current_ast_transforms() = isdefined(Base, :active_repl_backend) ? Base.active_repl_backend.ast_transforms : REPL.repl_ast_transforms
+
+function repl_eval_counter(hp)
+    length(hp.history)-hp.start_idx
+end
+
+function out_transform(x, repl=Base.active_repl)
+    return quote
+        julia_prompt = $repl.interface.modes[1]
+        n = $repl_eval_counter(julia_prompt.hist)
+        mod = $REPL.active_module()
+        if !isdefined(mod, :Out)
+            setglobal!(mod, :Out, Dict{Int, Any}())
+        end
+        local __temp_val = $x # workaround https://github.com/JuliaLang/julia/issues/464511
+        if __temp_val !== getglobal(mod, :Out) && __temp_val !== nothing # remove this?
+            getglobal(mod, :Out)[n] = __temp_val
+        end
+        __temp_val
+    end
+end
+
+function set_prompt(repl=Base.active_repl)
+    julia_prompt = repl.interface.modes[1]
+    julia_prompt.prompt = () -> string("In [", repl_eval_counter(julia_prompt.hist)+1, "]: ")
+end
+
+function set_output_prefix(repl=Base.active_repl)
+    julia_prompt = repl.interface.modes[1]
+    if REPL.hascolor(repl)
+        julia_prompt.output_prefix_prefix = Base.text_colors[:red]
+    end
+    julia_prompt.output_prefix = () -> string("Out[", repl_eval_counter(julia_prompt.hist), "]: ")
+end
+
+function __current_ast_transforms(repltask)
+    if repltask === nothing
+        isdefined(Base, :active_repl_backend) ? Base.active_repl_backend.ast_transforms : REPL.repl_ast_transforms
+    else
+        repltask.ast_transforms
+    end
+end
+
+
+function ipython_mode(repl=Base.active_repl, repltask=nothing)
+    set_prompt(repl)
+    set_output_prefix(repl)
+    push!(__current_ast_transforms(repltask), ast -> out_transform(ast, repl))
+    return
+end
+end
+
+import .IPython.ipython_mode
+
 end # module
diff --git a/stdlib/REPL/test/repl.jl b/stdlib/REPL/test/repl.jl
@@ -707,6 +707,11 @@ fake_repl() do stdin_write, stdout_read, repl
     wait(c)
     @test Main.A == 2
 
+    # Test removal of prefix in single statement paste
+    sendrepl2("\e[200~In [12]: A = 2.2\e[201~\n")
+    wait(c)
+    @test Main.A == 2.2
+
     # Test removal of prefix in multiple statement paste
     sendrepl2("""\e[200~
             julia> mutable struct T17599; a::Int; end
@@ -1548,3 +1553,37 @@ fake_repl() do stdin_write, stdout_read, repl
     LineEdit.edit_input(s, input_f)
     @test buffercontents(LineEdit.buffer(s)) == "1234αβ56γ"
 end
+
+# Non standard output_prefix, tested via `ipython_mode`
+fake_repl() do stdin_write, stdout_read, repl
+    repl.interface = REPL.setup_interface(repl)
+
+    backend = REPL.REPLBackend()
+    repltask = @async begin
+        REPL.run_repl(repl; backend)
+    end
+
+    REPL.ipython_mode(repl, backend)
+
+    global c = Condition()
+    sendrepl2(cmd) = write(stdin_write, "$cmd\n notify($(curmod_prefix)c)\n")
+
+    sendrepl2("\"z\" * \"z\"\n")
+    wait(c)
+    s = String(readuntil(stdout_read, "\"zz\""; keep=true))
+    @test contains(s, "In [1]")
+    @test contains(s, "Out[1]: \"zz\"")
+
+    sendrepl2("\"y\" * \"y\"\n")
+    wait(c)
+    s = String(readuntil(stdout_read, "\"yy\""; keep=true))
+    @test contains(s, "Out[3]: \"yy\"")
+
+    sendrepl2("Out[1] * Out[3]\n")
+    wait(c)
+    s = String(readuntil(stdout_read, "\"zzyy\""; keep=true))
+    @test contains(s, "Out[5]: \"zzyy\"")
+
+    write(stdin_write, '\x04')
+    Base.wait(repltask)
+end