Mirror Elixir env vars into os.environ on initialization

Author: jonatanklosko

c_src/python.cpp

@@ -56,6 +56,7 @@ DEF_SYMBOL(PyObject_GetIter)
 DEF_SYMBOL(PyObject_IsInstance)
 DEF_SYMBOL(PyObject_Repr)
 DEF_SYMBOL(PyObject_SetAttrString)
+DEF_SYMBOL(PyObject_SetItem)
 DEF_SYMBOL(PyObject_Str)
 DEF_SYMBOL(PySet_Add)
 DEF_SYMBOL(PySet_New)
@@ -130,6 +131,7 @@ void load_python_library(std::string path) {
   LOAD_SYMBOL(python_library, PyObject_IsInstance)
   LOAD_SYMBOL(python_library, PyObject_Repr)
   LOAD_SYMBOL(python_library, PyObject_SetAttrString)
+  LOAD_SYMBOL(python_library, PyObject_SetItem)
   LOAD_SYMBOL(python_library, PyObject_Str)
   LOAD_SYMBOL(python_library, PySet_Add)
   LOAD_SYMBOL(python_library, PySet_New)

c_src/python.hpp

@@ -110,6 +110,7 @@ extern PyObjectPtr (*PyObject_GetIter)(PyObjectPtr);
 extern int (*PyObject_IsInstance)(PyObjectPtr, PyObjectPtr);
 extern PyObjectPtr (*PyObject_Repr)(PyObjectPtr);
 extern int (*PyObject_SetAttrString)(PyObjectPtr, const char *, PyObjectPtr);
+extern int (*PyObject_SetItem)(PyObjectPtr, PyObjectPtr, PyObjectPtr);
 extern PyObjectPtr (*PyObject_Str)(PyObjectPtr);
 extern int (*PySet_Add)(PyObjectPtr, PyObjectPtr);
 extern PyObjectPtr (*PySet_New)(PyObjectPtr);

c_src/pythonx.cpp

@@ -9,6 +9,7 @@
 #include <string>
 #include <thread>
 #include <tuple>
+#include <unordered_map>
 
 #include "python.hpp"
 
@@ -287,7 +288,8 @@ ERL_NIF_TERM py_str_to_binary_term(ErlNifEnv *env, PyObjectPtr py_object) {
 fine::Ok<> init(ErlNifEnv *env, std::string python_dl_path,
                 ErlNifBinary python_home_path,
                 ErlNifBinary python_executable_path,
-                std::vector<ErlNifBinary> sys_paths) {
+                std::vector<ErlNifBinary> sys_paths,
+                std::vector<std::tuple<ErlNifBinary, ErlNifBinary>> envs) {
   auto init_guard = std::lock_guard<std::mutex>(init_mutex);
 
   if (is_initialized) {
@@ -377,6 +379,38 @@ fine::Ok<> init(ErlNifEnv *env, std::string python_dl_path,
     raise_if_failed(env, PyList_Append(py_sys_path, py_path));
   }
 
+  // We set env vars to match Elixir at the time of initialization.
+  // Note that the interpreter initializes its env vars from the OS
+  // process, however we also want to account for env vars set
+  // dynamically, for example via System.put_env/2.
+
+  auto py_os = PyImport_AddModule("os");
+  raise_if_failed(env, py_os);
+
+  auto py_os_environ = PyObject_GetAttrString(py_os, "environ");
+  raise_if_failed(env, py_os_environ);
+  auto py_os_environ_guard = PyDecRefGuard(py_os_environ);
+
+  auto py_os_environ_clear = PyObject_GetAttrString(py_os_environ, "clear");
+  raise_if_failed(env, py_os_environ_clear);
+  auto py_os_environ_clear_guard = PyDecRefGuard(py_os_environ_clear);
+  auto result = PyObject_CallNoArgs(py_os_environ_clear);
+  raise_if_failed(env, result);
+
+  for (const auto &[key, value] : envs) {
+    auto py_key = PyUnicode_FromStringAndSize(
+        reinterpret_cast<const char *>(key.data), key.size);
+    raise_if_failed(env, py_key);
+    auto py_key_guard = PyDecRefGuard(py_key);
+    auto py_value = PyUnicode_FromStringAndSize(
+        reinterpret_cast<const char *>(value.data), value.size);
+    raise_if_failed(env, py_value);
+    auto py_value_guard = PyDecRefGuard(py_value);
+
+    auto result = PyObject_SetItem(py_os_environ, py_key, py_value);
+    raise_if_failed(env, result);
+  }
+
   // Define global stdout and stdin overrides
 
   auto py_builtins = PyEval_GetBuiltins();

lib/pythonx.ex

@@ -47,6 +47,18 @@ defmodule Pythonx do
 
   For more configuration options, refer to the [uv documentation](https://docs.astral.sh/uv/concepts/projects/dependencies/).
 
+  > #### Environment variables {: .info}
+  >
+  > As part of the initialization, Python's `os.environ` is modified
+  > to match `System.get_env/1`. Therefore, `os.environ` accounts for
+  > prior changes to the environment, such as `System.put_env/2`.
+  >
+  > Subsequent changes to the environment, both via Elixir and Python,
+  > are not synchronized.
+  >
+  > Also note that contrarily to Elixir, changes to `os.environ` are
+  > automatically mirrored to the OS process environment.
+
   ## Options
 
     * `:force` - if true, runs with empty project cache. Defaults to `false`.
@@ -196,7 +208,18 @@ defmodule Pythonx do
       raise ArgumentError, "the given python executable does not exist: #{python_executable_path}"
     end
 
-    Pythonx.NIF.init(python_dl_path, python_home_path, python_executable_path, opts[:sys_paths])
+    envs =
+      for {k, v} <- :os.env() do
+        {IO.chardata_to_string(k), IO.chardata_to_string(v)}
+      end
+
+    Pythonx.NIF.init(
+      python_dl_path,
+      python_home_path,
+      python_executable_path,
+      opts[:sys_paths],
+      envs
+    )
   end
 
   @doc ~S'''

lib/pythonx/nif.ex

@@ -12,7 +12,9 @@ defmodule Pythonx.NIF do
     end
   end
 
-  def init(_python_dl_path, _python_home_path, _python_executable_path, _sys_paths), do: err!()
+  def init(_python_dl_path, _python_home_path, _python_executable_path, _sys_paths, _envs),
+    do: err!()
+
   def janitor_decref(_ptr), do: err!()
   def none_new(), do: err!()
   def false_new(), do: err!()

mix.lock

@@ -3,7 +3,7 @@
   "earmark_parser": {:hex, :earmark_parser, "1.4.42", "f23d856f41919f17cd06a493923a722d87a2d684f143a1e663c04a2b93100682", [:mix], [], "hexpm", "6915b6ca369b5f7346636a2f41c6a6d78b5af419d61a611079189233358b8b8b"},
   "elixir_make": {:hex, :elixir_make, "0.9.0", "6484b3cd8c0cee58f09f05ecaf1a140a8c97670671a6a0e7ab4dc326c3109726", [:mix], [], "hexpm", "db23d4fd8b757462ad02f8aa73431a426fe6671c80b200d9710caf3d1dd0ffdb"},
   "ex_doc": {:hex, :ex_doc, "0.36.1", "4197d034f93e0b89ec79fac56e226107824adcce8d2dd0a26f5ed3a95efc36b1", [:mix], [{:earmark_parser, "~> 1.4.42", [hex: :earmark_parser, repo: "hexpm", optional: false]}, {:makeup_c, ">= 0.1.0", [hex: :makeup_c, repo: "hexpm", optional: true]}, {:makeup_elixir, "~> 0.14 or ~> 1.0", [hex: :makeup_elixir, repo: "hexpm", optional: false]}, {:makeup_erlang, "~> 0.1 or ~> 1.0", [hex: :makeup_erlang, repo: "hexpm", optional: false]}, {:makeup_html, ">= 0.1.0", [hex: :makeup_html, repo: "hexpm", optional: true]}], "hexpm", "d7d26a7cf965dacadcd48f9fa7b5953d7d0cfa3b44fa7a65514427da44eafd89"},
-  "fine": {:hex, :fine, "0.1.2", "85cf7dd190c7c6c54c2840754ae977c9acc0417316255b674fad9f2678e4ecc7", [:mix], [], "hexpm", "9113531982c2b60dbea6c7233917ddf16806947cd7104b5d03011bf436ca3072"},
+  "fine": {:hex, :fine, "0.1.4", "b19a89c1476c7c57afb5f9314aed5960b5bc95d5277de4cb5ee8e1d1616ce379", [:mix], [], "hexpm", "be3324cc454a42d80951cf6023b9954e9ff27c6daa255483b3e8d608670303f5"},
   "makeup": {:hex, :makeup, "1.2.1", "e90ac1c65589ef354378def3ba19d401e739ee7ee06fb47f94c687016e3713d1", [:mix], [{:nimble_parsec, "~> 1.4", [hex: :nimble_parsec, repo: "hexpm", optional: false]}], "hexpm", "d36484867b0bae0fea568d10131197a4c2e47056a6fbe84922bf6ba71c8d17ce"},
   "makeup_elixir": {:hex, :makeup_elixir, "1.0.1", "e928a4f984e795e41e3abd27bfc09f51db16ab8ba1aebdba2b3a575437efafc2", [:mix], [{:makeup, "~> 1.0", [hex: :makeup, repo: "hexpm", optional: false]}, {:nimble_parsec, "~> 1.2.3 or ~> 1.3", [hex: :nimble_parsec, repo: "hexpm", optional: false]}], "hexpm", "7284900d412a3e5cfd97fdaed4f5ed389b8f2b4cb49efc0eb3bd10e2febf9507"},
   "makeup_erlang": {:hex, :makeup_erlang, "1.0.1", "c7f58c120b2b5aa5fd80d540a89fdf866ed42f1f3994e4fe189abebeab610839", [:mix], [{:makeup, "~> 1.0", [hex: :makeup, repo: "hexpm", optional: false]}], "hexpm", "8a89a1eeccc2d798d6ea15496a6e4870b75e014d1af514b1b71fa33134f57814"},

test/pythonx_test.exs

@@ -477,6 +477,22 @@ defmodule PythonxTest do
     end
   end
 
+  test "inherits env vars from elixir" do
+    # We set PYTHONX_TEST_ENV_VAR in test_helper.exs, before initializing
+    # Pythonx. That env var should be available to Python.
+
+    assert {result, %{}} =
+             Pythonx.eval(
+               """
+               import os
+               os.environ["PYTHONX_TEST_ENV_VAR"]
+               """,
+               %{}
+             )
+
+    assert repr(result) == "'value'"
+  end
+
   defp repr(object) do
     assert %Pythonx.Object{} = object
 

test/test_helper.exs

@@ -1,5 +1,7 @@
 python_minor = System.get_env("PYTHONX_TEST_PYTHON_MINOR", "13") |> String.to_integer()
 
+System.put_env("PYTHONX_TEST_ENV_VAR", "value")
+
 Pythonx.uv_init("""
 [project]
 name = "project"