Fix stderr not being captured in Python cells (#2941)

lib/livebook/runtime/definitions.ex

@@ -506,6 +506,6 @@ defmodule Livebook.Runtime.Definitions do
   def snippet_definitions(), do: @snippet_definitions
 
   def pythonx_dependency() do
-    %{dep: {:pythonx, "~> 0.3.0"}, config: []}
+    %{dep: {:pythonx, "~> 0.4.0"}, config: []}
   end
 end

lib/livebook/runtime/erl_dist/io_forward_gl.ex

@@ -42,8 +42,8 @@ defmodule Livebook.Runtime.ErlDist.IOForwardGL do
   def handle_info({:io_request, from, reply_as, req}, state) do
     case Process.info(from, :group_leader) do
       {:group_leader, group_leader} ->
-        # Forward the request to sender's group leader
-        # and instruct it to get back to us.
+        # Forward the request to sender's group leader and instruct
+        # it to get back to the sender.
         send(group_leader, {:io_request, from, reply_as, req})
 
       _ ->

lib/livebook/runtime/evaluator.ex

@@ -1005,7 +1005,25 @@ defmodule Livebook.Runtime.Evaluator do
       |> Map.replace!(:requires, [Pythonx])
       |> Map.replace!(:macros, [{Pythonx, [{:sigil_PY, 2}]}])
 
-    Macro.expand_once(quoted, env)
+    ast = Macro.expand_once(quoted, env)
+
+    # We modify the Pythonx.eval/2 call to specify the :stderr_device
+    # option. We want to Python stderr output to also be send to our
+    # group leader. By default it would be sent to our :standard_error,
+    # which sends it further to sender's group leader, however the
+    # sender is a process in the Pythonx supervision tree and has the
+    # default group leader.mix
+    Macro.prewalk(ast, fn
+      {{:., _, [{:__aliases__, _, [:Pythonx]}, :eval]} = target, meta, [code, globals]} ->
+        opts = [
+          stderr_device: {{:., [], [{:__aliases__, [], [:Process]}, :group_leader]}, [], []}
+        ]
+
+        {target, meta, [code, globals, opts]}
+
+      other ->
+        other
+    end)
   end
 
   defp eval_pyproject_toml(code, binding, env) do

mix.exs

@@ -125,7 +125,7 @@ defmodule Livebook.MixProject do
       {:floki, ">= 0.27.0", only: :test},
       {:bypass, "~> 2.1", only: :test},
       # So that we can test Python evaluation in the same node
-      {:pythonx, "~> 0.3.0", only: :test},
+      {:pythonx, "~> 0.4.0", only: :test},
       # ZTA deps
       {:jose, "~> 1.11.5"},
       {:req, "~> 0.5.8"},

mix.lock

@@ -47,7 +47,7 @@
   "plug_crypto": {:hex, :plug_crypto, "2.1.0", "f44309c2b06d249c27c8d3f65cfe08158ade08418cf540fd4f72d4d6863abb7b", [:mix], [], "hexpm", "131216a4b030b8f8ce0f26038bc4421ae60e4bb95c5cf5395e1421437824c4fa"},
   "pluggable": {:hex, :pluggable, "1.1.0", "7eba3bc70c0caf4d9056c63c882df8862f7534f0145da7ab3a47ca73e4adb1e4", [:mix], [], "hexpm", "d12eb00ea47b21e92cd2700d6fbe3737f04b64e71b63aad1c0accde87c751637"},
   "protobuf": {:hex, :protobuf, "0.13.0", "7a9d9aeb039f68a81717eb2efd6928fdf44f03d2c0dfdcedc7b560f5f5aae93d", [:mix], [{:jason, "~> 1.2", [hex: :jason, repo: "hexpm", optional: true]}], "hexpm", "21092a223e3c6c144c1a291ab082a7ead32821ba77073b72c68515aa51fef570"},
-  "pythonx": {:hex, :pythonx, "0.3.0", "61214452c993aa010a6674910ea4db96a55c5aa88e4562e615a89de59539ecf5", [:make, :mix], [{:cc_precompiler, "~> 0.1", [hex: :cc_precompiler, repo: "hexpm", optional: false]}, {:elixir_make, "~> 0.9", [hex: :elixir_make, repo: "hexpm", optional: false]}, {:fine, "~> 0.1.0", [hex: :fine, repo: "hexpm", optional: false]}], "hexpm", "b2402ed149844718f1b24a8a380143970e3e92e1bb132a42050b05750b02e449"},
+  "pythonx": {:hex, :pythonx, "0.4.0", "17528698aef5162ae83d8b4aa9df063cf15bb3acd3caa56b61eefc75ba7bad93", [:make, :mix], [{:cc_precompiler, "~> 0.1", [hex: :cc_precompiler, repo: "hexpm", optional: false]}, {:elixir_make, "~> 0.9", [hex: :elixir_make, repo: "hexpm", optional: false]}, {:fine, "~> 0.1.0", [hex: :fine, repo: "hexpm", optional: false]}], "hexpm", "801670ef6eb2eff2ff9ec00f70d14d4e4cd0ed0d6af02fe8c8e94954d56078ab"},
   "ranch": {:hex, :ranch, "1.8.0", "8c7a100a139fd57f17327b6413e4167ac559fbc04ca7448e9be9057311597a1d", [:make, :rebar3], [], "hexpm", "49fbcfd3682fab1f5d109351b61257676da1a2fdbe295904176d5e521a2ddfe5"},
   "req": {:hex, :req, "0.5.8", "50d8d65279d6e343a5e46980ac2a70e97136182950833a1968b371e753f6a662", [:mix], [{:brotli, "~> 0.3.1", [hex: :brotli, repo: "hexpm", optional: true]}, {:ezstd, "~> 1.0", [hex: :ezstd, repo: "hexpm", optional: true]}, {:finch, "~> 0.17", [hex: :finch, repo: "hexpm", optional: false]}, {:jason, "~> 1.0", [hex: :jason, repo: "hexpm", optional: false]}, {:mime, "~> 2.0.6 or ~> 2.1", [hex: :mime, repo: "hexpm", optional: false]}, {:nimble_csv, "~> 1.0", [hex: :nimble_csv, repo: "hexpm", optional: true]}, {:plug, "~> 1.0", [hex: :plug, repo: "hexpm", optional: true]}], "hexpm", "d7fc5898a566477e174f26887821a3c5082b243885520ee4b45555f5d53f40ef"},
   "telemetry": {:hex, :telemetry, "1.3.0", "fedebbae410d715cf8e7062c96a1ef32ec22e764197f70cda73d82778d61e7a2", [:rebar3], [], "hexpm", "7015fc8919dbe63764f4b4b87a95b7c0996bd539e0d499be6ec9d7f3875b79e6"},

test/livebook/live_markdown/export_test.exs

@@ -1171,7 +1171,7 @@ defmodule Livebook.LiveMarkdown.ExportTest do
             Notebook.Cell.new(:code)
             | source: """
               Mix.install([
-                {:pythonx, "~> 0.3.0"}
+                {:pythonx, "~> 0.4.0"}
               ])\
               """
           },
@@ -1193,7 +1193,7 @@ defmodule Livebook.LiveMarkdown.ExportTest do
 
       ```elixir
       Mix.install([
-        {:pythonx, "~> 0.3.0"}
+        {:pythonx, "~> 0.4.0"}
       ])
       ```
 

test/livebook/live_markdown/import_test.exs

@@ -1157,7 +1157,7 @@ defmodule Livebook.LiveMarkdown.ImportTest do
 
       ```elixir
       Mix.install([
-        {:pythonx, "~> 0.3.0"}
+        {:pythonx, "~> 0.4.0"}
       ])
       ```
 
@@ -1180,7 +1180,7 @@ defmodule Livebook.LiveMarkdown.ImportTest do
                      id: "setup",
                      source: """
                      Mix.install([
-                       {:pythonx, "~> 0.3.0"}
+                       {:pythonx, "~> 0.4.0"}
                      ])\
                      """
                    },

test/livebook/notebook/export/elixir_test.exs

@@ -203,7 +203,7 @@ defmodule Livebook.Notebook.Export.ElixirTest do
           Notebook.Cell.new(:code)
           | source: """
             Mix.install([
-              {:pythonx, "~> 0.3.0"}
+              {:pythonx, "~> 0.4.0"}
             ])\
             """
         },
@@ -226,7 +226,7 @@ defmodule Livebook.Notebook.Export.ElixirTest do
     # Title: My Notebook
 
     Mix.install([
-      {:pythonx, "~> 0.3.0"}
+      {:pythonx, "~> 0.4.0"}
     ])
 
     Pythonx.uv_init("""

test/livebook/runtime/evaluator_test.exs

@@ -1659,6 +1659,29 @@ defmodule Livebook.Runtime.EvaluatorTest do
              ModuleNotFoundError: No module named 'unknown'
              """
     end
+
+    test "captures standard output and sends it to the caller", %{evaluator: evaluator} do
+      code = """
+      print("hello from Python")
+      """
+
+      Evaluator.evaluate_code(evaluator, :python, code, :code_1, [])
+
+      assert_receive {:runtime_evaluation_output, :code_1,
+                      terminal_text("hello from Python\n", true)}
+    end
+
+    test "captures standard error and sends it to the caller", %{evaluator: evaluator} do
+      code = """
+      import sys
+      print("error from Python", file=sys.stderr)
+      """
+
+      Evaluator.evaluate_code(evaluator, :python, code, :code_1, [])
+
+      assert_receive {:runtime_evaluation_output, :code_1,
+                      terminal_text("error from Python\n", true)}
+    end
   end
 
   describe "formatting" do