From 6238a64335b56b09cab5eaff9f5aef24aca21a7a Mon Sep 17 00:00:00 2001
From: KristofferC <kristoffer.carlsson@juliacomputing.com>
Date: Wed, 24 Aug 2022 10:11:10 +0200
Subject: [PATCH] add an ipython mode on top of the output prefix functionality

---
 stdlib/REPL/docs/src/index.md | 14 ++++++++
 stdlib/REPL/src/REPL.jl       | 63 +++++++++++++++++++++++++++++++++--
 stdlib/REPL/test/repl.jl      | 39 ++++++++++++++++++++++
 3 files changed, 113 insertions(+), 3 deletions(-)

diff --git a/stdlib/REPL/docs/src/index.md b/stdlib/REPL/docs/src/index.md
index 203f377c9ba637..dbc3bd0fb36ddd 100644
--- 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/REPL.jl b/stdlib/REPL/src/REPL.jl
index f5b20b7acaa44d..3cdd166537c7f9 100644
--- a/stdlib/REPL/src/REPL.jl
+++ b/stdlib/REPL/src/REPL.jl
@@ -359,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)
@@ -1067,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
@@ -1393,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/46451
+        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
index fcc571d8a44efb..0312e59419b1b0 100644
--- 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