diff --git a/lib/livebook.ex b/lib/livebook.ex index 2cfb02c8d..d70165a42 100644 --- a/lib/livebook.ex +++ b/lib/livebook.ex @@ -30,7 +30,7 @@ defmodule Livebook do [ %{ - dependency: {:kino, "~> 0.6.1"}, + dependency: %{dep: {:kino, "~> 0.6.1"}, config: []}, description: "Interactive widgets for Livebook", name: "kino", url: "https://hex.pm/packages/kino", diff --git a/lib/livebook/runtime/dependencies.ex b/lib/livebook/runtime/dependencies.ex index cefc2831a..723216996 100644 --- a/lib/livebook/runtime/dependencies.ex +++ b/lib/livebook/runtime/dependencies.ex @@ -1,14 +1,37 @@ defmodule Livebook.Runtime.Dependencies do @moduledoc false + @doc """ + Adds the given list of dependencies to the setup code. + """ + @spec add_dependencies(String.t(), list(Livebook.Runtime.dependency())) :: + {:ok, String.t()} | {:error, String.t()} + def add_dependencies(code, dependencies) do + deps = Enum.map(dependencies, & &1.dep) + config = Enum.reduce(dependencies, [], &deep_merge(&2, &1.config)) + add_mix_deps(code, deps, config) + end + + defp deep_merge(left, right) do + if Keyword.keyword?(left) and Keyword.keyword?(right) do + Keyword.merge(left, right, fn _key, left, right -> deep_merge(left, right) end) + else + right + end + end + @doc """ Finds or adds a `Mix.install/2` call to `code` and modifies it to include the given Mix deps. + + A config may be given, in which case it is added or merged into the + `Mix.install/2` `:config` option. """ - @spec add_mix_deps(String.t(), list(tuple())) :: {:ok, String.t()} | {:error, String.t()} - def add_mix_deps(code, deps) do + @spec add_mix_deps(String.t(), list(tuple()), keyword()) :: + {:ok, String.t()} | {:error, String.t()} + def add_mix_deps(code, deps, config \\ []) do with {:ok, ast, comments} <- string_to_quoted_with_comments(code), - {:ok, ast} <- insert_deps(ast, deps), + {:ok, ast} <- insert_deps(ast, deps, config), do: {:ok, format(ast, comments)} end @@ -27,13 +50,18 @@ defmodule Livebook.Runtime.Dependencies do end end - defp insert_deps(ast, deps) do - with :error <- update_install(ast, deps) do + defp insert_deps(ast, deps, config) do + with :error <- update_install(ast, deps, config) do dep_nodes = Enum.map(deps, &dep_node/1) - install_node = - {{:., [], [{:__aliases__, [], [:Mix]}, :install]}, [], - [{:__block__, [newlines: 1], [dep_nodes]}]} + install_args = + [{:__block__, [newlines: 1], [dep_nodes]}] ++ + case config do + [] -> [] + config -> [[config: Macro.escape(config)]] + end + + install_node = {{:., [], [{:__aliases__, [], [:Mix]}, :install]}, [], install_args} {:ok, prepend_node(ast, install_node)} end @@ -57,17 +85,31 @@ defmodule Livebook.Runtime.Dependencies do defp update_install( {{:., _, [{:__aliases__, _, [:Mix]}, :install]} = target, meta1, [{:__block__, meta2, [dep_nodes]} | args]}, - deps + deps, + config ) do + args = + case {args, config} do + {args, []} -> + args + + {[], config} -> + [[config: Macro.escape(config)]] + + {[opts | other_args], config} -> + opts = map_maybe_block(opts, &update_config_in_opts(&1, config)) + [opts | other_args] + end + new_dep_nodes = for dep <- deps, not has_dep?(dep_nodes, dep), do: dep_node(dep) {:ok, {target, meta1, [{:__block__, meta2, [dep_nodes ++ new_dep_nodes]} | args]}} end - defp update_install({:__block__, meta, nodes}, deps) do + defp update_install({:__block__, meta, nodes}, deps, config) do {nodes, found} = Enum.map_reduce(nodes, _found = false, fn node, false -> - case update_install(node, deps) do + case update_install(node, deps, config) do {:ok, node} -> {node, true} _ -> {node, false} end @@ -83,7 +125,7 @@ defmodule Livebook.Runtime.Dependencies do end end - defp update_install(_node, _deps), do: :error + defp update_install(_node, _deps, _config), do: :error defp has_dep?(deps, dep) do name = elem(dep, 0) @@ -97,6 +139,41 @@ defmodule Livebook.Runtime.Dependencies do defp dep_node(dep), do: {:__block__, [], [Macro.escape(dep)]} + defp update_config_in_opts( + [{{:__block__, meta1, [:config]}, {:__block__, meta2, [current_config]}} | tail], + config + ) do + config = append_config(current_config, config) + [{{:__block__, meta1, [:config]}, {:__block__, meta2, [config]}} | tail] + end + + defp update_config_in_opts([head | tail], config), + do: [head | update_config_in_opts(tail, config)] + + defp update_config_in_opts([], config), do: [config: Macro.escape(config)] + + defp append_config(current_config, config) do + # Note: the current config has literals wrapped in a :__block__, + # while the new one doesn't. + + # If the given key is already in the config, we ignore it + existing_keys = for {{:__block__, _meta, [key]}, _value} <- current_config, do: key + config = Keyword.drop(config, existing_keys) + + # We need to wrap literals in a block, as in the rest of the AST + to_quoted_opts = [literal_encoder: &{:ok, {:__block__, &2, [&1]}}] + + {:__block__, _, [items]} = + config + |> Macro.to_string() + |> Code.string_to_quoted!(to_quoted_opts) + + current_config ++ items + end + + defp map_maybe_block({:__block__, meta, [value]}, fun), do: {:__block__, meta, [fun.(value)]} + defp map_maybe_block(value, fun), do: fun.(value) + @doc """ Parses a plain Elixir term from its string representation. @@ -225,14 +302,14 @@ defmodule Livebook.Runtime.Dependencies do end defp parse_package(package) do - {:ok, dependency} = parse_term(package["configs"]["mix.exs"]) + {:ok, dep} = parse_term(package["configs"]["mix.exs"]) %{ name: package["name"], version: package["latest_stable_version"] || package["latest_version"], description: package["meta"]["description"], url: package["html_url"], - dependency: dependency + dependency: %{dep: dep, config: []} } end diff --git a/lib/livebook/runtime/elixir_standalone.ex b/lib/livebook/runtime/elixir_standalone.ex index 781a32d81..ad393caac 100644 --- a/lib/livebook/runtime/elixir_standalone.ex +++ b/lib/livebook/runtime/elixir_standalone.ex @@ -17,9 +17,20 @@ defmodule Livebook.Runtime.ElixirStandalone do server_pid: pid() | nil } - kino_vega_lite = %{name: "kino_vega_lite", dependency: {:kino_vega_lite, "~> 0.1.4"}} - kino_db = %{name: "kino_db", dependency: {:kino_db, "~> 0.2.0"}} - kino_maplibre = %{name: "kino_maplibre", dependency: {:kino_maplibre, "~> 0.1.3"}} + kino_vega_lite = %{ + name: "kino_vega_lite", + dependency: %{dep: {:kino_vega_lite, "~> 0.1.4"}, config: []} + } + + kino_db = %{ + name: "kino_db", + dependency: %{dep: {:kino_db, "~> 0.2.0"}, config: []} + } + + kino_maplibre = %{ + name: "kino_maplibre", + dependency: %{dep: {:kino_maplibre, "~> 0.1.3"}, config: []} + } @extra_smart_cell_definitions [ %{ @@ -29,23 +40,44 @@ defmodule Livebook.Runtime.ElixirStandalone do variants: [ %{ name: "Amazon Athena", - packages: [kino_db, %{name: "req_athena", dependency: {:req_athena, "~> 0.1.0"}}] + packages: [ + kino_db, + %{ + name: "req_athena", + dependency: %{dep: {:req_athena, "~> 0.1.0"}, config: []} + } + ] }, %{ name: "Google BigQuery", - packages: [kino_db, %{name: "req_bigquery", dependency: {:req_bigquery, "~> 0.1.0"}}] + packages: [ + kino_db, + %{ + name: "req_bigquery", + dependency: %{dep: {:req_bigquery, "~> 0.1.0"}, config: []} + } + ] }, %{ name: "MySQL", - packages: [kino_db, %{name: "myxql", dependency: {:myxql, "~> 0.6.2"}}] + packages: [ + kino_db, + %{name: "myxql", dependency: %{dep: {:myxql, "~> 0.6.2"}, config: []}} + ] }, %{ name: "PostgreSQL", - packages: [kino_db, %{name: "postgrex", dependency: {:postgrex, "~> 0.16.3"}}] + packages: [ + kino_db, + %{name: "postgrex", dependency: %{dep: {:postgrex, "~> 0.16.3"}, config: []}} + ] }, %{ name: "SQLite", - packages: [kino_db, %{name: "exqlite", dependency: {:exqlite, "~> 0.11.0"}}] + packages: [ + kino_db, + %{name: "exqlite", dependency: %{dep: {:exqlite, "~> 0.11.0"}, config: []}} + ] } ] } @@ -214,7 +246,7 @@ defimpl Livebook.Runtime, for: Livebook.Runtime.ElixirStandalone do def fixed_dependencies?(_runtime), do: false def add_dependencies(_runtime, code, dependencies) do - Livebook.Runtime.Dependencies.add_mix_deps(code, dependencies) + Livebook.Runtime.Dependencies.add_dependencies(code, dependencies) end def search_packages(_runtime, send_to, search) do diff --git a/lib/livebook/runtime/embedded.ex b/lib/livebook/runtime/embedded.ex index 2ed49af1a..b8f4d4ac7 100644 --- a/lib/livebook/runtime/embedded.ex +++ b/lib/livebook/runtime/embedded.ex @@ -113,7 +113,7 @@ defimpl Livebook.Runtime, for: Livebook.Runtime.Embedded do end def add_dependencies(_runtime, code, dependencies) do - Livebook.Runtime.Dependencies.add_mix_deps(code, dependencies) + Livebook.Runtime.Dependencies.add_dependencies(code, dependencies) end def search_packages(_runtime, send_to, search) do diff --git a/test/livebook/runtime/dependencies_test.exs b/test/livebook/runtime/dependencies_test.exs index a82dee5c5..7dea5dc94 100644 --- a/test/livebook/runtime/dependencies_test.exs +++ b/test/livebook/runtime/dependencies_test.exs @@ -7,7 +7,31 @@ defmodule Livebook.Runtime.DependenciesTest do @jason {:jason, "~> 1.3.0"} - describe "add_mix_deps/2" do + describe "add_dependencies/3" do + test "adds dependencies and config" do + assert Dependencies.add_dependencies("", [ + %{dep: {:nx, "~> 0.4.0"}, config: []}, + %{ + dep: {:exla, "~> 0.4.0"}, + config: [nx: [default_defn_options: [compiler: EXLA]]] + }, + %{dep: {:torchx, "~> 0.4.0"}, config: [nx: [default_backend: Torchx.Backend]]} + ]) == + {:ok, + """ + Mix.install( + [ + {:nx, "~> 0.4.0"}, + {:exla, "~> 0.4.0"}, + {:torchx, "~> 0.4.0"} + ], + config: [nx: [default_defn_options: [compiler: EXLA], default_backend: Torchx.Backend]] + )\ + """} + end + end + + describe "add_mix_deps/3" do test "prepends Mix.install/2 call if there is none" do assert Dependencies.add_mix_deps("", [@jason]) == {:ok, @@ -188,6 +212,110 @@ defmodule Livebook.Runtime.DependenciesTest do | ^\ """} end + + test "adds config if specified" do + config = [nx: [default_backend: EXLA.Backend]] + + assert Dependencies.add_mix_deps("", [@jason], config) == + {:ok, + """ + Mix.install( + [ + {:jason, "~> 1.3.0"} + ], + config: [nx: [default_backend: EXLA.Backend]] + )\ + """} + + assert Dependencies.add_mix_deps( + """ + Mix.install([ + {:jason, "~> 1.3.0"} + ])\ + """, + [], + config + ) == + {:ok, + """ + Mix.install( + [ + {:jason, "~> 1.3.0"} + ], + config: [nx: [default_backend: EXLA.Backend]] + )\ + """} + end + + test "merges config in flat manner" do + assert Dependencies.add_mix_deps( + """ + Mix.install( + [ + {:jason, "~> 1.3.0"} + ], + config: [ + # Comment 1 + nx: [ + # Comment 2 + default_backend: Torchx.Backend + # Comment 3 + ], + test: [x: :y] + # Comment 4 + ] + )\ + """, + [], + nx: [ + default_defn_options: [compiler: EXLA] + ], + other: [ + default_defn_options: [compiler: EXLA] + ] + ) == + {:ok, + """ + Mix.install( + [ + {:jason, "~> 1.3.0"} + ], + config: [ + # Comment 1 + nx: [ + # Comment 2 + default_backend: Torchx.Backend + # Comment 3 + ], + test: [x: :y], + other: [default_defn_options: [compiler: EXLA]] + # Comment 4 + ] + )\ + """} + + assert Dependencies.add_mix_deps( + """ + Mix.install( + [ + {:jason, "~> 1.3.0"} + ], + [config: []] + )\ + """, + [], + nx: [default_backend: EXLA.Backend] + ) == + {:ok, + """ + Mix.install( + [ + {:jason, "~> 1.3.0"} + ], + config: [nx: [default_backend: EXLA.Backend]] + )\ + """} + end end describe "search_hex/2" do @@ -248,7 +376,7 @@ defmodule Livebook.Runtime.DependenciesTest do {:ok, [ %{ - dependency: {:ecto, "~> 3.7"}, + dependency: %{dep: {:ecto, "~> 3.7"}, config: []}, description: "A toolkit for data mapping and language integrated query for Elixir", name: "ecto", @@ -256,7 +384,7 @@ defmodule Livebook.Runtime.DependenciesTest do version: "3.7.2" }, %{ - dependency: {:ecto_sql, "~> 3.7"}, + dependency: %{dep: {:ecto_sql, "~> 3.7"}, config: []}, description: "SQL-based adapters for Ecto and database migrations", name: "ecto_sql", url: "https://hex.pm/packages/ecto_sql", diff --git a/test/livebook/session_test.exs b/test/livebook/session_test.exs index e801163fa..e0f75e4b8 100644 --- a/test/livebook/session_test.exs +++ b/test/livebook/session_test.exs @@ -202,7 +202,7 @@ defmodule Livebook.SessionTest do Session.subscribe(session.id) - Session.add_dependencies(session.pid, [{:jason, "~> 1.3.0"}]) + Session.add_dependencies(session.pid, [%{dep: {:jason, "~> 1.3.0"}, config: []}]) assert_receive {:operation, {:apply_cell_delta, "__server__", "setup", :primary, _delta, 1}} @@ -232,7 +232,7 @@ defmodule Livebook.SessionTest do Session.subscribe(session.id) - Session.add_dependencies(session.pid, [{:json, "~> 1.3.0"}]) + Session.add_dependencies(session.pid, [%{dep: {:jason, "~> 1.3.0"}, config: []}]) assert_receive {:error, "failed to add dependencies to the setup cell, reason:" <> _} end diff --git a/test/support/noop_runtime.ex b/test/support/noop_runtime.ex index dd0e65ce2..f6a0e91c1 100644 --- a/test/support/noop_runtime.ex +++ b/test/support/noop_runtime.ex @@ -34,7 +34,7 @@ defmodule Livebook.Runtime.NoopRuntime do def fixed_dependencies?(_), do: false def add_dependencies(_runtime, code, dependencies) do - Livebook.Runtime.Dependencies.add_mix_deps(code, dependencies) + Livebook.Runtime.Dependencies.add_dependencies(code, dependencies) end def search_packages(_, _, _), do: make_ref() diff --git a/test/test_helper.exs b/test/test_helper.exs index 94bcbbbd5..18b0595e4 100644 --- a/test/test_helper.exs +++ b/test/test_helper.exs @@ -21,7 +21,7 @@ defmodule Livebook.Runtime.Embedded.Packages do def list() do [ %{ - dependency: {:jason, "~> 1.3.0"}, + dependency: %{dep: {:jason, "~> 1.3.0"}, config: []}, description: "A blazing fast JSON parser and generator in pure Elixir", name: "jason", url: "https://hex.pm/packages/jason",