Apply imports and paths immediately to running interpreters

ensure_imported/1,2 and add_path/1 now apply changes immediately to
all running interpreters (contexts and event loops), not just future ones.

- Add apply_import_to_interpreters/1 and apply_path_to_interpreters/1
- Apply to all contexts via py_context_router:contexts()
- Apply to main event loop via py_event_loop:get_loop()
- Apply to pool event loops via py_event_loop_pool:get_all_loops()
- Add tests verifying immediate application

Author: benoitc

src/py_import.erl

@@ -15,17 +15,17 @@
 %%% @doc Import and path registry for Python interpreters.
 %%%
 %%% This module manages the global import and path registries that are
-%%% applied to all Python interpreters. When new interpreters are created,
-%%% they automatically get all registered imports and paths applied.
+%%% applied to all Python interpreters. Imports and paths are applied
+%%% immediately to all running interpreters and stored for new interpreters.
 %%%
 %%% == Examples ==
 %%%
 %%% ```
-%%% %% Register modules for import in all interpreters
+%%% %% Register modules for import in all interpreters (immediate + future)
 %%% ok = py_import:ensure_imported(json).
 %%% ok = py_import:ensure_imported(math, sqrt).
 %%%
-%%% %% Add paths to sys.path in all interpreters
+%%% %% Add paths to sys.path in all interpreters (immediate + future)
 %%% ok = py_import:add_path("/path/to/my/modules").
 %%%
 %%% %% Check registry contents
@@ -99,48 +99,42 @@ init() ->
 
 %% @doc Register a module for import in all interpreters.
 %%
-%% Adds the module to the global import registry. When new interpreters
-%% are created, they will automatically import all registered modules.
-%% The module will be imported lazily when first used.
+%% Imports the module immediately in all running interpreters and adds it
+%% to the registry for future interpreters.
 %%
 %% The `__main__' module is never cached.
 %%
 %% Example:
 %% ```
 %% ok = py_import:ensure_imported(json),
-%% {ok, Result} = py:call(json, dumps, [Data]).  %% Module imported on first use
+%% {ok, Result} = py:call(json, dumps, [Data]).
 %% '''
 %%
 %% @param Module Python module name
 %% @returns ok | {error, Reason}
 -spec ensure_imported(py_module()) -> ok | {error, term()}.
 ensure_imported(Module) ->
     ModuleBin = ensure_binary(Module),
-    %% Reject __main__
     case ModuleBin of
         <<"__main__">> ->
             {error, main_not_cacheable};
         _ ->
-            %% Add to global registry - module will be imported lazily
-            case ets:info(?IMPORT_REGISTRY) of
-                undefined -> ok;
-                _ -> ets:insert(?IMPORT_REGISTRY, {ModuleBin, all})
-            end,
+            ets:insert(?IMPORT_REGISTRY, {ModuleBin, all}),
+            apply_import_to_interpreters(ModuleBin),
             ok
     end.
 
 %% @doc Register a module/function for import in all interpreters.
 %%
-%% Adds the module/function to the global import registry. When new
-%% interpreters are created, they will automatically import the module.
-%% The module will be imported lazily when first used.
+%% Imports the module immediately in all running interpreters and adds it
+%% to the registry for future interpreters.
 %%
 %% The `__main__' module is never cached.
 %%
 %% Example:
 %% ```
 %% ok = py_import:ensure_imported(json, dumps),
-%% {ok, Result} = py:call(json, dumps, [Data]).  %% Module imported on first use
+%% {ok, Result} = py:call(json, dumps, [Data]).
 %% '''
 %%
 %% @param Module Python module name
@@ -150,16 +144,12 @@ ensure_imported(Module) ->
 ensure_imported(Module, Func) ->
     ModuleBin = ensure_binary(Module),
     FuncBin = ensure_binary(Func),
-    %% Reject __main__
     case ModuleBin of
         <<"__main__">> ->
             {error, main_not_cacheable};
         _ ->
-            %% Add to global registry - module will be imported lazily
-            case ets:info(?IMPORT_REGISTRY) of
-                undefined -> ok;
-                _ -> ets:insert(?IMPORT_REGISTRY, {ModuleBin, FuncBin})
-            end,
+            ets:insert(?IMPORT_REGISTRY, {ModuleBin, FuncBin}),
+            apply_import_to_interpreters(ModuleBin),
             ok
     end.
 
@@ -268,9 +258,8 @@ import_list() ->
 
 %% @doc Add a path to sys.path in all interpreters.
 %%
-%% Adds the path to the global path registry. When new interpreters
-%% are created, they will automatically have this path in sys.path.
-%% The path is inserted at the beginning of sys.path to take precedence.
+%% Adds the path immediately to sys.path in all running interpreters
+%% (contexts and event loops) and stores it for future interpreters.
 %%
 %% Example:
 %% ```
@@ -283,13 +272,9 @@ import_list() ->
 -spec add_path(string() | binary() | atom()) -> ok.
 add_path(Path) ->
     PathBin = ensure_binary(Path),
-    case ets:info(?PATH_REGISTRY) of
-        undefined -> ok;
-        _ ->
-            %% Use monotonic time as key to preserve insertion order
-            Key = erlang:monotonic_time(),
-            ets:insert(?PATH_REGISTRY, {Key, PathBin})
-    end,
+    Key = erlang:monotonic_time(),
+    ets:insert(?PATH_REGISTRY, {Key, PathBin}),
+    apply_path_to_interpreters(PathBin),
     ok.
 
 %% @doc Add multiple paths to sys.path in all interpreters.
@@ -367,3 +352,75 @@ is_path_added(Path) ->
 %% @private
 ensure_binary(S) ->
     py_util:to_binary(S).
+
+%% @private Apply import to all running interpreters (contexts + event loops)
+apply_import_to_interpreters(ModuleBin) ->
+    Imports = [{ModuleBin, all}],
+    %% Apply to all contexts
+    lists:foreach(
+        fun(Ctx) ->
+            try
+                Ref = py_context:get_nif_ref(Ctx),
+                py_nif:interp_apply_imports(Ref, Imports)
+            catch _:_ -> ok
+            end
+        end,
+        get_all_contexts()
+    ),
+    %% Apply to main event loop
+    case py_event_loop:get_loop() of
+        {ok, LoopRef} ->
+            catch py_nif:interp_apply_imports(LoopRef, Imports);
+        _ -> ok
+    end,
+    %% Apply to all pool event loops
+    case py_event_loop_pool:get_all_loops() of
+        {ok, Loops} ->
+            lists:foreach(
+                fun({LoopRef, _WorkerPid}) ->
+                    catch py_nif:interp_apply_imports(LoopRef, Imports)
+                end,
+                Loops
+            );
+        _ -> ok
+    end,
+    ok.
+
+%% @private Apply path to all running interpreters (contexts + event loops)
+apply_path_to_interpreters(PathBin) ->
+    Paths = [PathBin],
+    %% Apply to all contexts
+    lists:foreach(
+        fun(Ctx) ->
+            try
+                Ref = py_context:get_nif_ref(Ctx),
+                py_nif:interp_apply_paths(Ref, Paths)
+            catch _:_ -> ok
+            end
+        end,
+        get_all_contexts()
+    ),
+    %% Apply to main event loop
+    case py_event_loop:get_loop() of
+        {ok, LoopRef} ->
+            catch py_nif:interp_apply_paths(LoopRef, Paths);
+        _ -> ok
+    end,
+    %% Apply to all pool event loops
+    case py_event_loop_pool:get_all_loops() of
+        {ok, Loops} ->
+            lists:foreach(
+                fun({LoopRef, _WorkerPid}) ->
+                    catch py_nif:interp_apply_paths(LoopRef, Paths)
+                end,
+                Loops
+            );
+        _ -> ok
+    end,
+    ok.
+
+%% @private Get all context pids from all pools
+get_all_contexts() ->
+    DefaultCtxs = try py_context_router:contexts() catch _:_ -> [] end,
+    %% Could add other pools here if needed
+    DefaultCtxs.

test/py_import_SUITE.erl

@@ -42,7 +42,10 @@
     registry_import_in_sys_modules_test/1,
     context_import_in_sys_modules_test/1,
     %% Path registry tests
-    add_path_test/1
+    add_path_test/1,
+    %% Immediate application tests
+    import_applies_to_running_interpreter_test/1,
+    path_applies_to_running_interpreter_test/1
 ]).
 
 all() ->
@@ -76,7 +79,10 @@ groups() ->
         registry_import_in_sys_modules_test,
         context_import_in_sys_modules_test,
         %% Path registry tests
-        add_path_test
+        add_path_test,
+        %% Immediate application tests
+        import_applies_to_running_interpreter_test,
+        path_applies_to_running_interpreter_test
     ]}].
 
 init_per_suite(Config) ->
@@ -765,3 +771,77 @@ add_path_test(Config) ->
     ok = py_import:clear_paths(),
 
     ct:pal("add_path successfully registers paths and enables module imports").
+
+%% ============================================================================
+%% Immediate Application Tests
+%% ============================================================================
+
+%% @doc Test that ensure_imported applies immediately to running interpreters
+%%
+%% This verifies that calling ensure_imported on an already-running interpreter
+%% makes the module available without needing to create a new context.
+import_applies_to_running_interpreter_test(_Config) ->
+    %% Clear registry
+    ok = py_import:clear_imports(),
+
+    %% Verify 'zipfile' is NOT in sys.modules yet
+    ok = py:exec(<<"import sys; _zipfile_before = 'zipfile' in sys.modules">>),
+    {ok, BeforeImport} = py:eval(<<"_zipfile_before">>),
+    ?assertEqual(false, BeforeImport),
+
+    %% Now call ensure_imported - should apply immediately
+    ok = py_import:ensure_imported(zipfile),
+
+    %% Verify 'zipfile' IS now in sys.modules (without creating new context)
+    ok = py:exec(<<"import sys; _zipfile_after = 'zipfile' in sys.modules">>),
+    {ok, AfterImport} = py:eval(<<"_zipfile_after">>),
+    ?assertEqual(true, AfterImport),
+
+    %% Verify we can call functions from zipfile
+    {ok, _} = py:call(zipfile, 'is_zipfile', [<<"/nonexistent">>]),
+
+    ct:pal("ensure_imported applies immediately to running interpreter").
+
+%% @doc Test that add_path applies immediately to running interpreters
+%%
+%% This verifies that calling add_path on an already-running interpreter
+%% makes the path available in sys.path without needing to create a new context.
+path_applies_to_running_interpreter_test(Config) ->
+    %% Clear paths
+    ok = py_import:clear_paths(),
+
+    %% Create test module in priv_dir
+    PrivDir = ?config(priv_dir, Config),
+    ModuleDir = filename:join(PrivDir, "immediate_path_test"),
+    ok = filelib:ensure_dir(filename:join(ModuleDir, "dummy")),
+
+    %% Write a simple Python module
+    ModulePath = filename:join(ModuleDir, "immediate_test_mod.py"),
+    ModuleContent = <<"IMMEDIATE_TEST_VALUE = 42\n">>,
+    ok = file:write_file(ModulePath, ModuleContent),
+
+    ModuleDirBin = list_to_binary(ModuleDir),
+
+    %% Verify path is NOT in sys.path yet
+    CheckCode = <<"import sys; _path_before = '", ModuleDirBin/binary, "' in sys.path">>,
+    ok = py:exec(CheckCode),
+    {ok, BeforePath} = py:eval(<<"_path_before">>),
+    ?assertEqual(false, BeforePath),
+
+    %% Now call add_path - should apply immediately
+    ok = py_import:add_path(ModuleDir),
+
+    %% Verify path IS now in sys.path (without creating new context)
+    CheckAfterCode = <<"import sys; _path_after = '", ModuleDirBin/binary, "' in sys.path">>,
+    ok = py:exec(CheckAfterCode),
+    {ok, AfterPath} = py:eval(<<"_path_after">>),
+    ?assertEqual(true, AfterPath),
+
+    %% Verify we can import and use the module
+    {ok, Value} = py:eval(<<"__import__('immediate_test_mod').IMMEDIATE_TEST_VALUE">>),
+    ?assertEqual(42, Value),
+
+    %% Clean up
+    ok = py_import:clear_paths(),
+
+    ct:pal("add_path applies immediately to running interpreter").