Merge pull request #47 from benoitc/feature/import-caching

Add lazy module import caching API

Author: benoitc

c_src/py_event_loop.c

@@ -690,6 +690,62 @@ ERL_NIF_TERM nif_process_ready_tasks(ErlNifEnv *env, int argc,
 ERL_NIF_TERM nif_event_loop_set_py_loop(ErlNifEnv *env, int argc,
                                          const ERL_NIF_TERM argv[]);
 
+/* ============================================================================
+ * Module Import Caching
+ * ============================================================================ */
+
+/**
+ * @brief Import and cache a module in the event loop's interpreter
+ *
+ * Pre-imports the module and caches it for faster subsequent calls.
+ * The __main__ module is never cached (returns error).
+ *
+ * NIF: loop_import_module(LoopRef, Module) -> ok | {error, Reason}
+ */
+ERL_NIF_TERM nif_loop_import_module(ErlNifEnv *env, int argc,
+                                     const ERL_NIF_TERM argv[]);
+
+/**
+ * @brief Import a module and cache a specific function
+ *
+ * Pre-imports the module and caches the function reference.
+ * The __main__ module is never cached (returns error).
+ *
+ * NIF: loop_import_function(LoopRef, Module, Func) -> ok | {error, Reason}
+ */
+ERL_NIF_TERM nif_loop_import_function(ErlNifEnv *env, int argc,
+                                       const ERL_NIF_TERM argv[]);
+
+/**
+ * @brief Flush the import cache for an event loop's interpreter
+ *
+ * Clears the module/function cache for all namespaces in this loop.
+ *
+ * NIF: loop_flush_import_cache(LoopRef) -> ok
+ */
+ERL_NIF_TERM nif_loop_flush_import_cache(ErlNifEnv *env, int argc,
+                                          const ERL_NIF_TERM argv[]);
+
+/**
+ * @brief Get import cache statistics for the calling process's namespace
+ *
+ * Returns a map with count of cached entries.
+ *
+ * NIF: loop_import_stats(LoopRef) -> {ok, #{count => N}} | {error, Reason}
+ */
+ERL_NIF_TERM nif_loop_import_stats(ErlNifEnv *env, int argc,
+                                    const ERL_NIF_TERM argv[]);
+
+/**
+ * @brief List all cached imports in the calling process's namespace
+ *
+ * Returns a list of binary strings with cached module and function names.
+ *
+ * NIF: loop_import_list(LoopRef) -> {ok, [binary()]} | {error, Reason}
+ */
+ERL_NIF_TERM nif_loop_import_list(ErlNifEnv *env, int argc,
+                                   const ERL_NIF_TERM argv[]);
+
 /* ============================================================================
  * Internal Helper Functions
  * ============================================================================ */

c_src/py_event_loop.h

@@ -2475,27 +2475,42 @@ static void owngil_execute_call(py_context_t *ctx) {
         return;
     }
 
-    /* Get or import module */
-    PyObject *module = context_get_module(ctx, module_name);
-    if (module == NULL) {
-        ctx->response_term = make_py_error(ctx->shared_env);
-        ctx->response_ok = false;
-        enif_free(module_name);
-        enif_free(func_name_str);
-        return;
-    }
+    PyObject *module = NULL;
+    PyObject *func = NULL;
 
-    /* Get function */
-    PyObject *func = PyObject_GetAttrString(module, func_name_str);
-    enif_free(module_name);
-    enif_free(func_name_str);
+    /* Special handling for __main__ module - check ctx->globals first */
+    if (strcmp(module_name, "__main__") == 0) {
+        func = PyDict_GetItemString(ctx->globals, func_name_str);  /* Borrowed ref */
+        if (func != NULL) {
+            Py_INCREF(func);
+        }
+    }
 
     if (func == NULL) {
-        ctx->response_term = make_py_error(ctx->shared_env);
-        ctx->response_ok = false;
-        return;
+        /* Get or import module */
+        module = context_get_module(ctx, module_name);
+        if (module == NULL) {
+            ctx->response_term = make_py_error(ctx->shared_env);
+            ctx->response_ok = false;
+            enif_free(module_name);
+            enif_free(func_name_str);
+            return;
+        }
+
+        /* Get function */
+        func = PyObject_GetAttrString(module, func_name_str);
+        if (func == NULL) {
+            ctx->response_term = make_py_error(ctx->shared_env);
+            ctx->response_ok = false;
+            enif_free(module_name);
+            enif_free(func_name_str);
+            return;
+        }
     }
 
+    enif_free(module_name);
+    enif_free(func_name_str);
+
     /* Convert args */
     unsigned int args_len;
     if (!enif_get_list_length(ctx->shared_env, args_term, &args_len)) {
@@ -4251,18 +4266,31 @@ static ERL_NIF_TERM nif_context_call(ErlNifEnv *env, int argc, const ERL_NIF_TER
     bool prev_allow_suspension = tl_allow_suspension;
     tl_allow_suspension = true;
 
-    /* Get or import module */
-    PyObject *module = context_get_module(ctx, module_name);
-    if (module == NULL) {
-        result = make_py_error(env);
-        goto cleanup;
+    PyObject *module = NULL;
+    PyObject *func = NULL;
+
+    /* Special handling for __main__ module - check ctx->globals first */
+    if (strcmp(module_name, "__main__") == 0) {
+        func = PyDict_GetItemString(ctx->globals, func_name);  /* Borrowed ref */
+        if (func != NULL) {
+            Py_INCREF(func);
+        }
     }
 
-    /* Get function */
-    PyObject *func = PyObject_GetAttrString(module, func_name);
     if (func == NULL) {
-        result = make_py_error(env);
-        goto cleanup;
+        /* Get or import module */
+        module = context_get_module(ctx, module_name);
+        if (module == NULL) {
+            result = make_py_error(env);
+            goto cleanup;
+        }
+
+        /* Get function */
+        func = PyObject_GetAttrString(module, func_name);
+        if (func == NULL) {
+            result = make_py_error(env);
+            goto cleanup;
+        }
     }
 
     /* Convert args */
@@ -6764,6 +6792,12 @@ static ErlNifFunc nif_funcs[] = {
     /* Per-process namespace NIFs */
     {"event_loop_exec", 2, nif_event_loop_exec, ERL_NIF_DIRTY_JOB_IO_BOUND},
     {"event_loop_eval", 2, nif_event_loop_eval, ERL_NIF_DIRTY_JOB_IO_BOUND},
+    /* Module import caching NIFs */
+    {"loop_import_module", 2, nif_loop_import_module, ERL_NIF_DIRTY_JOB_IO_BOUND},
+    {"loop_import_function", 3, nif_loop_import_function, ERL_NIF_DIRTY_JOB_IO_BOUND},
+    {"loop_flush_import_cache", 1, nif_loop_flush_import_cache, 0},
+    {"loop_import_stats", 1, nif_loop_import_stats, 0},
+    {"loop_import_list", 1, nif_loop_import_list, 0},
     {"add_reader", 3, nif_add_reader, 0},
     {"remove_reader", 2, nif_remove_reader, 0},
     {"add_writer", 3, nif_add_writer, 0},

c_src/py_nif.c

@@ -56,6 +56,12 @@
     stream/4,
     stream_eval/1,
     stream_eval/2,
+    %% Module import caching
+    import/1,
+    import/2,
+    flush_imports/0,
+    import_stats/0,
+    import_list/0,
     version/0,
     memory_stats/0,
     gc/0,
@@ -327,6 +333,93 @@ exec(Ctx, Code) when is_pid(Ctx) ->
     EnvRef = get_local_env(Ctx),
     py_context:exec(Ctx, Code, EnvRef).
 
+%%% ============================================================================
+%%% Module Import Caching
+%%% ============================================================================
+
+%% @doc Import and cache a module in the current interpreter.
+%%
+%% The module is imported in the interpreter handling this process (via affinity).
+%% The `__main__' module is never cached in the interpreter cache.
+%%
+%% This is useful for pre-warming imports before making calls, ensuring the
+%% first call doesn't pay the import penalty.
+%%
+%% Example:
+%% ```
+%% ok = py:import(json),
+%% {ok, Result} = py:call(json, dumps, [Data]).  %% Uses cached module
+%% '''
+%%
+%% @param Module Python module name
+%% @returns ok | {error, Reason}
+-spec import(py_module()) -> ok | {error, term()}.
+import(Module) ->
+    py_event_loop_pool:import(Module).
+
+%% @doc Import and cache a module function in the current interpreter.
+%%
+%% Pre-imports the module and caches the function reference for faster
+%% subsequent calls. The `__main__' module is never cached.
+%%
+%% Example:
+%% ```
+%% ok = py:import(json, dumps),
+%% {ok, Result} = py:call(json, dumps, [Data]).  %% Uses cached function
+%% '''
+%%
+%% @param Module Python module name
+%% @param Func Function name to cache
+%% @returns ok | {error, Reason}
+-spec import(py_module(), py_func()) -> ok | {error, term()}.
+import(Module, Func) ->
+    py_event_loop_pool:import(Module, Func).
+
+%% @doc Flush import caches across all interpreters.
+%%
+%% Clears the module/function cache in all interpreters. Use this after
+%% modifying Python modules on disk to force re-import.
+%%
+%% @returns ok
+-spec flush_imports() -> ok.
+flush_imports() ->
+    py_event_loop_pool:flush_imports().
+
+%% @doc Get import cache statistics for the current interpreter.
+%%
+%% Returns a map with cache metrics for the interpreter handling this process.
+%%
+%% Example:
+%% ```
+%% {ok, #{count => 5}} = py:import_stats().
+%% '''
+%%
+%% @returns {ok, Stats} where Stats is a map with cache metrics
+-spec import_stats() -> {ok, map()} | {error, term()}.
+import_stats() ->
+    py_event_loop_pool:import_stats().
+
+%% @doc List all cached imports in the current interpreter.
+%%
+%% Returns a map of modules to their cached functions.
+%% Module names are binary keys, function lists are the values.
+%% An empty list means only the module is cached (no specific functions).
+%%
+%% Example:
+%% ```
+%% ok = py:import(json),
+%% ok = py:import(json, dumps),
+%% ok = py:import(json, loads),
+%% ok = py:import(math),
+%% {ok, #{<<"json">> => [<<"dumps">>, <<"loads">>],
+%%        <<"math">> => []}} = py:import_list().
+%% '''
+%%
+%% @returns {ok, #{Module => [Func]}} map of modules to functions
+-spec import_list() -> {ok, #{binary() => [binary()]}} | {error, term()}.
+import_list() ->
+    py_event_loop_pool:import_list().
+
 %%% ============================================================================
 %%% Asynchronous API
 %%% ============================================================================

src/py.erl

@@ -35,10 +35,14 @@
 
 -export([
     start_link/2,
+    new/1,
     stop/1,
+    destroy/1,
+    call/4,
     call/5,
     call/6,
     call/7,
+    eval/2,
     eval/3,
     eval/4,
     eval/5,
@@ -121,6 +125,38 @@ stop(Ctx) when is_pid(Ctx) ->
         ok
     end.
 
+%% @doc Create a new context with options map.
+%%
+%% Options:
+%% - `mode' - Context mode (auto | subinterp | worker | owngil), default: auto
+%%
+%% @param Opts Options map
+%% @returns {ok, Pid} | {error, Reason}
+-spec new(map()) -> {ok, context()} | {error, term()}.
+new(Opts) when is_map(Opts) ->
+    Mode = maps:get(mode, Opts, auto),
+    Id = erlang:unique_integer([positive]),
+    start_link(Id, Mode).
+
+%% @doc Alias for stop/1 for API consistency.
+-spec destroy(context()) -> ok.
+destroy(Ctx) ->
+    stop(Ctx).
+
+%% @doc Call a Python function with empty kwargs.
+%%
+%% This is a convenience wrapper for call/5 that defaults Kwargs to #{}.
+%%
+%% @param Ctx Context process
+%% @param Module Python module name
+%% @param Func Function name
+%% @param Args List of arguments
+%% @returns {ok, Result} | {error, Reason}
+-spec call(context(), atom() | binary(), atom() | binary(), list()) ->
+    {ok, term()} | {error, term()}.
+call(Ctx, Module, Func, Args) ->
+    call(Ctx, Module, Func, Args, #{}).
+
 %% @doc Call a Python function.
 %%
 %% @param Ctx Context process
@@ -181,6 +217,18 @@ call(Ctx, Module, Func, Args, Kwargs, Timeout, EnvRef) when is_pid(Ctx), is_refe
         {error, timeout}
     end.
 
+%% @doc Evaluate a Python expression with empty locals.
+%%
+%% This is a convenience wrapper for eval/3 that defaults Locals to #{}.
+%%
+%% @param Ctx Context process
+%% @param Code Python code to evaluate
+%% @returns {ok, Result} | {error, Reason}
+-spec eval(context(), binary() | string()) ->
+    {ok, term()} | {error, term()}.
+eval(Ctx, Code) ->
+    eval(Ctx, Code, #{}).
+
 %% @doc Evaluate a Python expression.
 %%
 %% @param Ctx Context process