[ Web Proxy ]
URL:
Viewing: https://raw.githubusercontent.com/princemaple/pythonx/main/c_src/pythonx/pythonx.cpp [Back]  [Original]

#include 
#include 
#include 
#include 
#include 
#include 
#include 
#include 
#include 
#include 
#include 

#include "python.hpp"

extern "C" void pythonx_handle_io_write(const char *message,
                                        const char *eval_info_bytes, bool type);

namespace pythonx {

using namespace python;

// State
std::mutex init_mutex;
bool is_initialized = false;
std::wstring python_home_path_w;
std::map compilation_cache;
std::mutex compilation_cache_mutex;

// Wrapper around the Python Global Interpreter Lock (GIL).
//
// To acquire the GIL, the caller simply needs to initialize a new
// guard object. Once the guard object's lifetime ends, the GIL is
// automatically released. This is the use of RAII [1] pattern,
// similarly to `std::lock_guard`.
//
// [1]: https://en.wikipedia.org/wiki/Resource_acquisition_is_initialization
class PyGILGuard {
  PyGILState_STATE state;

public:
  PyGILGuard() { this->state = PyGILState_Ensure(); }

  ~PyGILGuard() { PyGILState_Release(this->state); }
};

// Ensures the given object refcount is decremented when the guard
// goes out of scope.
class PyDecRefGuard {
  PyObjectPtr py_object;

public:
  PyDecRefGuard() : py_object(nullptr) {}
  PyDecRefGuard(PyObjectPtr py_object) : py_object(py_object) {}

  ~PyDecRefGuard() {
    if (this->py_object != nullptr) {
      Py_DecRef(this->py_object);
    }
  }

  PyDecRefGuard &operator=(PyObjectPtr py_object) {
    this->py_object = py_object;
    return *this;
  }
};

void ensure_initialized() {
  auto init_guard = std::lock_guard(init_mutex);

  if (!is_initialized) {
    throw std::runtime_error("Python interpreter has not been initialized");
  }
}

namespace atoms {
auto ElixirPythonxError = fine::Atom("Elixir.Pythonx.Error");
auto ElixirPythonxJanitor = fine::Atom("Elixir.Pythonx.Janitor");
auto ElixirPythonxObject = fine::Atom("Elixir.Pythonx.Object");
auto decref = fine::Atom("decref");
auto integer = fine::Atom("integer");
auto list = fine::Atom("list");
auto map = fine::Atom("map");
auto map_set = fine::Atom("map_set");
auto output = fine::Atom("output");
auto ref = fine::Atom("ref");
auto traceback = fine::Atom("traceback");
auto tuple = fine::Atom("tuple");
auto type = fine::Atom("type");
auto value = fine::Atom("value");
} // namespace atoms

struct ExObjectResource {
  PyObjectPtr py_object;

  ExObjectResource(PyObjectPtr py_object) : py_object(py_object) {}

  void destructor(ErlNifEnv *env) {
    // Decrementing refcount requires GIL and we should not block in
    // the destructor, so we send a message to a known process and let
    // it decrement the refcount for us. Also see [1].
    //
    // [1]:https://erlangforums.com/t/how-to-deal-with-destructors-that-can-take-a-while-to-run-and-possibly-block-the-scheduler/4290

    if (!is_initialized) {
      // If we allow multiple initializations, we need to add a counter
      // and check that py_object comes from the current initialization
      return;
    }

    auto ptr = reinterpret_cast(this->py_object);

    auto janitor_name = fine::encode(env, atoms::ElixirPythonxJanitor);
    ErlNifPid janitor_pid;
    if (enif_whereis_pid(env, janitor_name, &janitor_pid)) {
      auto msg_env = enif_alloc_env();
      auto msg = fine::encode(msg_env, std::make_tuple(atoms::decref, ptr));
      enif_send(env, &janitor_pid, msg_env, msg);
      enif_free_env(msg_env);
    } else {
      std::cerr py_object, ex_key.resource->py_object,
                     ex_value.resource->py_object);
  raise_if_failed(env, result);

  return fine::Ok();
}

FINE_NIF(dict_set_item, ERL_NIF_DIRTY_JOB_CPU_BOUND);

ExObject tuple_new(ErlNifEnv *env, uint64_t size) {
  ensure_initialized();
  auto gil_guard = PyGILGuard();

  auto py_tuple = PyTuple_New(size);
  raise_if_failed(env, py_tuple);

  return ExObject(fine::make_resource(py_tuple));
}

FINE_NIF(tuple_new, ERL_NIF_DIRTY_JOB_CPU_BOUND);

fine::Ok tuple_set_item(ErlNifEnv *env, ExObject ex_object, uint64_t index,
                          ExObject ex_value) {
  ensure_initialized();
  auto gil_guard = PyGILGuard();

  auto result = PyTuple_SetItem(ex_object.resource->py_object, index,
                                ex_value.resource->py_object);
  raise_if_failed(env, result);

  // PyTuple_SetItem steals a reference, so we add one back
  Py_IncRef(ex_value.resource->py_object);

  return fine::Ok();
}

FINE_NIF(tuple_set_item, ERL_NIF_DIRTY_JOB_CPU_BOUND);

ExObject list_new(ErlNifEnv *env, uint64_t size) {
  ensure_initialized();
  auto gil_guard = PyGILGuard();

  auto py_tuple = PyList_New(size);
  raise_if_failed(env, py_tuple);

  return ExObject(fine::make_resource(py_tuple));
}

FINE_NIF(list_new, ERL_NIF_DIRTY_JOB_CPU_BOUND);

fine::Ok list_set_item(ErlNifEnv *env, ExObject ex_object, uint64_t index,
                         ExObject ex_value) {
  ensure_initialized();
  auto gil_guard = PyGILGuard();

  auto result = PyList_SetItem(ex_object.resource->py_object, index,
                               ex_value.resource->py_object);
  raise_if_failed(env, result);

  // PyList_SetItem steals a reference, so we add one back
  Py_IncRef(ex_value.resource->py_object);

  return fine::Ok();
}

FINE_NIF(list_set_item, ERL_NIF_DIRTY_JOB_CPU_BOUND);

ExObject set_new(ErlNifEnv *env) {
  ensure_initialized();
  auto gil_guard = PyGILGuard();

  auto py_set = PySet_New(NULL);
  raise_if_failed(env, py_set);

  return ExObject(fine::make_resource(py_set));
}

FINE_NIF(set_new, ERL_NIF_DIRTY_JOB_CPU_BOUND);

fine::Ok set_add(ErlNifEnv *env, ExObject ex_object, ExObject ex_key) {
  ensure_initialized();
  auto gil_guard = PyGILGuard();

  auto result =
      PySet_Add(ex_object.resource->py_object, ex_key.resource->py_object);
  raise_if_failed(env, result);

  return fine::Ok();
}

FINE_NIF(set_add, ERL_NIF_DIRTY_JOB_CPU_BOUND);

ExObject object_repr(ErlNifEnv *env, ExObject ex_object) {
  ensure_initialized();
  auto gil_guard = PyGILGuard();

  auto py_repr = PyObject_Repr(ex_object.resource->py_object);
  raise_if_failed(env, py_repr);

  return ExObject(fine::make_resource(py_repr));
}

FINE_NIF(object_repr, ERL_NIF_DIRTY_JOB_CPU_BOUND);

fine::Term format_exception(ErlNifEnv *env, ExError error) {
  ensure_initialized();
  auto gil_guard = PyGILGuard();

  auto py_traceback_module = PyImport_ImportModule("traceback");
  raise_if_failed(env, py_traceback_module);
  auto py_traceback_module_guard = PyDecRefGuard(py_traceback_module);

  auto format_exception =
      PyObject_GetAttrString(py_traceback_module, "format_exception");
  raise_if_failed(env, format_exception);
  auto format_exception_guard = PyDecRefGuard(format_exception);

  auto py_type = error.type.resource->py_object;
  auto py_value = error.value.resource->py_object;
  auto py_traceback = error.traceback.resource->py_object;

  auto format_exception_args = PyTuple_Pack(3, py_type, py_value, py_traceback);
  raise_if_failed(env, format_exception_args);
  auto format_exception_args_guard = PyDecRefGuard(format_exception_args);

  auto py_lines = PyObject_Call(format_exception, format_exception_args, NULL);
  raise_if_failed(env, py_lines);
  auto py_lines_guard = PyDecRefGuard(py_lines);

  auto size = PyList_Size(py_lines);
  raise_if_failed(env, size);

  auto terms = std::vector();
  terms.reserve(size);

  for (Py_ssize_t i = 0; i < size; i++) {
    auto py_line = PyList_GetItem(py_lines, i);
    raise_if_failed(env, py_line);

    terms.push_back(py_object_to_binary_term(env, py_line));
  }

  return enif_make_list_from_array(env, terms.data(),
                                   static_cast(size));
}

FINE_NIF(format_exception, ERL_NIF_DIRTY_JOB_CPU_BOUND);

fine::Term decode_once(ErlNifEnv *env, ExObject ex_object) {
  ensure_initialized();
  auto gil_guard = PyGILGuard();

  auto py_object = ex_object.resource->py_object;

  auto is_none = Py_IsNone(py_object);
  raise_if_failed(env, is_none);
  if (is_none) {
    return fine::encode(env, std::nullopt);
  }

  auto is_true = Py_IsTrue(py_object);
  raise_if_failed(env, is_true);
  if (is_true) {
    return fine::encode(env, true);
  }

  auto is_false = Py_IsFalse(py_object);
  raise_if_failed(env, is_false);
  if (is_false) {
    return fine::encode(env, false);
  }

  auto py_builtins = PyEval_GetBuiltins();
  raise_if_failed(env, py_builtins);

  auto py_int_type = PyDict_GetItemString(py_builtins, "int");
  raise_if_failed(env, py_int_type);
  auto is_long = PyObject_IsInstance(py_object, py_int_type);
  raise_if_failed(env, is_long);
  if (is_long) {
    int overflow;
    auto integer = PyLong_AsLongLongAndOverflow(py_object, &overflow);

    if (PyErr_Occurred() != NULL) {
      raise_py_error(env);
    }

    if (overflow == 0) {
      return enif_make_int64(env, integer);
    }

    // Integer over 64 bits

    auto py_str = PyObject_Str(py_object);
    raise_if_failed(env, py_str);
    auto py_str_guard = PyDecRefGuard(py_str);

    auto binary_term = py_object_to_binary_term(env, py_str);

    return fine::encode(
        env, std::make_tuple(atoms::integer, fine::Term(binary_term)));
  }

  auto py_float_type = PyDict_GetItemString(py_builtins, "float");
  raise_if_failed(env, py_float_type);
  auto is_float = PyObject_IsInstance(py_object, py_float_type);
  raise_if_failed(env, is_float);
  if (is_float) {
    double number = PyFloat_AsDouble(py_object);
    if (PyErr_Occurred() != NULL) {
      raise_py_error(env);
    }

    return enif_make_double(env, number);
  }

  auto py_tuple_type = PyDict_GetItemString(py_builtins, "tuple");
  raise_if_failed(env, py_tuple_type);
  auto is_tuple = PyObject_IsInstance(py_object, py_tuple_type);
  raise_if_failed(env, is_tuple);
  if (is_tuple) {
    auto size = PyTuple_Size(py_object);
    raise_if_failed(env, size);

    auto terms = std::vector();
    terms.reserve(size);

    for (Py_ssize_t i = 0; i < size; i++) {
      auto py_item = PyTuple_GetItem(py_object, i);
      raise_if_failed(env, py_item);
      Py_IncRef(py_item);
      auto ex_item = ExObject(fine::make_resource(py_item));
      terms.push_back(fine::encode(env, ex_item));
    }

    auto items = enif_make_list_from_array(env, terms.data(),
                                           static_cast(size));
    return fine::encode(env, std::make_tuple(atoms::tuple, fine::Term(items)));
  }

  auto py_list_type = PyDict_GetItemString(py_builtins, "list");
  raise_if_failed(env, py_list_type);
  auto is_list = PyObject_IsInstance(py_object, py_list_type);
  raise_if_failed(env, is_list);
  if (is_list) {
    auto size = PyList_Size(py_object);
    raise_if_failed(env, size);

    auto terms = std::vector();
    terms.reserve(size);

    for (Py_ssize_t i = 0; i < size; i++) {
      auto py_item = PyList_GetItem(py_object, i);
      raise_if_failed(env, py_item);
      Py_IncRef(py_item);
      auto ex_item = ExObject(fine::make_resource(py_item));
      terms.push_back(fine::encode(env, ex_item));
    }

    auto items = enif_make_list_from_array(env, terms.data(),
                                           static_cast(size));
    return fine::encode(env, std::make_tuple(atoms::list, fine::Term(items)));
  }

  auto py_dict_type = PyDict_GetItemString(py_builtins, "dict");
  raise_if_failed(env, py_dict_type);
  auto is_dict = PyObject_IsInstance(py_object, py_dict_type);
  raise_if_failed(env, is_dict);
  if (is_dict) {
    auto size = PyDict_Size(py_object);
    raise_if_failed(env, size);

    auto terms = std::vector();
    terms.reserve(size);

    PyObjectPtr py_key, py_value;
    Py_ssize_t pos = 0;

    while (PyDict_Next(py_object, &pos, &py_key, &py_value)) {
      Py_IncRef(py_key);
      auto ex_key = ExObject(fine::make_resource(py_key));

      Py_IncRef(py_value);
      auto ex_value = ExObject(fine::make_resource(py_value));

      terms.push_back(fine::encode(env, std::make_tuple(ex_key, ex_value)));
    }

    auto items = enif_make_list_from_array(env, terms.data(),
                                           static_cast(size));
    return fine::encode(env, std::make_tuple(atoms::map, fine::Term(items)));
  }

  auto py_str_type = PyDict_GetItemString(py_builtins, "str");
  raise_if_failed(env, py_str_type);
  auto is_unicode = PyObject_IsInstance(py_object, py_str_type);
  raise_if_failed(env, is_unicode);
  if (is_unicode) {
    return py_object_to_binary_term(env, py_object);
  }

  auto py_bytes_type = PyDict_GetItemString(py_builtins, "bytes");
  raise_if_failed(env, py_bytes_type);
  auto is_bytes = PyObject_IsInstance(py_object, py_bytes_type);
  raise_if_failed(env, is_bytes);
  if (is_bytes) {
    Py_ssize_t size;
    char *buffer;
    auto result = PyBytes_AsStringAndSize(py_object, &buffer, &size);
    raise_if_failed(env, result);

    ERL_NIF_TERM binary_term;
    auto data = enif_make_new_binary(env, size, &binary_term);
    memcpy(data, buffer, size);

    return binary_term;
  }

  auto py_set_type = PyDict_GetItemString(py_builtins, "set");
  raise_if_failed(env, py_set_type);
  auto is_set = PyObject_IsInstance(py_object, py_set_type);
  raise_if_failed(env, is_set);
  auto py_frozenset_type = PyDict_GetItemString(py_builtins, "frozenset");
  raise_if_failed(env, py_frozenset_type);
  auto is_frozenset = PyObject_IsInstance(py_object, py_frozenset_type);
  raise_if_failed(env, is_frozenset);
  if (is_set || is_frozenset) {
    auto size = PySet_Size(py_object);
    raise_if_failed(env, size);

    auto terms = std::vector();
    terms.reserve(size);

    auto py_iter = PyObject_GetIter(py_object);
    raise_if_failed(env, py_iter);
    auto py_iter_guard = PyDecRefGuard(py_iter);

    PyObjectPtr py_item = NULL;

    while ((py_item = PyIter_Next(py_iter)) != NULL) {
      // Note that PyIter_Next already returns a new reference
      auto ex_item = ExObject(fine::make_resource(py_item));
      terms.push_back(fine::encode(env, ex_item));
    }

    if (PyErr_Occurred() != NULL) {
      raise_py_error(env);
    }

    auto items = enif_make_list_from_array(env, terms.data(),
                                           static_cast(size));
    return fine::encode(env,
                        std::make_tuple(atoms::map_set, fine::Term(items)));
  }

  // None of the built-ins, return %Pythonx.Object{} as is
  return fine::encode(env, ex_object);
}

FINE_NIF(decode_once, ERL_NIF_DIRTY_JOB_CPU_BOUND);

std::tuple compile(ErlNifEnv *env,
                                             ErlNifBinary code) {
  // Python code can be compiled in either "exec" mode (multiple
  // statements with no result value) or "eval" mode (single expression
  // with a result value). We want our eval API to accept arbitrary
  // Python code with multiple statements, while also returning the
  // final result. To achieve this we parse the code using Python
  // standard library and check if the last statement is an expression.
  // If that is the case, we split the code, "exec" the statements
  // and "eval" the final expression separately.
  //
  // For the reference, below is a Python code corresponding to the
  // described logic. Technically we could "exec" that Python code,
  // but in order to avoid extra overhead we call the Python functions
  // directly via the C API.
  //
  //     import ast
  //
  //     module = ast.parse(code, "", mode="exec")
  //
  //     body_code = None
  //     last_expr_code = None
  //
  //     if module.body:
  //       last_statement = module.body[-1]
  //
  //       if isinstance(last_statement, ast.Expr):
  //         expr = ast.Expression(module.body.pop().value)
  //         # Copy positional information to the expression root node
  //         expr.lineno = last_statement.lineno
  //         expr.col_offset = last_statement.col_offset
  //         expr.end_col_offset = last_statement.end_col_offset
  //         last_expr_code = compile(expr, filename, mode="eval")
  //
  //     if module.body:
  //       body_code = compile(module, filename, mode="exec")
  //
  // The body and last expression is then evaluated separately, as in:
  //
  //     if body_code:
  //       eval(body_code)
  //
  //     if last_expr_code:
  //       result = eval(last_expr_code)
  //     else:
  //       result = None

  PyObjectPtr py_body_code = nullptr;
  PyObjectPtr py_last_expr_code = nullptr;
  auto py_last_expr_code_guard = PyDecRefGuard();

  auto py_ast = PyImport_ImportModule("ast");
  raise_if_failed(env, py_ast);
  auto py_ast_guard = PyDecRefGuard(py_ast);

  auto py_parse = PyObject_GetAttrString(py_ast, "parse");
  raise_if_failed(env, py_parse);
  auto py_parse_guard = PyDecRefGuard(py_parse);

  auto py_code = PyUnicode_FromStringAndSize(
      reinterpret_cast(code.data), code.size);
  raise_if_failed(env, py_code);
  auto py_code_guard = PyDecRefGuard(py_code);

  auto py_file_string = PyUnicode_FromStringAndSize("", 8);
  raise_if_failed(env, py_file_string);
  auto py_file_string_guard = PyDecRefGuard(py_file_string);

  auto py_exec_string = PyUnicode_FromStringAndSize("exec", 4);
  raise_if_failed(env, py_exec_string);
  auto py_exec_string_guard = PyDecRefGuard(py_exec_string);

  auto py_parse_args = PyTuple_Pack(3, py_code, py_file_string, py_exec_string);
  raise_if_failed(env, py_parse_args);
  auto py_parse_args_guard = PyDecRefGuard(py_parse_args);

  auto py_module_ast = PyObject_Call(py_parse, py_parse_args, NULL);
  raise_if_failed(env, py_module_ast);
  auto py_module_ast_guard = PyDecRefGuard(py_module_ast);

  auto py_builtins = PyEval_GetBuiltins();
  raise_if_failed(env, py_builtins);

  auto py_compile = PyDict_GetItemString(py_builtins, "compile");
  raise_if_failed(env, py_compile);

  auto py_module_body = PyObject_GetAttrString(py_module_ast, "body");
  raise_if_failed(env, py_module_body);
  auto py_module_body_guard = PyDecRefGuard(py_module_body);

  auto py_module_body_size = PyList_Size(py_module_body);
  raise_if_failed(env, py_module_body_size);

  if (py_module_body_size > 0) {
    auto py_last_expr = PyList_GetItem(py_module_body, py_module_body_size - 1);
    raise_if_failed(env, py_last_expr);

    auto py_Expr = PyObject_GetAttrString(py_ast, "Expr");
    raise_if_failed(env, py_Expr);
    auto py_Expr_guard = PyDecRefGuard(py_Expr);

    auto is_Expr_instance = PyObject_IsInstance(py_last_expr, py_Expr);
    raise_if_failed(env, is_Expr_instance);

    if (is_Expr_instance) {
      auto py_module_body_pop = PyObject_GetAttrString(py_module_body, "pop");
      raise_if_failed(env, py_module_body_pop);
      auto py_module_body_pop_guard = PyDecRefGuard(py_module_body_pop);
      py_module_body_size -= 1;

      py_last_expr = PyObject_CallNoArgs(py_module_body_pop);
      raise_if_failed(env, py_last_expr);
      auto py_last_statement_guard = PyDecRefGuard(py_last_expr);

      auto py_last_expr_value = PyObject_GetAttrString(py_last_expr, "value");
      raise_if_failed(env, py_last_expr_value);
      auto py_last_expr_value_guard = PyDecRefGuard(py_last_expr_value);

      auto py_Expression = PyObject_GetAttrString(py_ast, "Expression");
      raise_if_failed(env, py_Expression);
      auto py_Expression_guard = PyDecRefGuard(py_Expression);

      auto py_Expression_args = PyTuple_Pack(1, py_last_expr_value);
      raise_if_failed(env, py_Expression_args);
      auto py_Expression_args_guard = PyDecRefGuard(py_Expression_args);

      auto py_expr = PyObject_Call(py_Expression, py_Expression_args, NULL);
      raise_if_failed(env, py_expr);
      auto py_expr_guard = PyDecRefGuard(py_expr);

      for (const auto &attr_name : {"lineno", "col_offset", "end_col_offset"}) {
        auto attr_value = PyObject_GetAttrString(py_last_expr, attr_name);
        raise_if_failed(env, attr_value);
        auto attr_value_guard = PyDecRefGuard(attr_value);

        raise_if_failed(env,
                        PyObject_SetAttrString(py_expr, attr_name, attr_value));
      }

      auto py_eval_string = PyUnicode_FromStringAndSize("eval", 4);
      raise_if_failed(env, py_eval_string);
      auto py_eval_string_guard = PyDecRefGuard(py_eval_string);

      auto py_compile_args =
          PyTuple_Pack(3, py_expr, py_file_string, py_eval_string);
      raise_if_failed(env, py_compile_args);
      auto py_compile_args_guard = PyDecRefGuard(py_compile_args);

      py_last_expr_code = PyObject_Call(py_compile, py_compile_args, NULL);
      raise_if_failed(env, py_last_expr_code);

      py_last_expr_code_guard = py_last_expr_code;
    }
  }

  if (py_module_body_size > 0) {
    auto py_compile_args =
        PyTuple_Pack(3, py_module_ast, py_file_string, py_exec_string);
    raise_if_failed(env, py_compile_args);
    auto py_compile_args_guard = PyDecRefGuard(py_compile_args);

    py_body_code = PyObject_Call(py_compile, py_compile_args, NULL);
    raise_if_failed(env, py_body_code);
  }

  py_last_expr_code_guard = nullptr;

  return std::make_tuple(py_body_code, py_last_expr_code);
}

std::tuple
eval(ErlNifEnv *env, ErlNifBinary code, std::string code_md5,
     std::vector globals,
     ErlNifPid stdout_device, ErlNifPid stderr_device) {
  ensure_initialized();

  // Step 1: compile (or get cached result)

  PyObjectPtr py_body_code = nullptr;
  PyObjectPtr py_last_expr_code = nullptr;

  {
    // Note that it is important that we don't hold GIL while trying
    // to acquire the mutex, otherwise we could deadlock.
    auto guard = std::lock_guard(compilation_cache_mutex);

    if (compilation_cache.find(code_md5) == compilation_cache.end()) {
      auto gil_guard = PyGILGuard();
      auto compiled = compile(env, code);
      compilation_cache[code_md5] = compiled;
    }

    auto compiled = compilation_cache[code_md5];
    py_body_code = std::get(compiled);
    py_last_expr_code = std::get(compiled);
  }

  auto gil_guard = PyGILGuard();

  // Step 2: prepare globals

  // For globals, we create a new module named __main__ and use its
  // dict as globals (extended with the given entries). It corresponds
  // to the following Python code:
  //
  //     import types
  //     import sys
  //
  //     main_module = types.ModuleType("__main__")
  //     main_module.__dict__["__builtins__"] = builtins
  //     sys.modules["__main__"] = main_module
  //
  //     main_module.__dict__["x"] = 1
  //     main_module.__dict__["y"] = 2
  //
  //     eval(..., main_module.__dict__)
  //
  // Note that a more straightforward approach would be to get the
  // default __main__ module, copy its dict and use that as globals.
  // However, to better mirror actual Python execution, we want to
  // use an actual module dict as globals.
  //
  // A practical scenario where this matters is pickling an object
  // of a class defined via evaluation. The pickle module consults
  // sys.modules (in this case sys.modules["__main__"]) and looks up
  // the class or function name. If we use a plain dict as globals,
  // the class will be defined only in that dict and such lookups wil
  // fail.
  //
  // However, it is worth noting that the current approach is not
  // perfect as it can fail under race conditions. Evaluations may
  // happen concurrently (if one of them releases GIL, for example,
  // by calling time.sleep) and there can be other threads started
  // by evaluation. If a new evaluation sets sys.modules["__main__"]
  // and yields, an older evaluation or thread may resume and at that
  // point the value of sys.modules["__main__"] is no longer accurate.
  // For more details see [1].
  //
  // [1]: https://github.com/marimo-team/marimo/pull/811

  auto py_types = PyImport_ImportModule("types");
  raise_if_failed(env, py_types);
  auto py_types_guard = PyDecRefGuard(py_types);

  auto py_ModuleType = PyObject_GetAttrString(py_types, "ModuleType");
  raise_if_failed(env, py_ModuleType);
  auto py_ModuleType_guard = PyDecRefGuard(py_ModuleType);

  auto py_main_module_string = PyUnicode_FromStringAndSize("__main__", 8);
  raise_if_failed(env, py_main_module_string);
  auto py_main_module_string_guard = PyDecRefGuard(py_main_module_string);

  auto py_ModuleType_args = PyTuple_Pack(1, py_main_module_string);
  raise_if_failed(env, py_ModuleType_args);
  auto py_ModuleType_args_guard = PyDecRefGuard(py_ModuleType_args);

  auto py_main_module = PyObject_Call(py_ModuleType, py_ModuleType_args, NULL);
  raise_if_failed(env, py_main_module);
  auto py_main_module_guard = PyDecRefGuard(py_main_module);

  auto py_globals = PyModule_GetDict(py_main_module);
  raise_if_failed(env, py_globals);

  auto py_builtins = PyEval_GetBuiltins();
  raise_if_failed(env, py_builtins);

  raise_if_failed(
      env, PyDict_SetItemString(py_globals, "__builtins__", py_builtins));

  // The IO capture consists of the following steps:
  //
  //   1. We dump EvalInfo into Python bytes and store it in globals
  //      as __pythonx_eval_info_bytes__.
  //
  //   2. When IO happens, our custom sys.stdout.write (overriden on
  //      init) retrieves the info from globals and calls the
  //      pythonx_handle_io_write C function, passing the info.
  //
  //   3. The pythonx_handle_io_write C function loads EvalInfo and
  //      uses it to send messages to the specified process.
  //
  // Each step has a bit more specifics, but they are explained alongside
  // the corresponding code.

  // Here we dump the EvalInfo struct into Python bytes object and
  // put it in globals. We need to pass all of the data, rather than
  // a pointer, because it may be used during IO from Python threads,
  // even after the NIF finished. The ErlNifPid struct is opaque, so
  // we need to store struct memory contents anyway, hence we do this
  // for the whole EvalInfo struct at once.
  //
  // Note that copying struct memory contents into a byte buffer and
  // vice versa is safe, as long as the struct is always allocated
  // using its own type (to guarantee proper alignment). Contrarily,
  // casting a buffer (such as char*) to struct pointer is not safe.
  auto eval_info = EvalInfo{};
  eval_info.stdout_device = stdout_device;
  eval_info.stderr_device = stderr_device;
  eval_info.env = env;
  eval_info.thread_id = std::this_thread::get_id();

  auto py_eval_info_bytes = PyBytes_FromStringAndSize(
      reinterpret_cast(&eval_info), sizeof(EvalInfo));
  raise_if_failed(env, py_eval_info_bytes);

  raise_if_failed(env, PyDict_SetItemString(py_globals,
                                            "__pythonx_eval_info_bytes__",
                                            py_eval_info_bytes));

  auto py_sys = PyImport_AddModule("sys");
  raise_if_failed(env, py_sys);

  auto py_modules = PyObject_GetAttrString(py_sys, "modules");
  raise_if_failed(env, py_modules);
  auto py_modules_guard = PyDecRefGuard(py_modules);

  raise_if_failed(env,
                  PyDict_SetItemString(py_modules, "__main__", py_main_module));

  auto py_globals_initial = PyDict_Copy(py_globals);
  raise_if_failed(env, py_globals_initial);
  auto py_globals_guard = PyDecRefGuard(py_globals_initial);

  for (const auto &[key, value] : globals) {
    auto py_key = PyUnicode_FromStringAndSize(
        reinterpret_cast(key.data), key.size);
    raise_if_failed(env, py_key);

    auto result = PyDict_SetItem(py_globals, py_key, value.resource->py_object);
    Py_DecRef(py_key);
    raise_if_failed(env, result);
  }

  // Step 3: eval body and expression

  if (py_body_code != nullptr) {
    auto py_body_result = PyEval_EvalCode(py_body_code, py_globals, py_globals);
    raise_if_failed(env, py_body_result);
    Py_DecRef(py_body_result);
  }

  auto result = std::optional();

  if (py_last_expr_code != nullptr) {
    auto py_result = PyEval_EvalCode(py_last_expr_code, py_globals, py_globals);
    raise_if_failed(env, py_result);
    result = ExObject(fine::make_resource(py_result));
  }

  // Step 4: flat-decode globals

  std::vector key_terms;
  std::vector value_terms;

  PyObjectPtr py_key, py_value;
  Py_ssize_t pos = 0;

  auto py_str_type = PyDict_GetItemString(py_builtins, "str");
  raise_if_failed(env, py_str_type);

  while (PyDict_Next(py_globals, &pos, &py_key, &py_value)) {
    // If the key was present in the default globals, ignore it
    if (PyDict_GetItem(py_globals_initial, py_key) != NULL) {
      continue;
    }

    auto is_unicode = PyObject_IsInstance(py_key, py_str_type);
    raise_if_failed(env, is_unicode);
    // If the key is not a string, ignore it. This can happen if the
    // globals dict is modified directly, but that's not expected.
    if (!is_unicode) {
      continue;
    }

    auto key_term = py_object_to_binary_term(env, py_key);
    key_terms.push_back(key_term);

    // Incref before making the resource
    Py_IncRef(py_value);
    auto ex_value = ExObject(fine::make_resource(py_value));
    value_terms.push_back(fine::encode(env, ex_value));
  }

  ERL_NIF_TERM map;
  if (!enif_make_map_from_arrays(env, key_terms.data(), value_terms.data(),
                                 key_terms.size(), &map)) {
    throw std::runtime_error("failed to make a map");
  }

  return std::make_tuple(result, map);
}

FINE_NIF(eval, ERL_NIF_DIRTY_JOB_CPU_BOUND);

} // namespace pythonx

FINE_INIT("Elixir.Pythonx.NIF");

// Below are functions we call from Python code

extern "C" void pythonx_handle_io_write(const char *message,
                                        const char *eval_info_bytes,
                                        bool type) {
  // Note that we allocate EvalInfo first, so it will have the proper
  // alignment and memcpy simply restores the original struct state.
  auto eval_info = pythonx::EvalInfo{};
  std::memcpy(&eval_info, eval_info_bytes, sizeof(pythonx::EvalInfo));

  auto env = enif_alloc_env();

  // The enif_whereis_pid and enif_send functions require passing the
  // caller env. Stdout write may be called by the evaluated code from
  // the NIF call, but it may also be called by a Python thread, after
  // the NIF call already finished. Since the BEAM uses OS threads for
  // its schedulers, we can simply check if this function is invoked
  // in the same thread as the NIF, or in a different one (Python thread).
  bool is_main_thread = std::this_thread::get_id() == eval_info.thread_id;
  auto caller_env = is_main_thread ? eval_info.env : NULL;

  // Note that we send the output to Pythonx.Janitor and it then sends
  // it to the device. We do this to avoid IO replies being sent to
  // the calling Elixir process (which would be unexpected).
  auto janitor_name = fine::encode(env, pythonx::atoms::ElixirPythonxJanitor);
  ErlNifPid janitor_pid;
  if (enif_whereis_pid(caller_env, janitor_name, &janitor_pid)) {
    auto device = type == 0 ? eval_info.stdout_device : eval_info.stderr_device;

    auto msg = fine::encode(env, std::make_tuple(pythonx::atoms::output,
                                                 std::string(message), device));
    enif_send(caller_env, &janitor_pid, env, msg);
    enif_free_env(env);
  } else {
    std::cerr 

Web Proxy Viewer  |  New URL  |  Original Page