From 9ae2f02124948b5a0be4b89bd9353a6b2644d5a8 Mon Sep 17 00:00:00 2001
From: Ian Butterworth <i.r.butterworth@gmail.com>
Date: Tue, 27 Apr 2021 18:50:21 -0400
Subject: [PATCH] Enable REPL to offer to install missing packages if install
 hooks are provided (#39026)

---
 stdlib/REPL/src/REPL.jl  | 37 +++++++++++++++++++++++++++++++++++++
 stdlib/REPL/test/repl.jl | 27 +++++++++++++++++++++++++++
 2 files changed, 64 insertions(+)

diff --git a/stdlib/REPL/src/REPL.jl b/stdlib/REPL/src/REPL.jl
index d17482ec453cba..0c5a6c6267cffb 100644
--- a/stdlib/REPL/src/REPL.jl
+++ b/stdlib/REPL/src/REPL.jl
@@ -123,6 +123,12 @@ const softscope! = softscope
 
 const repl_ast_transforms = Any[softscope] # defaults for new REPL backends
 
+# Allows an external package to add hooks into the code loading.
+# The hook should take a Vector{Symbol} of package names and
+# return true if all packages could be installed, false if not
+# to e.g. install packages on demand
+const install_packages_hooks = Any[]
+
 function eval_user_input(@nospecialize(ast), backend::REPLBackend)
     lasterr = nothing
     Base.sigatomic_begin()
@@ -133,6 +139,9 @@ function eval_user_input(@nospecialize(ast), backend::REPLBackend)
                 put!(backend.response_channel, Pair{Any, Bool}(lasterr, true))
             else
                 backend.in_eval = true
+                if !isempty(install_packages_hooks)
+                    check_for_missing_packages_and_run_hooks(ast)
+                end
                 for xf in backend.ast_transforms
                     ast = Base.invokelatest(xf, ast)
                 end
@@ -155,6 +164,34 @@ function eval_user_input(@nospecialize(ast), backend::REPLBackend)
     nothing
 end
 
+function check_for_missing_packages_and_run_hooks(ast)
+    mods = modules_to_be_loaded(ast)
+    filter!(mod -> isnothing(Base.identify_package(String(mod))), mods) # keep missing modules
+    if !isempty(mods)
+        for f in install_packages_hooks
+            Base.invokelatest(f, mods) && return
+        end
+    end
+end
+
+function modules_to_be_loaded(ast, mods = Symbol[])
+    if ast.head in [:using, :import]
+        for arg in ast.args
+            if first(arg.args) isa Symbol # i.e. `Foo`
+                if first(arg.args) != :. # don't include local imports
+                    push!(mods, first(arg.args))
+                end
+            else # i.e. `Foo: bar`
+                push!(mods, first(first(arg.args).args))
+            end
+        end
+    end
+    for arg in ast.args
+        arg isa Expr && modules_to_be_loaded(arg, mods)
+    end
+    return mods
+end
+
 """
     start_repl_backend(repl_channel::Channel, response_channel::Channel)
 
diff --git a/stdlib/REPL/test/repl.jl b/stdlib/REPL/test/repl.jl
index cb2d086a30d7f2..677d184f99eabd 100644
--- a/stdlib/REPL/test/repl.jl
+++ b/stdlib/REPL/test/repl.jl
@@ -1295,3 +1295,30 @@ Base.wait(frontend_task)
 macro throw_with_linenumbernode(err)
     Expr(:block, LineNumberNode(42, Symbol("test.jl")), :(() -> throw($err)))
 end
+
+@testset "Install missing packages via hooks" begin
+    @testset "Parse AST for packages" begin
+        mods = REPL.modules_to_be_loaded(Meta.parse("using Foo"))
+        @test mods == [:Foo]
+        mods = REPL.modules_to_be_loaded(Meta.parse("import Foo"))
+        @test mods == [:Foo]
+        mods = REPL.modules_to_be_loaded(Meta.parse("using Foo, Bar"))
+        @test mods == [:Foo, :Bar]
+        mods = REPL.modules_to_be_loaded(Meta.parse("import Foo, Bar"))
+        @test mods == [:Foo, :Bar]
+
+        mods = REPL.modules_to_be_loaded(Meta.parse("if false using Foo end"))
+        @test mods == [:Foo]
+        mods = REPL.modules_to_be_loaded(Meta.parse("if false if false using Foo end end"))
+        @test mods == [:Foo]
+        mods = REPL.modules_to_be_loaded(Meta.parse("if false using Foo, Bar end"))
+        @test mods == [:Foo, :Bar]
+        mods = REPL.modules_to_be_loaded(Meta.parse("if false using Foo: bar end"))
+        @test mods == [:Foo]
+
+        mods = REPL.modules_to_be_loaded(Meta.parse("import Foo.bar as baz"))
+        @test mods == [:Foo]
+        mods = REPL.modules_to_be_loaded(Meta.parse("using .Foo"))
+        @test mods == []
+    end
+end