Set proper file in env when the notebook is persisted (#95)

* Set proper file in env when the notebook is persisted

* Fix typo

lib/livebook/evaluator.ex

@@ -64,10 +64,16 @@ defmodule Livebook.Evaluator do
 
   Evaluation response is sent to the process identified by `send_to` as `{:evaluation_response, ref, response}`.
   Note that response is transformed with the configured formatter (identity by default).
+
+  ## Options
+
+  * `:file` - file to which the evaluated code belongs. Most importantly,
+    this has an impact on the value of `__DIR__`.
   """
-  @spec evaluate_code(t(), pid(), String.t(), ref(), ref()) :: :ok
+  @spec evaluate_code(t(), pid(), String.t(), ref(), ref(), keyword()) :: :ok
-  def evaluate_code(evaluator, send_to, code, ref, prev_ref \\ :initial) when ref != :initial do
+  def evaluate_code(evaluator, send_to, code, ref, prev_ref \\ :initial, opts \\ [])
-    GenServer.cast(evaluator, {:evaluate_code, send_to, code, ref, prev_ref})
+      when ref != :initial do
+    GenServer.cast(evaluator, {:evaluate_code, send_to, code, ref, prev_ref, opts})
   end
 
   @doc """
@@ -107,11 +113,14 @@ defmodule Livebook.Evaluator do
   end
 
   @impl true
-  def handle_cast({:evaluate_code, send_to, code, ref, prev_ref}, state) do
+  def handle_cast({:evaluate_code, send_to, code, ref, prev_ref, opts}, state) do
     Evaluator.IOProxy.configure(state.io_proxy, send_to, ref)
 
     context = Map.get(state.contexts, prev_ref, state.contexts.initial)
 
+    file = Keyword.get(opts, :file, "nofile")
+    context = put_in(context.env.file, file)
+
     case eval(code, context.binding, context.env) do
       {:ok, result, binding, env} ->
         result_context = %{binding: binding, env: env}

lib/livebook/runtime.ex

@@ -41,9 +41,21 @@ defprotocol Livebook.Runtime do
 
   * `{:evaluation_stdout, ref, string}` - output captured during evaluation
   * `{:evaluation_response, ref, response}` - final result of the evaluation
+
+  ## Options
+
+  * `:file` - file to which the evaluated code belongs. Most importantly,
+    this has an impact on the value of `__DIR__`.
   """
-  @spec evaluate_code(t(), String.t(), ref(), ref(), ref()) :: :ok
+  @spec evaluate_code(t(), String.t(), ref(), ref(), ref(), keyword()) :: :ok
-  def evaluate_code(runtime, code, container_ref, evaluation_ref, prev_evaluation_ref \\ :initial)
+  def evaluate_code(
+        runtime,
+        code,
+        container_ref,
+        evaluation_ref,
+        prev_evaluation_ref \\ :initial,
+        opts \\ []
+      )
 
   @doc """
   Disposes of evaluation identified by the given ref.

lib/livebook/runtime/attached.ex

@@ -49,13 +49,21 @@ defimpl Livebook.Runtime, for: Livebook.Runtime.Attached do
     ErlDist.Manager.stop(runtime.node)
   end
 
-  def evaluate_code(runtime, code, container_ref, evaluation_ref, prev_evaluation_ref \\ :initial) do
+  def evaluate_code(
+        runtime,
+        code,
+        container_ref,
+        evaluation_ref,
+        prev_evaluation_ref \\ :initial,
+        opts \\ []
+      ) do
     ErlDist.Manager.evaluate_code(
       runtime.node,
       code,
       container_ref,
       evaluation_ref,
-      prev_evaluation_ref
+      prev_evaluation_ref,
+      opts
     )
   end
 

lib/livebook/runtime/elixir_standalone.ex

@@ -73,13 +73,21 @@ defimpl Livebook.Runtime, for: Livebook.Runtime.ElixirStandalone do
     ErlDist.Manager.stop(runtime.node)
   end
 
-  def evaluate_code(runtime, code, container_ref, evaluation_ref, prev_evaluation_ref \\ :initial) do
+  def evaluate_code(
+        runtime,
+        code,
+        container_ref,
+        evaluation_ref,
+        prev_evaluation_ref \\ :initial,
+        opts \\ []
+      ) do
     ErlDist.Manager.evaluate_code(
       runtime.node,
       code,
       container_ref,
       evaluation_ref,
-      prev_evaluation_ref
+      prev_evaluation_ref,
+      opts
     )
   end
 

lib/livebook/runtime/erl_dist/manager.ex

@@ -49,12 +49,25 @@ defmodule Livebook.Runtime.ErlDist.Manager do
 
   See `Evaluator` for more details.
   """
-  @spec evaluate_code(node(), String.t(), Evaluator.ref(), Evaluator.ref(), Evaluator.ref()) ::
+  @spec evaluate_code(
-          :ok
+          node(),
-  def evaluate_code(node, code, container_ref, evaluation_ref, prev_evaluation_ref \\ :initial) do
+          String.t(),
+          Evaluator.ref(),
+          Evaluator.ref(),
+          Evaluator.ref(),
+          keyword()
+        ) :: :ok
+  def evaluate_code(
+        node,
+        code,
+        container_ref,
+        evaluation_ref,
+        prev_evaluation_ref \\ :initial,
+        opts \\ []
+      ) do
     GenServer.cast(
       {@name, node},
-      {:evaluate_code, code, container_ref, evaluation_ref, prev_evaluation_ref}
+      {:evaluate_code, code, container_ref, evaluation_ref, prev_evaluation_ref, opts}
     )
   end
 
@@ -153,7 +166,7 @@ defmodule Livebook.Runtime.ErlDist.Manager do
   end
 
   def handle_cast(
-        {:evaluate_code, code, container_ref, evaluation_ref, prev_evaluation_ref},
+        {:evaluate_code, code, container_ref, evaluation_ref, prev_evaluation_ref, opts},
         state
       ) do
     state = ensure_evaluator(state, container_ref)
@@ -163,7 +176,8 @@ defmodule Livebook.Runtime.ErlDist.Manager do
       state.owner,
       code,
       evaluation_ref,
-      prev_evaluation_ref
+      prev_evaluation_ref,
+      opts
     )
 
     {:noreply, state}

lib/livebook/runtime/mix_standalone.ex

@@ -109,13 +109,21 @@ defimpl Livebook.Runtime, for: Livebook.Runtime.MixStandalone do
     ErlDist.Manager.stop(runtime.node)
   end
 
-  def evaluate_code(runtime, code, container_ref, evaluation_ref, prev_evaluation_ref \\ :initial) do
+  def evaluate_code(
+        runtime,
+        code,
+        container_ref,
+        evaluation_ref,
+        prev_evaluation_ref \\ :initial,
+        opts \\ []
+      ) do
     ErlDist.Manager.evaluate_code(
       runtime.node,
       code,
       container_ref,
       evaluation_ref,
-      prev_evaluation_ref
+      prev_evaluation_ref,
+      opts
     )
   end
 

lib/livebook/session.ex

@@ -546,7 +546,10 @@ defmodule Livebook.Session do
         [] -> :initial
       end
 
-    Runtime.evaluate_code(state.data.runtime, cell.source, :main, cell.id, prev_ref)
+    file = (state.data.path || "") <> "#cell"
+    opts = [file: file]
+
+    Runtime.evaluate_code(state.data.runtime, cell.source, :main, cell.id, prev_ref, opts)
 
     state
   end

test/livebook/evaluator_test.exs

@@ -8,7 +8,7 @@ defmodule Livebook.EvaluatorTest do
     %{evaluator: evaluator}
   end
 
-  describe "evaluate_code/4" do
+  describe "evaluate_code/6" do
     test "given a valid code returns evaluation result", %{evaluator: evaluator} do
       code = """
       x = 1
@@ -107,6 +107,17 @@ defmodule Livebook.EvaluatorTest do
                        1000
       end)
     end
+
+    test "given file option sets it in evaluation environment", %{evaluator: evaluator} do
+      code = """
+      __DIR__
+      """
+
+      opts = [file: "/path/dir/file"]
+      Evaluator.evaluate_code(evaluator, self(), code, :code_1, :initial, opts)
+
+      assert_receive {:evaluation_response, :code_1, {:ok, "/path/dir"}}
+    end
   end
 
   describe "forget_evaluation/2" do

test/livebook/runtime/erl_dist/manager_test.exs

@@ -28,7 +28,7 @@ defmodule Livebook.Runtime.ErlDist.ManagerTest do
     end
   end
 
-  describe "evaluate_code/2" do
+  describe "evaluate_code/6" do
     test "spawns a new evaluator when necessary" do
       Manager.start()
       Manager.set_owner(node(), self())

test/support/single_evaluator_runtime.ex

@@ -30,9 +30,17 @@ defimpl Livebook.Runtime, for: LivebookTest.Runtime.SingleEvaluator do
         code,
         _container_ref,
         evaluation_ref,
-        prev_evaluation_ref \\ :initial
+        prev_evaluation_ref \\ :initial,
+        opts \\ []
       ) do
-    Evaluator.evaluate_code(runtime.evaluator, self(), code, evaluation_ref, prev_evaluation_ref)
+    Evaluator.evaluate_code(
+      runtime.evaluator,
+      self(),
+      code,
+      evaluation_ref,
+      prev_evaluation_ref,
+      opts
+    )
 
     :ok
   end