diff --git a/c_src/py_nif.c b/c_src/py_nif.c index 2feee6f..839d34f 100644 --- a/c_src/py_nif.c +++ b/c_src/py_nif.c @@ -4707,6 +4707,249 @@ static ERL_NIF_TERM nif_create_local_env(ErlNifEnv *env, int argc, const ERL_NIF return enif_make_tuple2(env, ATOM_OK, ref); } +/** + * @brief Create a new Python environment with initialization code + * + * nif_new_env_with_code(Code) -> {ok, EnvRef} | {error, Reason} + * + * Creates a new Python environment with globals/locals dicts and executes + * the provided initialization code. The environment can be used independently + * of any context - it runs in the main interpreter. + * + * This is useful for creating named environments that can be shared + * across processes and set as the current environment for py:eval/exec. + */ +static ERL_NIF_TERM nif_new_env_with_code(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]) { + (void)argc; + + if (!runtime_is_running()) { + return make_error(env, "python_not_running"); + } + + ErlNifBinary code_bin; + if (!enif_inspect_binary(env, argv[0], &code_bin)) { + return make_error(env, "invalid_code"); + } + + /* Allocate environment resource */ + py_env_resource_t *res = enif_alloc_resource(PY_ENV_RESOURCE_TYPE, + sizeof(py_env_resource_t)); + if (res == NULL) { + return make_error(env, "alloc_failed"); + } + + res->globals = NULL; + res->locals = NULL; + res->interp_id = 0; + res->pool_slot = -1; + + /* Acquire GIL for main interpreter */ + PyGILState_STATE gstate = PyGILState_Ensure(); + + /* Store interpreter info for destructor */ + PyInterpreterState *interp = PyInterpreterState_Get(); + if (interp != NULL) { + res->interp_id = PyInterpreterState_GetID(interp); + } + + /* Create globals dict with builtins */ + res->globals = PyDict_New(); + if (res->globals == NULL) { + PyGILState_Release(gstate); + enif_release_resource(res); + return make_error(env, "globals_failed"); + } + + /* Add __builtins__ */ + PyObject *builtins = PyEval_GetBuiltins(); + if (builtins != NULL) { + PyDict_SetItemString(res->globals, "__builtins__", builtins); + } + + /* Add __name__ = '__main__' */ + PyObject *main_name = PyUnicode_FromString("__main__"); + if (main_name != NULL) { + PyDict_SetItemString(res->globals, "__name__", main_name); + Py_DECREF(main_name); + } + + /* Add erlang module */ + PyObject *erlang = PyImport_ImportModule("erlang"); + if (erlang != NULL) { + PyDict_SetItemString(res->globals, "erlang", erlang); + Py_DECREF(erlang); + } + + /* Use the same dict for locals (module-level execution) */ + res->locals = res->globals; + Py_INCREF(res->locals); + + /* Execute initialization code */ + char *code = enif_alloc(code_bin.size + 1); + if (code == NULL) { + Py_DECREF(res->globals); + Py_DECREF(res->locals); + res->globals = NULL; + res->locals = NULL; + PyGILState_Release(gstate); + enif_release_resource(res); + return make_error(env, "alloc_failed"); + } + memcpy(code, code_bin.data, code_bin.size); + code[code_bin.size] = '\0'; + + PyObject *py_result = PyRun_String(code, Py_file_input, res->globals, res->globals); + enif_free(code); + + if (py_result == NULL) { + ERL_NIF_TERM error = make_py_error(env); + Py_DECREF(res->globals); + Py_DECREF(res->locals); + res->globals = NULL; + res->locals = NULL; + PyGILState_Release(gstate); + enif_release_resource(res); + return error; + } + Py_DECREF(py_result); + + PyGILState_Release(gstate); + + ERL_NIF_TERM ref = enif_make_resource(env, res); + enif_release_resource(res); /* Ref now owns it */ + + return enif_make_tuple2(env, ATOM_OK, ref); +} + +/** + * @brief Evaluate a Python expression using an environment + * + * nif_env_eval(EnvRef, Code) -> {ok, Result} | {error, Reason} + * + * Evaluates a Python expression using the provided environment's globals. + * This allows evaluation against a named environment without needing a context. + */ +static ERL_NIF_TERM nif_env_eval(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]) { + (void)argc; + + if (!runtime_is_running()) { + return make_error(env, "python_not_running"); + } + + py_env_resource_t *penv; + if (!enif_get_resource(env, argv[0], PY_ENV_RESOURCE_TYPE, (void **)&penv)) { + return make_error(env, "invalid_env"); + } + + ErlNifBinary code_bin; + if (!enif_inspect_binary(env, argv[1], &code_bin)) { + return make_error(env, "invalid_code"); + } + + if (penv->globals == NULL) { + return make_error(env, "env_not_initialized"); + } + + /* Acquire GIL */ + PyGILState_STATE gstate = PyGILState_Ensure(); + + /* Verify interpreter ownership */ + PyInterpreterState *current_interp = PyInterpreterState_Get(); + if (current_interp != NULL && penv->interp_id != PyInterpreterState_GetID(current_interp)) { + PyGILState_Release(gstate); + return make_error(env, "wrong_interpreter"); + } + + /* Copy code to null-terminated string */ + char *code = enif_alloc(code_bin.size + 1); + if (code == NULL) { + PyGILState_Release(gstate); + return make_error(env, "alloc_failed"); + } + memcpy(code, code_bin.data, code_bin.size); + code[code_bin.size] = '\0'; + + /* Evaluate expression */ + PyObject *py_result = PyRun_String(code, Py_eval_input, penv->globals, penv->globals); + enif_free(code); + + ERL_NIF_TERM result; + if (py_result == NULL) { + result = make_py_error(env); + } else { + ERL_NIF_TERM term_result = py_to_term(env, py_result); + Py_DECREF(py_result); + result = enif_make_tuple2(env, ATOM_OK, term_result); + } + + PyGILState_Release(gstate); + return result; +} + +/** + * @brief Execute Python statements using an environment + * + * nif_env_exec(EnvRef, Code) -> ok | {error, Reason} + * + * Executes Python statements using the provided environment's globals. + * This allows execution against a named environment without needing a context. + */ +static ERL_NIF_TERM nif_env_exec(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]) { + (void)argc; + + if (!runtime_is_running()) { + return make_error(env, "python_not_running"); + } + + py_env_resource_t *penv; + if (!enif_get_resource(env, argv[0], PY_ENV_RESOURCE_TYPE, (void **)&penv)) { + return make_error(env, "invalid_env"); + } + + ErlNifBinary code_bin; + if (!enif_inspect_binary(env, argv[1], &code_bin)) { + return make_error(env, "invalid_code"); + } + + if (penv->globals == NULL) { + return make_error(env, "env_not_initialized"); + } + + /* Acquire GIL */ + PyGILState_STATE gstate = PyGILState_Ensure(); + + /* Verify interpreter ownership */ + PyInterpreterState *current_interp = PyInterpreterState_Get(); + if (current_interp != NULL && penv->interp_id != PyInterpreterState_GetID(current_interp)) { + PyGILState_Release(gstate); + return make_error(env, "wrong_interpreter"); + } + + /* Copy code to null-terminated string */ + char *code = enif_alloc(code_bin.size + 1); + if (code == NULL) { + PyGILState_Release(gstate); + return make_error(env, "alloc_failed"); + } + memcpy(code, code_bin.data, code_bin.size); + code[code_bin.size] = '\0'; + + /* Execute statements */ + PyObject *py_result = PyRun_String(code, Py_file_input, penv->globals, penv->globals); + enif_free(code); + + ERL_NIF_TERM result; + if (py_result == NULL) { + result = make_py_error(env); + } else { + Py_DECREF(py_result); + result = ATOM_OK; + } + + PyGILState_Release(gstate); + return result; +} + /** * @brief Execute Python statements using a process-local environment * @@ -6836,6 +7079,9 @@ static ErlNifFunc nif_funcs[] = { {"context_eval", 4, nif_context_eval_with_env, ERL_NIF_DIRTY_JOB_CPU_BOUND}, {"context_call", 6, nif_context_call_with_env, ERL_NIF_DIRTY_JOB_CPU_BOUND}, {"create_local_env", 1, nif_create_local_env, 0}, + {"new_env_with_code", 1, nif_new_env_with_code, ERL_NIF_DIRTY_JOB_CPU_BOUND}, + {"env_eval", 2, nif_env_eval, ERL_NIF_DIRTY_JOB_CPU_BOUND}, + {"env_exec", 2, nif_env_exec, ERL_NIF_DIRTY_JOB_CPU_BOUND}, {"context_call_method", 4, nif_context_call_method, ERL_NIF_DIRTY_JOB_CPU_BOUND}, {"context_to_term", 1, nif_context_to_term, 0}, {"context_interp_id", 1, nif_context_interp_id, 0}, diff --git a/src/py.erl b/src/py.erl index b78b2b8..2fce673 100644 --- a/src/py.erl +++ b/src/py.erl @@ -140,7 +140,15 @@ dup_fd/1, %% Pool registration API register_pool/2, - unregister_pool/1 + unregister_pool/1, + %% Named environment API + new_env/1, + new_env/2, + set_env/1, + get_env/0, + get_env/1, + list_envs/0, + destroy_env/1 ]). -type py_result() :: {ok, term()} | {error, term()}. @@ -158,6 +166,12 @@ %% Process dictionary key for local Python environment -define(LOCAL_ENV_KEY, py_local_env). +%% Process dictionary key for current named environment +-define(CURRENT_ENV_KEY, py_current_env). + +%% Persistent term key for named environments registry +-define(PT_NAMED_ENVS, {py, named_envs}). + %% @doc Get or create a process-local Python environment for a context. %% %% Each Erlang process can have Python environments per interpreter. @@ -252,13 +266,17 @@ call(Module, Func, Args, Kwargs, Timeout) -> %% @private %% Call using a named pool with semaphore protection -%% Uses the process-local environment from the calling process +%% Uses named env if set, otherwise the process-local environment do_pool_call(Pool, Module, Func, Args, Kwargs, Timeout) -> case py_semaphore:acquire(Timeout) of ok -> try Ctx = py_context_router:get_context(Pool), - EnvRef = get_local_env(Ctx), + %% Check named env first, then fall back to process-local env + EnvRef = case get_env() of + undefined -> get_local_env(Ctx); + NamedEnv -> NamedEnv + end, py_context:call(Ctx, Module, Func, Args, Kwargs, Timeout, EnvRef) after py_semaphore:release() @@ -269,7 +287,8 @@ do_pool_call(Pool, Module, Func, Args, Kwargs, Timeout) -> %% @doc Evaluate a Python expression and return the result. %% -%% In worker mode, evaluation uses the process-local Python environment. +%% If a named environment has been set via set_env/1, uses that environment. +%% Otherwise, uses the process-local Python environment via the context system. %% Variables defined via exec are visible in eval within the same process. -spec eval(string() | binary()) -> py_result(). eval(Code) -> @@ -299,24 +318,40 @@ eval(Ctx, Code, Locals) when is_pid(Ctx), is_map(Locals) -> EnvRef = get_local_env(Ctx), py_context:eval(Ctx, Code, Locals, infinity, EnvRef); eval(Code, Locals, Timeout) -> - %% Always route through context process - it handles callbacks inline using - %% suspension-based approach (no separate callback handler, no blocking) - Ctx = py_context_router:get_context(), - EnvRef = get_local_env(Ctx), - py_context:eval(Ctx, Code, Locals, Timeout, EnvRef). + %% Check if process has a named environment set + case get_env() of + undefined -> + %% Fall back to context-based execution + Ctx = py_context_router:get_context(), + EnvRef = get_local_env(Ctx), + py_context:eval(Ctx, Code, Locals, Timeout, EnvRef); + EnvRef -> + %% Use the named environment directly + %% Note: Locals and Timeout are ignored for named envs - use exec to modify + CodeBin = ensure_binary(Code), + py_nif:env_eval(EnvRef, CodeBin) + end. %% @doc Execute Python statements (no return value expected). %% -%% In worker mode, the code runs in a process-local Python environment. +%% If a named environment has been set via set_env/1, uses that environment. +%% Otherwise, uses the process-local Python environment via the context system. %% Variables defined via exec persist within the calling Erlang process. %% In subinterpreter mode, each context has its own isolated namespace. -spec exec(string() | binary()) -> ok | {error, term()}. exec(Code) -> - %% Always route through context process - it handles callbacks inline using - %% suspension-based approach (no separate callback handler, no blocking) - Ctx = py_context_router:get_context(), - EnvRef = get_local_env(Ctx), - py_context:exec(Ctx, Code, EnvRef). + %% Check if process has a named environment set + case get_env() of + undefined -> + %% Fall back to context-based execution + Ctx = py_context_router:get_context(), + EnvRef = get_local_env(Ctx), + py_context:exec(Ctx, Code, EnvRef); + EnvRef -> + %% Use the named environment directly + CodeBin = ensure_binary(Code), + py_nif:env_exec(EnvRef, CodeBin) + end. %% @doc Execute Python statements using a specific context. %% @@ -567,9 +602,16 @@ async_call(Module, Func, Args) -> %% @doc Call a Python async function with keyword arguments. -spec async_call(py_module(), py_func(), py_args(), py_kwargs()) -> py_ref(). async_call(Module, Func, Args, Kwargs) -> - Ref = make_ref(), - py_async_pool:request({async_call, Ref, self(), Module, Func, Args, Kwargs}), - Ref. + %% If named env is set, use create_task which has env support + case get_env() of + undefined -> + Ref = make_ref(), + py_async_pool:request({async_call, Ref, self(), Module, Func, Args, Kwargs}), + Ref; + _EnvRef -> + %% Use event loop pool's create_task which supports named envs + py_event_loop_pool:create_task(Module, Func, Args, Kwargs) + end. %% @doc Wait for an async call to complete. -spec async_await(py_ref()) -> py_result(). @@ -1491,3 +1533,118 @@ unregister_pool(Module) when is_atom(Module) -> unregister_pool({Module, Func}) when is_atom(Module), is_atom(Func) -> py_context_router:unregister_pool(Module, Func). +%%% ============================================================================ +%%% Named Environment API +%%% +%%% Named environments allow creating reusable Python environments with +%%% pre-initialized modules, functions, and variables. These environments +%%% can be shared across processes and set as the default for py:eval/exec. +%%% ============================================================================ + +%% @doc Create a new environment with initialization code. +%% +%% Creates a Python environment and executes the provided code to initialize +%% it with imports, functions, and variables. The environment is isolated +%% and runs in the main interpreter. +%% +%% Example: +%% ``` +%% {ok, Env} = py:new_env(<<"import numpy as np">>). +%% ''' +-spec new_env(Code :: binary() | iolist()) -> {ok, reference()} | {error, term()}. +new_env(Code) -> + new_env(Code, #{}). + +%% @doc Create a new environment with initialization code and options. +%% +%% Options: +%% - `name' - Register the environment with a name for lookup +%% +%% Example: +%% ``` +%% {ok, Env} = py:new_env(<<"import pandas as pd">>, #{name => data_env}). +%% ''' +-spec new_env(Code :: binary() | iolist(), Opts :: map()) -> {ok, reference()} | {error, term()}. +new_env(Code, Opts) -> + CodeBin = iolist_to_binary(Code), + case py_nif:new_env_with_code(CodeBin) of + {ok, EnvRef} -> + case maps:get(name, Opts, undefined) of + undefined -> + {ok, EnvRef}; + Name when is_atom(Name) -> + register_env(Name, EnvRef), + {ok, EnvRef} + end; + Error -> Error + end. + +%% @doc Set the default environment for this process. +%% +%% After setting, py:eval/exec will use this environment if no +%% explicit environment is provided. +%% +%% Example: +%% ``` +%% ok = py:set_env(data_env). %% by name +%% ok = py:set_env(EnvRef). %% by reference +%% ''' +-spec set_env(Env :: reference() | atom()) -> ok | {error, term()}. +set_env(Name) when is_atom(Name) -> + case get_env(Name) of + {ok, EnvRef} -> + put(?CURRENT_ENV_KEY, EnvRef), + ok; + Error -> Error + end; +set_env(EnvRef) when is_reference(EnvRef) -> + put(?CURRENT_ENV_KEY, EnvRef), + ok. + +%% @doc Get the current environment for this process. +%% +%% Returns undefined if no environment has been set. +-spec get_env() -> reference() | undefined. +get_env() -> + get(?CURRENT_ENV_KEY). + +%% @doc Lookup a named environment. +%% +%% Returns {ok, EnvRef} if found, {error, not_found} otherwise. +-spec get_env(Name :: atom()) -> {ok, reference()} | {error, not_found}. +get_env(Name) when is_atom(Name) -> + Envs = persistent_term:get(?PT_NAMED_ENVS, #{}), + case maps:get(Name, Envs, undefined) of + undefined -> {error, not_found}; + EnvRef -> {ok, EnvRef} + end. + +%% @doc List all named environments. +%% +%% Returns a list of {Name, EnvRef} tuples. +-spec list_envs() -> [{atom(), reference()}]. +list_envs() -> + maps:to_list(persistent_term:get(?PT_NAMED_ENVS, #{})). + +%% @doc Destroy an environment. +%% +%% Removes the environment from the named registry if it was registered. +%% The environment's Python resources will be freed when all references +%% to it are garbage collected. +-spec destroy_env(reference()) -> ok. +destroy_env(EnvRef) -> + %% Remove from registry if named + Envs = persistent_term:get(?PT_NAMED_ENVS, #{}), + NewEnvs = maps:filter(fun(_K, V) -> V =/= EnvRef end, Envs), + case NewEnvs =:= Envs of + true -> ok; + false -> persistent_term:put(?PT_NAMED_ENVS, NewEnvs) + end, + ok. + +%% @private Register a named environment +-spec register_env(atom(), reference()) -> ok. +register_env(Name, EnvRef) -> + Envs = persistent_term:get(?PT_NAMED_ENVS, #{}), + persistent_term:put(?PT_NAMED_ENVS, Envs#{Name => EnvRef}), + ok. diff --git a/src/py_event_loop.erl b/src/py_event_loop.erl index 65093f9..397237a 100644 --- a/src/py_event_loop.erl +++ b/src/py_event_loop.erl @@ -167,8 +167,8 @@ create_task(Module, Func, Args, Kwargs) -> Caller = self(), ModuleBin = py_util:to_binary(Module), FuncBin = py_util:to_binary(Func), - %% Check if there's a process-local env (from py:exec) and use it - ok = case get_process_env() of + %% Check named env first, then process-local env + ok = case get_effective_env() of undefined -> py_nif:submit_task(LoopRef, Caller, Ref, ModuleBin, FuncBin, Args, Kwargs); EnvRef -> @@ -193,6 +193,15 @@ get_process_env() -> EnvRef end. +%% @doc Get the effective environment for task submission. +%% Checks named env (py:get_env()) first, then falls back to process-local env. +-spec get_effective_env() -> reference() | undefined. +get_effective_env() -> + case py:get_env() of + undefined -> get_process_env(); + EnvRef -> EnvRef + end. + %% @doc Wait for an async task result. %% %% Blocks until the result is received or timeout is reached. @@ -232,7 +241,8 @@ spawn_task(Module, Func, Args, Kwargs) -> {ok, LoopRef} = get_loop(), Ref = make_ref(), %% Get env from caller's process BEFORE spawning receiver - CallerEnv = get_process_env(), + %% Check named env first, then process-local env + CallerEnv = get_effective_env(), %% Spawn a process that will receive and discard the result Receiver = erlang:spawn(fun() -> receive diff --git a/src/py_event_loop_pool.erl b/src/py_event_loop_pool.erl index b2f12f0..406acdb 100644 --- a/src/py_event_loop_pool.erl +++ b/src/py_event_loop_pool.erl @@ -43,7 +43,9 @@ await/1, await/2, %% Per-process namespace API exec/1, exec/2, - eval/1, eval/2 + eval/1, eval/2, + %% Pool-wide operations + exec_all/1 ]). %% Legacy API @@ -176,7 +178,8 @@ create_task_on_loop(LoopRef, Module, Func, Args, Kwargs) -> Caller = self(), ModuleBin = py_util:to_binary(Module), FuncBin = py_util:to_binary(Func), - ok = case py_event_loop:get_process_env() of + %% Check named env first, then fall back to process env + ok = case get_effective_env() of undefined -> py_nif:submit_task(LoopRef, Caller, Ref, ModuleBin, FuncBin, Args, Kwargs); EnvRef -> @@ -241,7 +244,8 @@ spawn_task_regular(Module, Func, Args) -> case get_loop() of {ok, LoopRef} -> Ref = make_ref(), - CallerEnv = py_event_loop:get_process_env(), + %% Check named env first, then fall back to process env + CallerEnv = get_effective_env(), Receiver = erlang:spawn(fun() -> receive {async_result, _, _} -> ok @@ -296,12 +300,20 @@ await(Ref, Timeout) -> %% -spec exec(Code :: binary() | iolist()) -> ok | {error, term()}. exec(Code) -> - case get_loop() of - {ok, LoopRef} -> - exec(LoopRef, Code); - {error, not_available} -> - %% Fallback to default event loop - py_event_loop:exec(Code) + %% Check if process has a named environment set + case py:get_env() of + undefined -> + %% Original behavior - per-process namespace on event loop + case get_loop() of + {ok, LoopRef} -> + exec(LoopRef, Code); + {error, not_available} -> + %% Fallback to default event loop + py_event_loop:exec(Code) + end; + _EnvRef -> + %% Use current named env via py module + py:exec(Code) end. -spec exec(LoopRef :: reference(), Code :: binary() | iolist()) -> ok | {error, term()}. @@ -320,18 +332,58 @@ exec(LoopRef, Code) -> %% -spec eval(Expr :: binary() | iolist()) -> {ok, term()} | {error, term()}. eval(Expr) -> - case get_loop() of - {ok, LoopRef} -> - eval(LoopRef, Expr); - {error, not_available} -> - %% Fallback to default event loop - py_event_loop:eval(Expr) + %% Check if process has a named environment set + case py:get_env() of + undefined -> + %% Original behavior - per-process namespace on event loop + case get_loop() of + {ok, LoopRef} -> + eval(LoopRef, Expr); + {error, not_available} -> + %% Fallback to default event loop + py_event_loop:eval(Expr) + end; + _EnvRef -> + %% Use current named env via py module + py:eval(Expr) end. -spec eval(LoopRef :: reference(), Expr :: binary() | iolist()) -> {ok, term()} | {error, term()}. eval(LoopRef, Expr) -> py_nif:event_loop_eval(LoopRef, Expr). +%% @doc Execute Python code on ALL workers in the pool. +%% +%% Useful for initializing shared state, importing modules, or defining +%% functions that should be available across all workers. +%% +%% Returns ok if all workers succeed, or {error, Errors} with a list of +%% {Index, Error} tuples for any workers that failed. +%% +%% Example: +%% 
+%% ok = py_event_loop_pool:exec_all(<<"
+%%     import numpy as np
+%%     SHARED_CONFIG = {'version': 1}
+%% ">>)
+%%
+-spec exec_all(Code :: binary() | iolist()) -> ok | {error, [{pos_integer(), term()}]}. +exec_all(Code) -> + case pool_size() of + 0 -> {error, not_available}; + N -> + Loops = persistent_term:get(?PT_LOOPS), + Results = [begin + {LoopRef, _WorkerPid} = element(Idx, Loops), + {Idx, py_nif:event_loop_exec(LoopRef, Code)} + end || Idx <- lists:seq(1, N)], + Errors = [{Idx, Err} || {Idx, {error, Err}} <- Results], + case Errors of + [] -> ok; + _ -> {error, Errors} + end + end. + %%% ============================================================================ %%% Legacy API %%% ============================================================================ @@ -536,6 +588,18 @@ get_loop_by_index(Idx) -> Loops = persistent_term:get(?PT_LOOPS), element(Idx, Loops). +%% @private Get the effective environment for task submission. +%% Checks named env first (py:get_env()), then falls back to process env. +-spec get_effective_env() -> reference() | undefined. +get_effective_env() -> + case py:get_env() of + undefined -> + %% Fall back to event loop process env + py_event_loop:get_process_env(); + EnvRef -> + EnvRef + end. + %% @private Check if OWN_GIL mode is enabled -spec is_owngil_enabled() -> boolean(). is_owngil_enabled() -> diff --git a/src/py_nif.erl b/src/py_nif.erl index da7a02f..014648e 100644 --- a/src/py_nif.erl +++ b/src/py_nif.erl @@ -181,6 +181,10 @@ context_exec/3, context_call_method/4, create_local_env/1, + %% Named environment API + new_env_with_code/1, + env_eval/2, + env_exec/2, context_to_term/1, context_interp_id/1, context_set_callback_handler/2, @@ -1443,6 +1447,40 @@ context_call_method(_ContextRef, _ObjRef, _Method, _Args) -> create_local_env(_CtxRef) -> ?NIF_STUB. +%%% ============================================================================ +%%% Named Environment API +%%% ============================================================================ + +%% @doc Create a new Python environment with initialization code. +%% +%% Creates an isolated Python environment with globals/locals dicts and +%% executes the provided initialization code. The environment runs in the +%% main interpreter and can be shared across processes. +%% +%% @param Code Python code to execute during initialization +%% @returns {ok, EnvRef} | {error, Reason} +-spec new_env_with_code(binary()) -> {ok, reference()} | {error, term()}. +new_env_with_code(_Code) -> + ?NIF_STUB. + +%% @doc Evaluate a Python expression using a named environment. +%% +%% @param EnvRef Environment reference from new_env_with_code/1 +%% @param Code Python expression to evaluate +%% @returns {ok, Result} | {error, Reason} +-spec env_eval(reference(), binary()) -> {ok, term()} | {error, term()}. +env_eval(_EnvRef, _Code) -> + ?NIF_STUB. + +%% @doc Execute Python statements using a named environment. +%% +%% @param EnvRef Environment reference from new_env_with_code/1 +%% @param Code Python statements to execute +%% @returns ok | {error, Reason} +-spec env_exec(reference(), binary()) -> ok | {error, term()}. +env_exec(_EnvRef, _Code) -> + ?NIF_STUB. + %% @doc Convert a Python object reference to an Erlang term. %% %% The reference carries the interpreter ID, allowing automatic routing diff --git a/test/py_event_loop_pool_SUITE.erl b/test/py_event_loop_pool_SUITE.erl index 6ae4737..7e79917 100644 --- a/test/py_event_loop_pool_SUITE.erl +++ b/test/py_event_loop_pool_SUITE.erl @@ -33,7 +33,14 @@ test_exec_basic/1, test_eval_basic/1, test_exec_eval_namespace/1, - test_exec_define_function/1 + test_exec_define_function/1, + test_exec_all/1, + + %% Named environment tests + test_named_env/1, + test_multiple_envs/1, + test_env_in_processes/1, + test_list_envs/1 ]). all() -> @@ -50,7 +57,12 @@ all() -> test_exec_basic, test_eval_basic, test_exec_eval_namespace, - test_exec_define_function + test_exec_define_function, + test_exec_all, + test_named_env, + test_multiple_envs, + test_env_in_processes, + test_list_envs ]. init_per_suite(Config) -> @@ -205,3 +217,100 @@ def pool_test_multiply(a, b): Ref = py_event_loop_pool:create_task('__main__', pool_test_multiply, [7, 6]), {ok, 42} = py_event_loop_pool:await(Ref, 5000), ok. + +test_exec_all(_Config) -> + %% Execute on all workers - initializes this process's namespace on each loop + ok = py_event_loop_pool:exec_all(<<"EXEC_ALL_VAR = 42">>), + + %% Verify from the same process - should see the variable on its assigned loop + {ok, 42} = py_event_loop_pool:eval(<<"EXEC_ALL_VAR">>), + + %% Define a function across all workers + ok = py_event_loop_pool:exec_all(<<" +def exec_all_multiply(a, b): + return a * b +">>), + + %% Should be able to call the function + Ref = py_event_loop_pool:create_task('__main__', exec_all_multiply, [7, 8]), + {ok, 56} = py_event_loop_pool:await(Ref, 5000), + + %% Test error case - invalid Python should fail on all loops + {error, Errors} = py_event_loop_pool:exec_all(<<"this is not valid python syntax $$$">>), + true = is_list(Errors), + true = length(Errors) > 0, + + ok. + +%% ============================================================================ +%% Named Environment Tests +%% ============================================================================ + +test_named_env(_Config) -> + %% Create named environment + {ok, _Ref} = py:new_env(<<" +import math +ENV_VAR = 'test_env' +def env_func(x): return x * 2 +">>, #{name => test_env}), + + %% Set as current env + ok = py:set_env(test_env), + + %% Functions available + {ok, <<"test_env">>} = py:eval(<<"ENV_VAR">>), + {ok, 42} = py:eval(<<"env_func(21)">>), + {ok, Pi} = py:eval(<<"math.pi">>), + true = is_float(Pi), + + %% Clean up - clear current env from process dictionary + erase(py_current_env), + + ok. + +test_multiple_envs(_Config) -> + %% Create two different environments + {ok, _} = py:new_env(<<"X = 'env1'">>, #{name => env1}), + {ok, _} = py:new_env(<<"X = 'env2'">>, #{name => env2}), + + %% Switch between them + ok = py:set_env(env1), + {ok, <<"env1">>} = py:eval(<<"X">>), + + ok = py:set_env(env2), + {ok, <<"env2">>} = py:eval(<<"X">>), + + %% Clean up + erase(py_current_env), + + ok. + +test_env_in_processes(_Config) -> + {ok, _} = py:new_env(<<"SHARED = 42">>, #{name => shared_env}), + + Parent = self(), + Pids = [spawn(fun() -> + ok = py:set_env(shared_env), + R = py:eval(<<"SHARED">>), + Parent ! {self(), R} + end) || _ <- lists:seq(1, 4)], + + Results = [receive {Pid, R} -> R after 5000 -> timeout end || Pid <- Pids], + ct:log("Named env multiprocess results: ~p", [Results]), + + lists:foreach(fun(R) -> + {ok, 42} = R + end, Results), + + ok. + +test_list_envs(_Config) -> + {ok, _} = py:new_env(<<"A = 1">>, #{name => list_env_a}), + {ok, _} = py:new_env(<<"B = 2">>, #{name => list_env_b}), + + Envs = py:list_envs(), + ct:log("Listed envs: ~p", [Envs]), + true = lists:keymember(list_env_a, 1, Envs), + true = lists:keymember(list_env_b, 1, Envs), + + ok.