Add py:reload/1 to reload Python modules across all workers

Adds broadcast/1 function to py_pool for sending requests to all
workers simultaneously. The reload/1 function uses importlib.reload()
to refresh modules from disk, useful during development.

Author: benoitc

src/py.erl

@@ -88,7 +88,9 @@
     state_incr/1,
     state_incr/2,
     state_decr/1,
-    state_decr/2
+    state_decr/2,
+    %% Module reload
+    reload/1
 ]).
 
 -type py_result() :: {ok, term()} | {error, term()}.
@@ -574,3 +576,45 @@ state_decr(Key) ->
 -spec state_decr(term(), integer()) -> integer().
 state_decr(Key, Amount) ->
     py_state:decr(Key, Amount).
+
+%%% ============================================================================
+%%% Module Reload
+%%% ============================================================================
+
+%% @doc Reload a Python module across all workers.
+%% This uses importlib.reload() to refresh the module from disk.
+%% Useful during development when Python code changes.
+%%
+%% Note: This only affects already-imported modules. If the module
+%% hasn't been imported in a worker yet, the reload is a no-op for that worker.
+%%
+%% Example:
+%% ```
+%% %% After modifying mymodule.py on disk:
+%% ok = py:reload(mymodule).
+%% '''
+%%
+%% Returns ok if reload succeeded in all workers, or {error, Reasons}
+%% if any workers failed.
+-spec reload(py_module()) -> ok | {error, [{worker, term()}]}.
+reload(Module) ->
+    ModuleBin = ensure_binary(Module),
+    %% Build Python code that:
+    %% 1. Checks if module is loaded in sys.modules
+    %% 2. If yes, reloads it with importlib.reload()
+    %% 3. Returns the module name or None if not loaded
+    Code = <<"__import__('importlib').reload(__import__('sys').modules['",
+             ModuleBin/binary,
+             "']) if '", ModuleBin/binary, "' in __import__('sys').modules else None">>,
+    %% Broadcast to all workers
+    Request = {eval, undefined, undefined, Code, #{}},
+    Results = py_pool:broadcast(Request),
+    %% Check if any failed
+    Errors = lists:filtermap(fun
+        ({ok, _}) -> false;
+        ({error, Reason}) -> {true, Reason}
+    end, Results),
+    case Errors of
+        [] -> ok;
+        _ -> {error, [{worker, E} || E <- Errors]}
+    end.

src/py_pool.erl

@@ -24,6 +24,7 @@
 -export([
     start_link/1,
     request/1,
+    broadcast/1,
     get_stats/0
 ]).
 
@@ -55,6 +56,12 @@ start_link(NumWorkers) ->
 request(Request) ->
     gen_server:cast(?MODULE, {request, Request}).
 
+%% @doc Broadcast a request to all workers.
+%% Returns a list of results from each worker.
+-spec broadcast(term()) -> [{ok, term()} | {error, term()}].
+broadcast(Request) ->
+    gen_server:call(?MODULE, {broadcast, Request}, infinity).
+
 %% @doc Get pool statistics.
 -spec get_stats() -> map().
 get_stats() ->
@@ -94,6 +101,12 @@ handle_call(get_stats, _From, State) ->
     },
     {reply, Stats, State};
 
+handle_call({broadcast, Request}, _From, State) ->
+    %% Send request to all workers and collect results
+    Workers = queue:to_list(State#state.workers),
+    Results = broadcast_to_workers(Workers, Request),
+    {reply, Results, State};
+
 handle_call(_Request, _From, State) ->
     {reply, {error, unknown_request}, State}.
 
@@ -150,3 +163,33 @@ extract_ref_caller({call, Ref, Caller, _, _, _, _}) -> {Ref, Caller, call};
 extract_ref_caller({eval, Ref, Caller, _, _}) -> {Ref, Caller, eval};
 extract_ref_caller({exec, Ref, Caller, _}) -> {Ref, Caller, exec};
 extract_ref_caller({stream, Ref, Caller, _, _, _, _}) -> {Ref, Caller, stream}.
+
+%% @private
+%% Send a request to all workers and collect results
+broadcast_to_workers(Workers, RequestTemplate) ->
+    Self = self(),
+    %% Send requests to all workers in parallel
+    Refs = lists:map(fun(Worker) ->
+        Ref = make_ref(),
+        Request = inject_ref_caller(RequestTemplate, Ref, Self),
+        Worker ! {py_request, Request},
+        Ref
+    end, Workers),
+    %% Collect all responses
+    lists:map(fun(Ref) ->
+        receive
+            {py_response, Ref, Result} -> Result;
+            {py_error, Ref, Error} -> {error, Error}
+        after 30000 ->
+            {error, timeout}
+        end
+    end, Refs).
+
+%% @private
+%% Inject a reference and caller into a request template
+inject_ref_caller({exec, _Ref, _Caller, Code}, NewRef, NewCaller) ->
+    {exec, NewRef, NewCaller, Code};
+inject_ref_caller({eval, _Ref, _Caller, Code, Locals}, NewRef, NewCaller) ->
+    {eval, NewRef, NewCaller, Code, Locals};
+inject_ref_caller({eval, _Ref, _Caller, Code, Locals, Timeout}, NewRef, NewCaller) ->
+    {eval, NewRef, NewCaller, Code, Locals, Timeout}.

test/py_SUITE.erl

@@ -46,7 +46,8 @@
     test_semaphore_timeout/1,
     test_semaphore_rate_limiting/1,
     test_overload_protection/1,
-    test_shared_state/1
+    test_shared_state/1,
+    test_reload/1
 ]).
 
 all() ->
@@ -87,7 +88,8 @@ all() ->
         test_semaphore_timeout,
         test_semaphore_rate_limiting,
         test_overload_protection,
-        test_shared_state
+        test_shared_state,
+        test_reload
     ].
 
 init_per_suite(Config) ->
@@ -891,3 +893,25 @@ assert val == 2, f'Expected 2, got {val}'
     ok = py:state_clear(),
 
     ok.
+
+%% Test module reload across all workers
+test_reload(_Config) ->
+    %% First, ensure json module is imported in at least one worker
+    {ok, _} = py:call(json, dumps, [[1, 2, 3]]),
+
+    %% Now reload it - should succeed across all workers
+    ok = py:reload(json),
+
+    %% Verify the module still works after reload
+    {ok, <<"[1, 2, 3]">>} = py:call(json, dumps, [[1, 2, 3]]),
+
+    %% Test reload of a module that might not be loaded (should not error)
+    ok = py:reload(collections),
+
+    %% Test reload with binary module name
+    ok = py:reload(<<"os">>),
+
+    %% Test reload with string module name
+    ok = py:reload("sys"),
+
+    ok.