Capture evaluator standard output and send to the caller

lib/live_book/evaluator.ex

@@ -13,6 +13,10 @@ defmodule LiveBook.Evaluator do
 
   use GenServer
 
+  import LiveBook.Utils
+
+  alias LiveBook.Evaluator
+
   @type t :: GenServer.server()
 
   @type state :: %{
@@ -29,7 +33,7 @@ defmodule LiveBook.Evaluator do
   """
   @type ref :: term()
 
-  # API
+  ## API
 
   def start_link(opts \\ []) do
     GenServer.start_link(__MODULE__, opts)
@@ -49,7 +53,7 @@ defmodule LiveBook.Evaluator do
     GenServer.call(evaluator, {:evaluate_code, code, ref, prev_ref}, :infinity)
   end
 
-  # Callbacks
+  ## Callbacks
 
   @impl true
   def init(_opts) do
@@ -66,20 +70,25 @@ defmodule LiveBook.Evaluator do
   end
 
   @impl true
-  def handle_call({:evaluate_code, code, ref, prev_ref}, _from, state) do
+  def handle_call({:evaluate_code, code, ref, prev_ref}, {from, _}, state) do
     context = state.contexts[prev_ref]
 
-    case eval(code, context.binding, context.env) do
-      {:ok, result, binding, env} ->
-        result_context = %{binding: binding, env: env}
+    {:ok, io} = Evaluator.IOProxy.start_link(from, ref)
 
-        new_contexts = Map.put(state.contexts, ref, result_context)
-        new_state = %{state | contexts: new_contexts}
+    # Use the dedicated IO device as the group leader,
+    # so that it handles all :stdio operations.
+    with_group_leader io do
+      case eval(code, context.binding, context.env) do
+        {:ok, result, binding, env} ->
+          result_context = %{binding: binding, env: env}
+          new_contexts = Map.put(state.contexts, ref, result_context)
+          new_state = %{state | contexts: new_contexts}
 
-        {:reply, {:ok, result}, new_state}
+          {:reply, {:ok, result}, new_state}
 
-      {:error, exception} ->
-        {:reply, {:error, exception}, state}
+        {:error, exception} ->
+          {:reply, {:error, exception}, state}
+      end
     end
   end
 

lib/live_book/evaluator/io_proxy.ex

@@ -0,0 +1,149 @@
+defmodule LiveBook.Evaluator.IOProxy do
+  @moduledoc false
+
+  # An IO device process used by `Evaluator` as its `:stdio`.
+  #
+  # The process implements [The Erlang I/O Protocol](https://erlang.org/doc/apps/stdlib/io_protocol.html)
+  # and can be thought of as a *virtual* IO device.
+  #
+  # Upon receiving an IO requests, the process sends a message
+  # the `target` process specified during initialization.
+  # Currently only output requests are supported.
+  #
+  # The implementation is based the build-in `StringIO`,
+  # so check it out for more reference.
+
+  use GenServer
+
+  alias LiveBook.Evaluator
+
+  ## API
+
+  @doc """
+  Starts the IO device process.
+
+  For all supported requests a message is sent to `target`,
+  so this device serves as a proxy. The given evaluation
+  reference (`ref`) is also sent in all messages.
+
+  The possible messages are:
+
+    * `{:evaluator_stdout, ref, string}` - for output requests,
+      where `ref` is the given evaluation reference and `string` is the output.
+  """
+  @spec start_link(pid(), Evaluator.ref()) :: GenServer.on_start()
+  def start_link(target, ref) do
+    GenServer.start_link(__MODULE__, {target, ref})
+  end
+
+  ## Callbacks
+
+  @impl true
+  def init({target, ref}) do
+    {:ok, %{encoding: :unicode, target: target, ref: ref}}
+  end
+
+  @impl true
+  def handle_info({:io_request, from, reply_as, req}, state) do
+    {reply, state} = io_request(req, state)
+    io_reply(from, reply_as, reply)
+    {:noreply, state}
+  end
+
+  defp io_request({:put_chars, chars} = req, state) do
+    put_chars(:latin1, chars, req, state)
+  end
+
+  defp io_request({:put_chars, mod, fun, args} = req, state) do
+    put_chars(:latin1, apply(mod, fun, args), req, state)
+  end
+
+  defp io_request({:put_chars, encoding, chars} = req, state) do
+    put_chars(encoding, chars, req, state)
+  end
+
+  defp io_request({:put_chars, encoding, mod, fun, args} = req, state) do
+    put_chars(encoding, apply(mod, fun, args), req, state)
+  end
+
+  defp io_request({:get_chars, _prompt, count}, state) when count >= 0 do
+    {{:error, :enotsup}, state}
+  end
+
+  defp io_request({:get_chars, _encoding, _prompt, count}, state) when count >= 0 do
+    {{:error, :enotsup}, state}
+  end
+
+  defp io_request({:get_line, _prompt}, state) do
+    {{:error, :enotsup}, state}
+  end
+
+  defp io_request({:get_line, _encoding, _prompt}, state) do
+    {{:error, :enotsup}, state}
+  end
+
+  defp io_request({:get_until, _prompt, _mod, _fun, _args}, state) do
+    {{:error, :enotsup}, state}
+  end
+
+  defp io_request({:get_until, _encoding, _prompt, _mod, _fun, _args}, state) do
+    {{:error, :enotsup}, state}
+  end
+
+  defp io_request({:get_password, _encoding}, state) do
+    {{:error, :enotsup}, state}
+  end
+
+  defp io_request({:setopts, [encoding: encoding]}, state) when encoding in [:latin1, :unicode] do
+    {:ok, %{state | encoding: encoding}}
+  end
+
+  defp io_request({:setopts, _opts}, state) do
+    {{:error, :enotsup}, state}
+  end
+
+  defp io_request(:getopts, state) do
+    {[binary: true, encoding: state.encoding], state}
+  end
+
+  defp io_request({:get_geometry, :columns}, state) do
+    {{:error, :enotsup}, state}
+  end
+
+  defp io_request({:get_geometry, :rows}, state) do
+    {{:error, :enotsup}, state}
+  end
+
+  defp io_request({:requests, reqs}, state) do
+    io_requests(reqs, {:ok, state})
+  end
+
+  defp io_request(_, state) do
+    {{:error, :request}, state}
+  end
+
+  defp io_requests([req | rest], {:ok, state}) do
+    io_requests(rest, io_request(req, state))
+  end
+
+  defp io_requests(_, result) do
+    result
+  end
+
+  defp put_chars(encoding, chars, req, state) do
+    case :unicode.characters_to_binary(chars, encoding, state.encoding) do
+      string when is_binary(string) ->
+        send(state.target, {:evaluator_stdout, state.ref, string})
+        {:ok, state}
+
+      {_, _, _} ->
+        {{:error, req}, state}
+    end
+  rescue
+    ArgumentError -> {{:error, req}, state}
+  end
+
+  defp io_reply(from, reply_as, reply) do
+    send(from, {:io_reply, reply_as, reply})
+  end
+end

lib/live_book/utils.ex

@@ -10,4 +10,23 @@ defmodule LiveBook.Utils do
   def random_id() do
     :crypto.strong_rand_bytes(20) |> Base.encode32(case: :lower)
   end
+
+  @doc """
+  Wraps the given expression so that it's executed
+  with the given process as the group leader.
+
+  Always restores the original group leader afterwards.
+  """
+  defmacro with_group_leader(group_leader, do: expression) do
+    quote do
+      original_gl = Process.group_leader()
+
+      try do
+        Process.group_leader(self(), unquote(group_leader))
+        unquote(expression)
+      after
+        Process.group_leader(self(), original_gl)
+      end
+    end
+  end
 end

test/live_book/evaluator/io_proxy_test.exs

@@ -0,0 +1,35 @@
+defmodule LiveBook.Evaluator.IOProxyTest do
+  use ExUnit.Case, async: true
+
+  alias LiveBook.Evaluator.IOProxy
+
+  setup do
+    {:ok, io} = IOProxy.start_link(self(), :ref)
+    %{io: io}
+  end
+
+  # Test the basic ways users interact with :stdio
+
+  test "IO.puts", %{io: io} do
+    IO.puts(io, "hey")
+    assert_received {:evaluator_stdout, :ref, "hey\n"}
+  end
+
+  test "IO.write", %{io: io} do
+    IO.write(io, "hey")
+    assert_received {:evaluator_stdout, :ref, "hey"}
+  end
+
+  test "IO.inspect", %{io: io} do
+    IO.inspect(io, %{}, [])
+    assert_received {:evaluator_stdout, :ref, "%{}\n"}
+  end
+
+  test "IO.read", %{io: io} do
+    assert IO.read(io, :all) == {:error, :enotsup}
+  end
+
+  test "IO.gets", %{io: io} do
+    assert IO.gets(io, "> ") == {:error, :enotsup}
+  end
+end

test/live_book/evaluator_test.exs

@@ -4,7 +4,7 @@ defmodule LiveBook.EvaluatorTest do
   alias LiveBook.Evaluator
 
   setup do
-    evaluator = start_supervised!(Evaluator)
+    {:ok, evaluator} = Evaluator.start_link()
     %{evaluator: evaluator}
   end
 
@@ -29,12 +29,24 @@ defmodule LiveBook.EvaluatorTest do
       assert {:error, %CompileError{description: "undefined function x/0"}} = result
     end
 
-    test "given prev_ref does sees previous evaluation context", %{evaluator: evaluator} do
+    test "given prev_ref sees previous evaluation context", %{evaluator: evaluator} do
       Evaluator.evaluate_code(evaluator, "x = 1", :code_1)
 
       result = Evaluator.evaluate_code(evaluator, "x", :code_2, :code_1)
 
       assert result == {:ok, 1}
     end
+
+    test "captures standard output and sends it to the caller", %{evaluator: evaluator} do
+      Evaluator.evaluate_code(evaluator, ~s{IO.puts("hey")}, :code_1)
+
+      assert_received {:evaluator_stdout, :code_1, "hey\n"}
+    end
+
+    test "using standard input results in an immediate error", %{evaluator: evaluator} do
+      result = Evaluator.evaluate_code(evaluator, ~s{IO.gets("> ")}, :code_1)
+
+      assert result == {:ok, {:error, :enotsup}}
+    end
   end
 end