From b6b6403b06d2fcb76d0263dab66540fe9b0d1e25 Mon Sep 17 00:00:00 2001 From: Ivan Paljak Date: Thu, 27 Feb 2020 15:13:26 +0100 Subject: [PATCH] Implement PyEdgesIterator and integrate with PyVertex Reviewers: teon.banek Reviewed By: teon.banek Subscribers: pullbot Differential Revision: https://phabricator.memgraph.io/D2692 --- include/mgp.py | 16 +++- src/query/procedure/py_module.cpp | 118 +++++++++++++++++++++++++++++- 2 files changed, 131 insertions(+), 3 deletions(-) diff --git a/include/mgp.py b/include/mgp.py index d580b5b6d..8ce6d5232 100644 --- a/include/mgp.py +++ b/include/mgp.py @@ -226,14 +226,26 @@ class Vertex: '''Raise InvalidVertexError.''' if not self.is_valid(): raise InvalidVertexError() - raise NotImplementedError() + edges_it = self._vertex.iter_in_edges() + edge = edges_it.get() + while edge is not None: + yield Edge(edge) + if not self.is_valid(): + raise InvalidVertexError() + edge = edges_it.next() @property def out_edges(self) -> typing.Iterable[Edge]: '''Raise InvalidVertexError.''' if not self.is_valid(): raise InvalidVertexError() - raise NotImplementedError() + edges_it = self._vertex.iter_out_edges() + edge = edges_it.get() + while edge is not None: + yield Edge(edge) + if not self.is_valid(): + raise InvalidVertexError() + edge = edges_it.next() def __eq__(self, other) -> bool: '''Raise InvalidVertexError''' diff --git a/src/query/procedure/py_module.cpp b/src/query/procedure/py_module.cpp index 1bdf799da..4b9a03201 100644 --- a/src/query/procedure/py_module.cpp +++ b/src/query/procedure/py_module.cpp @@ -44,6 +44,7 @@ void PyVerticesIteratorDealloc(PyVerticesIterator *self) { // execution, so we may cause a double free issue. if (self->py_graph->graph) mgp_vertices_iterator_destroy(self->it); Py_DECREF(self->py_graph); + Py_TYPE(self)->tp_free(self); } PyObject *PyVerticesIteratorGet(PyVerticesIterator *self, @@ -87,6 +88,66 @@ static PyTypeObject PyVerticesIteratorType = { .tp_dealloc = reinterpret_cast(PyVerticesIteratorDealloc), }; +struct PyEdgesIterator { + PyObject_HEAD + mgp_edges_iterator *it; + PyGraph *py_graph; +}; + +PyObject *MakePyEdge(mgp_edge *edge, PyGraph *py_graph); + +void PyEdgesIteratorDealloc(PyEdgesIterator *self) { + CHECK(self->it); + CHECK(self->py_graph); + // Avoid invoking `mgp_edges_iterator_destroy` if we are not in valid + // execution context. The query execution should free all memory used during + // execution, so we may cause a double free issue. + if (self->py_graph->graph) mgp_edges_iterator_destroy(self->it); + Py_DECREF(self->py_graph); + Py_TYPE(self)->tp_free(self); +} + +PyObject *PyEdgesIteratorGet(PyEdgesIterator *self, + PyObject *Py_UNUSED(ignored)) { + CHECK(self->it); + CHECK(self->py_graph); + CHECK(self->py_graph->graph); + const auto *edge = mgp_edges_iterator_get(self->it); + if (!edge) Py_RETURN_NONE; + return MakePyEdge(mgp_edge_copy(edge, self->py_graph->memory), + self->py_graph); +} + +PyObject *PyEdgesIteratorNext(PyEdgesIterator *self, + PyObject *Py_UNUSED(ignored)) { + CHECK(self->it); + CHECK(self->py_graph); + CHECK(self->py_graph->graph); + const auto *edge = mgp_edges_iterator_next(self->it); + if (!edge) Py_RETURN_NONE; + return MakePyEdge(mgp_edge_copy(edge, self->py_graph->memory), + self->py_graph); +} + +static PyMethodDef PyEdgesIteratorMethods[] = { + {"get", reinterpret_cast(PyEdgesIteratorGet), METH_NOARGS, + "Get the current edge pointed to by the iterator or return None."}, + {"next", reinterpret_cast(PyEdgesIteratorNext), METH_NOARGS, + "Advance the iterator to the next edge and return it."}, + {nullptr}, +}; + +static PyTypeObject PyEdgesIteratorType = { + PyVarObject_HEAD_INIT(nullptr, 0) + .tp_name = "_mgp.EdgesIterator", + .tp_doc = "Wraps struct mgp_edges_iterator.", + .tp_basicsize = sizeof(PyEdgesIterator), + .tp_flags = Py_TPFLAGS_DEFAULT, + .tp_new = PyType_GenericNew, + .tp_methods = PyEdgesIteratorMethods, + .tp_dealloc = reinterpret_cast(PyEdgesIteratorDealloc), +}; + PyObject *PyGraphIsValid(PyGraph *self, PyObject *Py_UNUSED(ignored)) { return PyBool_FromLong(!!self->graph); } @@ -118,7 +179,7 @@ PyObject *PyGraphIterVertices(PyGraph *self, PyObject *Py_UNUSED(ignored)) { } auto *py_vertices_it = PyObject_New(PyVerticesIterator, &PyVerticesIteratorType); - if (!vertices_it) { + if (!py_vertices_it) { PyErr_SetString(PyExc_MemoryError, "Unable to allocate _mgp.VerticesIterator."); return nullptr; @@ -627,6 +688,56 @@ PyObject *PyVertexLabelAt(PyVertex *self, PyObject *args) { return PyUnicode_FromString(label.name); } +PyObject *PyVertexIterInEdges(PyVertex *self, PyObject *Py_UNUSED(ignored)) { + CHECK(self); + CHECK(self->vertex); + CHECK(self->py_graph); + CHECK(self->py_graph->graph); + auto *edges_it = + mgp_vertex_iter_in_edges(self->vertex, self->py_graph->memory); + if (!edges_it) { + PyErr_SetString(PyExc_MemoryError, + "Unable to allocate mgp_edges_iterator for in edges."); + return nullptr; + } + auto *py_edges_it = PyObject_New(PyEdgesIterator, &PyEdgesIteratorType); + if (!py_edges_it) { + PyErr_SetString(PyExc_MemoryError, + "Unable to allocate _mgp.EdgesIterator for in edges."); + return nullptr; + } + py_edges_it->it = edges_it; + Py_INCREF(self->py_graph); + py_edges_it->py_graph = self->py_graph; + return PyObject_Init(reinterpret_cast(py_edges_it), + &PyEdgesIteratorType); +} + +PyObject *PyVertexIterOutEdges(PyVertex *self, PyObject *Py_UNUSED(ignored)) { + CHECK(self); + CHECK(self->vertex); + CHECK(self->py_graph); + CHECK(self->py_graph->graph); + auto *edges_it = + mgp_vertex_iter_out_edges(self->vertex, self->py_graph->memory); + if (!edges_it) { + PyErr_SetString(PyExc_MemoryError, + "Unable to allocate mgp_edges_iterator for out edges."); + return nullptr; + } + auto *py_edges_it = PyObject_New(PyEdgesIterator, &PyEdgesIteratorType); + if (!py_edges_it) { + PyErr_SetString(PyExc_MemoryError, + "Unable to allocate _mgp.EdgesIterator for out edges."); + return nullptr; + } + py_edges_it->it = edges_it; + Py_INCREF(self->py_graph); + py_edges_it->py_graph = self->py_graph; + return PyObject_Init(reinterpret_cast(py_edges_it), + &PyEdgesIteratorType); +} + static PyMethodDef PyVertexMethods[] = { {"is_valid", reinterpret_cast(PyVertexIsValid), METH_NOARGS, "Return True if Vertex is in valid context and may be used."}, @@ -636,6 +747,10 @@ static PyMethodDef PyVertexMethods[] = { METH_NOARGS, "Return number of lables of a vertex."}, {"label_at", reinterpret_cast(PyVertexLabelAt), METH_VARARGS, "Return label of a vertex on a given index."}, + {"iter_in_edges", reinterpret_cast(PyVertexIterInEdges), + METH_NOARGS, "Return _mgp.EdgesIterator for in edges"}, + {"iter_out_edges", reinterpret_cast(PyVertexIterOutEdges), + METH_NOARGS, "Return _mgp.EdgesIterator for out edges"}, {nullptr}}; PyObject *PyVertexRichCompare(PyObject *self, PyObject *other, int op); @@ -696,6 +811,7 @@ PyObject *PyInitMgpModule() { }; if (!register_type(&PyVerticesIteratorType, "VerticesIterator")) return nullptr; + if (!register_type(&PyEdgesIteratorType, "EdgesIterator")) return nullptr; if (!register_type(&PyGraphType, "Graph")) return nullptr; if (!register_type(&PyEdgeType, "Edge")) return nullptr; if (!register_type(&PyQueryProcType, "Proc")) return nullptr;