diff --git a/lib/ex_ast/pattern.ex b/lib/ex_ast/pattern.ex index 4d82998..50bfd7d 100644 --- a/lib/ex_ast/pattern.ex +++ b/lib/ex_ast/pattern.ex @@ -333,12 +333,40 @@ defmodule ExAST.Pattern do defp collect_import_directive(_args, acc), do: acc defp import_only(opts) do - case Keyword.get(opts, :only) do - only when is_list(only) -> only - _ -> :all + case opt_value(opts, :only) do + nil -> :all + only_ast -> parse_only_list(only_ast) end end + defp opt_value(opts, key) do + Enum.find_value(opts, fn + {{:__block__, _, [^key]}, value} -> value + {^key, value} -> value + _other -> nil + end) + end + + defp parse_only_list(ast) do + case unwrap_block(ast) do + list when is_list(list) -> Enum.flat_map(list, &parse_only_entry/1) + # `only: :functions` / `:macros` carry no per-function membership. + _other -> :all + end + end + + defp parse_only_entry({name_ast, arity_ast}) do + case {unwrap_block(name_ast), unwrap_block(arity_ast)} do + {name, arity} when is_atom(name) and is_integer(arity) -> [{name, arity}] + _ -> [] + end + end + + defp parse_only_entry(_other), do: [] + + defp unwrap_block({:__block__, _, [inner]}), do: inner + defp unwrap_block(other), do: other + defp put_import(acc, parts, only) do imports = Map.get(acc, {__MODULE__, :imports}, []) Map.put(acc, {__MODULE__, :imports}, [{parts, only} | imports]) @@ -352,8 +380,10 @@ defmodule ExAST.Pattern do end) end - defp import_matches?(:all, _name, _arity), do: true - defp import_matches?(only, name, arity) when is_list(only), do: Keyword.get(only, name) == arity + # A bare `import Mod` gives no membership info, so expanding local calls to + # `Mod.fun` would corrupt unrelated calls — only `only:` imports are safe. + defp import_matches?(:all, _name, _arity), do: false + defp import_matches?(only, name, arity) when is_list(only), do: {name, arity} in only defp import_matches?(_only, _name, _arity), do: false defp import_expandable_call?(form) do diff --git a/test/ex_ast/pattern_test.exs b/test/ex_ast/pattern_test.exs index 699dc07..46d9ed2 100644 --- a/test/ex_ast/pattern_test.exs +++ b/test/ex_ast/pattern_test.exs @@ -287,6 +287,110 @@ defmodule ExAST.PatternTest do assert [_] = ExAST.Patcher.find_all(source, "Ecto.Query.from(_, _)") end + test "import :only only expands the listed functions" do + source = """ + import Ecto.Query, only: [from: 2] + + from(u in User, where: u.id == 1) + where(query, [u], u.id == 1) + """ + + # `from/2` is imported and resolves to its qualified name... + assert [_] = ExAST.Patcher.find_all(source, "Ecto.Query.from(_, _)") + # ...but `where/3` is not listed, so it stays a local call. + assert [] = ExAST.Patcher.find_all(source, "Ecto.Query.where(_, _, _)") + assert [_] = ExAST.Patcher.find_all(source, "where(_, _, _)") + end + + test "import :only expands every listed arity of the same function" do + source = """ + import Ecto.Query, only: [from: 1, from: 2] + + from(u in User) + from(u in User, where: u.id == 1) + """ + + assert [_] = ExAST.Patcher.find_all(source, "Ecto.Query.from(_)") + assert [_] = ExAST.Patcher.find_all(source, "Ecto.Query.from(_, _)") + end + + test "import with options but no :only does not expand calls" do + source = """ + import Ecto.Query, warn: false + + from(u in User, where: u.id == 1) + """ + + assert [] = ExAST.Patcher.find_all(source, "Ecto.Query.from(_, _)") + assert [_] = ExAST.Patcher.find_all(source, "from(_, _)") + end + + test "import only: :functions does not expand calls" do + source = """ + import Ecto.Query, only: :functions + + from(u in User, where: u.id == 1) + """ + + assert [] = ExAST.Patcher.find_all(source, "Ecto.Query.from(_, _)") + assert [_] = ExAST.Patcher.find_all(source, "from(_, _)") + end + + test "malformed :only entries are ignored" do + # `:from` is not a `fun: arity` pair and `where: :two` has a non-integer + # arity, so neither expands. + source = """ + import Ecto.Query, only: [:from, where: :two] + + from(u in User, where: u.id == 1) + where(query, [u], u.id == 1) + """ + + assert [] = ExAST.Patcher.find_all(source, "Ecto.Query.from(_, _)") + assert [] = ExAST.Patcher.find_all(source, "Ecto.Query.where(_, _, _)") + end + + test "a bare import does not break matching of unrelated calls" do + source = """ + defmodule Demo do + import Enum + + def add(a, b, c) do + String.upcase("x") + a + b + c + end + end + """ + + assert [_] = ExAST.Patcher.find_all(source, "def add(_, _, _) do ... end") + assert [_] = ExAST.Patcher.find_all(source, "def _(_, _, _) do ... end") + assert [_] = ExAST.Patcher.find_all(source, "add(_, _, _)") + assert [_] = ExAST.Patcher.find_all(source, "String.upcase(_)") + assert [_] = ExAST.Patcher.find_all(source, "def _ do ... end") + end + + test "a bare import does not break Logger.info matching (bug report repro)" do + source = """ + defmodule Demo do + import Enum + require Logger + def run, do: Logger.info("hi") + end + """ + + assert [_] = ExAST.Patcher.find_all(source, "Logger.info(...)") + end + + test "import :only expands calls on a raw quoted AST" do + ast = + quote do + import Ecto.Query, only: [from: 2] + from(u in User, where: u.id == 1) + end + + assert [_] = ExAST.Patcher.find_all(ast, "Ecto.Query.from(_, _)") + end + test "alias" do assert {:ok, caps} = match!("alias MyApp.Accounts.User", "alias mod") assert Map.has_key?(caps, :mod)