Dispatch Python output rendering to kino (#2954)

lib/livebook/runtime/definitions.ex

@@ -509,5 +509,9 @@ defmodule Livebook.Runtime.Definitions do
     %{dep: {:pythonx, "~> 0.4.2"}, config: []}
   end
 
+  def kino_pythonx_dependency() do
+    %{dep: {:kino_pythonx, github: "livebook-dev/kino_pythonx"}, config: []}
+  end
+
   def pythonx_requirement(), do: "~> 0.4.0"
 end

lib/livebook/runtime/evaluator.ex

@@ -498,6 +498,8 @@ defmodule Livebook.Runtime.Evaluator do
     state = put_context(state, ref, new_context)
     output = Evaluator.Formatter.format_result(language, result)
 
+    after_evaluation(language)
+
     metadata = %{
       errored: error_result?(result),
       interrupted: interrupt_result?(result),
@@ -934,6 +936,7 @@ defmodule Livebook.Runtime.Evaluator do
   end
 
   @compile {:no_warn_undefined, {Pythonx, :eval, 2}}
+  @compile {:no_warn_undefined, {Pythonx, :uv_init, 1}}
   @compile {:no_warn_undefined, {Pythonx, :decode, 1}}
 
   defp eval_python(code, binding, env) do
@@ -1027,14 +1030,30 @@ defmodule Livebook.Runtime.Evaluator do
 
   defp eval_pyproject_toml(code, binding, env) do
     with :ok <- ensure_pythonx() do
-      quoted = {{:., [], [{:__aliases__, [alias: false], [:Pythonx]}, :uv_init]}, [], [code]}
-
       {result, _diagnostics} =
         Code.with_diagnostics([log: true], fn ->
           try do
-            {value, binding, env} =
-              Code.eval_quoted_with_env(quoted, binding, env, prune_binding: true)
+            Pythonx.uv_init(code)
 
+            # The default matplotlib backend relies on OS-specific GUI
+            # and crashes when embedding Python. For this reason, we
+            # configure a non-interactive backend that only allows
+            # exporting figures as images. In general we want to avoid
+            # special casing like this, but given how common matplotlib
+            # is, it does make sense to streamline the experience.
+            # We set the backend using env var, instead of calling
+            # plt.backend(...), because importing the module for the
+            # first time is slow, so we prefer to avoid that as part
+            # of setup.
+            Pythonx.eval(
+              """
+              import os
+              os.environ["MPLBACKEND"] = "Agg"
+              """,
+              %{}
+            )
+
+            value = :ok
             result = {:ok, value, binding, env}
             code_markers = []
             {result, code_markers}
@@ -1081,6 +1100,27 @@ defmodule Livebook.Runtime.Evaluator do
 
   defp pythonx_version(), do: List.to_string(Application.spec(:pythonx)[:vsn])
 
+  defp after_evaluation(:python) do
+    if ensure_pythonx() == :ok do
+      # With matplotlib the charts are built imperatively, by modifying
+      # a global figure state. We clear the global state after the
+      # evaluation, otherwise re-evaluating cells draws on top of the
+      # previous figure. We do this only if matplotlib is imported.
+      Pythonx.eval(
+        """
+        import sys
+
+        if "matplotlib" in sys.modules:
+          import matplotlib.pyplot as plt
+          plt.close("all")
+        """,
+        %{}
+      )
+    end
+  end
+
+  defp after_evaluation(_language), do: :ok
+
   defp identifier_dependencies(context, tracer_info, prev_context) do
     identifiers_used = MapSet.new()
     identifiers_defined = %{}

lib/livebook/runtime/evaluator/formatter.ex

@@ -2,6 +2,7 @@ defmodule Livebook.Runtime.Evaluator.Formatter do
   require Logger
 
   @compile {:no_warn_undefined, {Kino.Render, :to_livebook, 1}}
+  @compile {:no_warn_undefined, {Kino.Render, :impl_for, 1}}
   @compile {:no_warn_undefined, {Pythonx, :eval, 2}}
   @compile {:no_warn_undefined, {Pythonx, :decode, 1}}
 
@@ -57,8 +58,22 @@ defmodule Livebook.Runtime.Evaluator.Formatter do
   end
 
   def format_result(:python, {:ok, value}) do
-    repr_string = Pythonx.eval("repr(value)", %{"value" => value}) |> elem(0) |> Pythonx.decode()
-    %{type: :terminal_text, text: repr_string, chunk: false}
+    if Code.ensure_loaded?(Kino.Render) do
+      try do
+        if Kino.Render.impl_for(value) == Kino.Render.Any do
+          to_repr_output(value)
+        else
+          Kino.Render.to_livebook(value)
+        end
+      catch
+        kind, error ->
+          formatted = format_error(kind, error, __STACKTRACE__)
+          Logger.error(formatted)
+          to_repr_output(value)
+      end
+    else
+      to_repr_output(value)
+    end
   end
 
   def format_result(:python, {:error, _kind, error, _stacktrace})
@@ -108,6 +123,13 @@ defmodule Livebook.Runtime.Evaluator.Formatter do
     end
   end
 
+  defp to_repr_output(value) do
+    repr_string =
+      Pythonx.eval("repr(value)", %{"value" => value}) |> elem(0) |> Pythonx.decode()
+
+    %{type: :terminal_text, text: repr_string, chunk: false}
+  end
+
   defp to_inspect_output(value, opts \\ []) do
     try do
       inspected = inspect(value, inspect_opts(opts))

lib/livebook/session.ex

@@ -1243,7 +1243,12 @@ defmodule Livebook.Session do
   end
 
   def handle_cast({:enable_language, client_pid, language}, state) do
-    case do_add_dependencies(state, [Livebook.Runtime.Definitions.pythonx_dependency()]) do
+    dependencies = [
+      Livebook.Runtime.Definitions.pythonx_dependency(),
+      Livebook.Runtime.Definitions.kino_pythonx_dependency()
+    ]
+
+    case do_add_dependencies(state, dependencies) do
       {:ok, state} ->
         client_id = client_id(state, client_pid)