Fix bare imports breaking pattern matching with arguments

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

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)