Add local_for_callback option to Macro.Env.expand_import (#14620)

Author: lukaszsamson

lib/elixir/lib/macro/env.ex

@@ -96,7 +96,8 @@ defmodule Macro.Env do
         ]
 
   @type expand_import_opts :: [
-          allow_locals: boolean(),
+          allow_locals:
+            boolean() | (Macro.metadata(), atom(), arity(), t() -> function() | false),
           check_deprecations: boolean(),
           trace: boolean()
         ]
@@ -555,8 +556,15 @@ defmodule Macro.Env do
 
   ## Options
 
-    * `:allow_locals` - when set to `false`, it does not attempt to capture
-      local macros defined in the current module in `env`
+    * `:allow_locals` - controls how local macros are resolved.
+      Defaults to `true`.
+      
+      - When `false`, does not attempt to capture local macros defined in the
+        current module in `env`
+      - When `true`, uses a default resolver that looks for public macros in
+        the current module
+      - When a function, uses the function as a custom local resolver. The function
+        must have the signature: `(meta, name, arity, env) -> function() | false`
 
     * `:check_deprecations` - when set to `false`, does not check for deprecations
       when expanding macros
@@ -580,10 +588,16 @@ defmodule Macro.Env do
         trace = Keyword.get(opts, :trace, true)
         module = env.module
 
+        # When allow_locals is a callback, we don't need to pass module macros as extra
+        # because the callback will handle local macro resolution
         extra =
-          case allow_locals and function_exported?(module, :__info__, 1) do
-            true -> [{module, module.__info__(:macros)}]
-            false -> []
+          if is_function(allow_locals, 4) do
+            []
+          else
+            case allow_locals and function_exported?(module, :__info__, 1) do
+              true -> [{module, module.__info__(:macros)}]
+              false -> []
+            end
           end
 
         case :elixir_dispatch.expand_import(meta, name, arity, env, extra, allow_locals, trace) do

lib/elixir/src/elixir_dispatch.erl

@@ -172,7 +172,13 @@ expand_import(Meta, Name, Arity, E, Extra, AllowLocals, Trace) ->
       do_expand_import(Dispatch, Meta, Name, Arity, Module, E, Trace);
 
     _ ->
-      Local = AllowLocals andalso elixir_def:local_for(Meta, Name, Arity, [defmacro, defmacrop], E),
+      Local = case AllowLocals of
+        false -> false;
+        true  -> elixir_def:local_for(Meta, Name, Arity, [defmacro, defmacrop], E);
+        Fun when is_function(Fun, 4) ->
+          %% If we have a custom local resolver, use it.
+          Fun(Meta, Name, Arity, E)
+      end,
 
       case Dispatch of
         %% There is a local and an import. This is a conflict unless