diff --git a/src/py/py.hpp b/src/py/py.hpp index fe0615346..030233ca6 100644 --- a/src/py/py.hpp +++ b/src/py/py.hpp @@ -71,6 +71,13 @@ class Object final { operator bool() const { return ptr_; } + /// Release the ownership on this PyObject, i.e. we steal the reference. + PyObject *Steal() { + auto *p = ptr_; + ptr_ = nullptr; + return p; + } + /// Equivalent to `str(o)` in Python. /// /// Returned Object is nullptr if an error occurred. diff --git a/src/query/CMakeLists.txt b/src/query/CMakeLists.txt index f7023af74..eb1b26fc9 100644 --- a/src/query/CMakeLists.txt +++ b/src/query/CMakeLists.txt @@ -27,6 +27,7 @@ set(mg_query_sources plan/variable_start_planner.cpp procedure/mg_procedure_impl.cpp procedure/module.cpp + procedure/py_module.cpp typed_value.cpp) add_library(mg-query STATIC ${mg_query_sources}) diff --git a/src/query/procedure/py_module.cpp b/src/query/procedure/py_module.cpp new file mode 100644 index 000000000..40b30c680 --- /dev/null +++ b/src/query/procedure/py_module.cpp @@ -0,0 +1,53 @@ +#include "query/procedure/py_module.hpp" + +#include "query/procedure/mg_procedure_impl.hpp" + +namespace query::procedure { + +py::Object MgpValueToPyObject(const mgp_value &value) { + switch (mgp_value_get_type(&value)) { + case MGP_VALUE_TYPE_NULL: + Py_INCREF(Py_None); + return py::Object(Py_None); + case MGP_VALUE_TYPE_BOOL: + return py::Object(PyBool_FromLong(mgp_value_get_bool(&value))); + case MGP_VALUE_TYPE_INT: + return py::Object(PyLong_FromLongLong(mgp_value_get_int(&value))); + case MGP_VALUE_TYPE_DOUBLE: + return py::Object(PyFloat_FromDouble(mgp_value_get_double(&value))); + case MGP_VALUE_TYPE_STRING: + return py::Object(PyUnicode_FromString(mgp_value_get_string(&value))); + case MGP_VALUE_TYPE_LIST: { + const auto *list = mgp_value_get_list(&value); + const size_t len = mgp_list_size(list); + py::Object py_list(PyList_New(len)); + CHECK(py_list); + for (size_t i = 0; i < len; ++i) { + auto elem = MgpValueToPyObject(*mgp_list_at(list, i)); + CHECK(elem); + // Explicitly convert `py_list`, which is `py::Object`, via static_cast. + // Then the macro will cast it to `PyList *`. + PyList_SET_ITEM(static_cast(py_list), i, elem.Steal()); + } + return py_list; + } + case MGP_VALUE_TYPE_MAP: { + const auto *map = mgp_value_get_map(&value); + py::Object py_dict(PyDict_New()); + CHECK(py_dict); + for (const auto &[key, val] : map->items) { + auto py_val = MgpValueToPyObject(val); + CHECK(py_val); + // Unlike PyList_SET_ITEM, PyDict_SetItem does not steal the value. + CHECK(PyDict_SetItemString(py_dict, key.c_str(), py_val) == 0); + } + return py_dict; + } + case MGP_VALUE_TYPE_VERTEX: + case MGP_VALUE_TYPE_EDGE: + case MGP_VALUE_TYPE_PATH: + throw utils::NotYetImplemented("MgpValueToPyObject"); + } +} + +} // namespace query::procedure diff --git a/src/query/procedure/py_module.hpp b/src/query/procedure/py_module.hpp new file mode 100644 index 000000000..c39b48014 --- /dev/null +++ b/src/query/procedure/py_module.hpp @@ -0,0 +1,13 @@ +/// @file +/// Functions and types for loading Query Modules written in Python. +#pragma once + +#include "py/py.hpp" + +struct mgp_value; + +namespace query::procedure { + +py::Object MgpValueToPyObject(const mgp_value &); + +} // namespace query::procedure diff --git a/tests/unit/CMakeLists.txt b/tests/unit/CMakeLists.txt index 28ffa90c5..188dc1d72 100644 --- a/tests/unit/CMakeLists.txt +++ b/tests/unit/CMakeLists.txt @@ -77,6 +77,10 @@ target_include_directories(${test_prefix}query_procedure_mgp_type PRIVATE ${CMAK add_unit_test(query_procedure_mgp_module.cpp) target_link_libraries(${test_prefix}query_procedure_mgp_module mg-query) target_include_directories(${test_prefix}query_procedure_mgp_module PRIVATE ${CMAKE_SOURCE_DIR}/include) + +add_unit_test(query_procedure_py_module.cpp) +target_link_libraries(${test_prefix}query_procedure_py_module mg-query) +target_include_directories(${test_prefix}query_procedure_py_module PRIVATE ${CMAKE_SOURCE_DIR}/include) # END query/procedure add_unit_test(query_profile.cpp) diff --git a/tests/unit/query_procedure_py_module.cpp b/tests/unit/query_procedure_py_module.cpp new file mode 100644 index 000000000..120164127 --- /dev/null +++ b/tests/unit/query_procedure_py_module.cpp @@ -0,0 +1,78 @@ +#include + +#include + +#include "query/procedure/mg_procedure_impl.hpp" +#include "query/procedure/py_module.hpp" + +TEST(PyModule, MgpValueToPyObject) { + mgp_memory memory{utils::NewDeleteResource()}; + auto *list = mgp_list_make_empty(42, &memory); + { + // Create a list: [null, false, true, 42, 0.1, "some text"] + auto primitive_values = {mgp_value_make_null(&memory), + mgp_value_make_bool(0, &memory), + mgp_value_make_bool(1, &memory), + mgp_value_make_int(42, &memory), + mgp_value_make_double(0.1, &memory), + mgp_value_make_string("some text", &memory)}; + for (auto *val : primitive_values) { + mgp_list_append(list, val); + mgp_value_destroy(val); + } + } + auto *list_val = mgp_value_make_list(list); + auto *map = mgp_map_make_empty(&memory); + mgp_map_insert(map, "list", list_val); + mgp_value_destroy(list_val); + auto *map_val = mgp_value_make_map(map); + auto gil = py::EnsureGIL(); + auto py_dict = query::procedure::MgpValueToPyObject(*map_val); + mgp_value_destroy(map_val); + // We should now have in Python: + // {"list": [None, False, True, 42, 0.1, "some text"]} + ASSERT_TRUE(PyDict_Check(py_dict)); + EXPECT_EQ(PyDict_Size(py_dict), 1); + PyObject *key = nullptr; + PyObject *value = nullptr; + Py_ssize_t pos = 0; + while (PyDict_Next(py_dict, &pos, &key, &value)) { + ASSERT_TRUE(PyUnicode_Check(key)); + EXPECT_EQ(std::string(PyUnicode_AsUTF8(key)), "list"); + ASSERT_TRUE(PyList_Check(value)); + ASSERT_EQ(PyList_Size(value), 6); + EXPECT_EQ(PyList_GetItem(value, 0), Py_None); + EXPECT_EQ(PyList_GetItem(value, 1), Py_False); + EXPECT_EQ(PyList_GetItem(value, 2), Py_True); + auto *py_long = PyList_GetItem(value, 3); + ASSERT_TRUE(PyLong_Check(py_long)); + EXPECT_EQ(PyLong_AsLong(py_long), 42); + auto *py_float = PyList_GetItem(value, 4); + ASSERT_TRUE(PyFloat_Check(py_float)); + EXPECT_EQ(PyFloat_AsDouble(py_float), 0.1); + auto *py_str = PyList_GetItem(value, 5); + ASSERT_TRUE(PyUnicode_Check(py_str)); + EXPECT_EQ(std::string(PyUnicode_AsUTF8(py_str)), "some text"); + } + // TODO: Vertex, Edge and Path values +} + +int main(int argc, char **argv) { + ::testing::InitGoogleTest(&argc, argv); + // Initialize Python + auto *program_name = Py_DecodeLocale(argv[0], nullptr); + CHECK(program_name); + // Set program name, so Python can find its way to runtime libraries relative + // to executable. + Py_SetProgramName(program_name); + Py_InitializeEx(0 /* = initsigs */); + PyEval_InitThreads(); + int test_result; + Py_BEGIN_ALLOW_THREADS; + test_result = RUN_ALL_TESTS(); + Py_END_ALLOW_THREADS; + // Shutdown Python + Py_Finalize(); + PyMem_RawFree(program_name); + return test_result; +}