query_modules/mgp_networkx.py
Rama master · réplica de lectura del repositorio original.
import collections
import sys
import mgp
try:
import networkx as nx
except ImportError as import_error:
sys.stderr.write(
(f"NOTE: Please install networkx to be able to" f"use graph_analyzer module. Using Python: {sys.version}")
)
raise import_error
class MemgraphAdjlistOuterDict(collections.abc.Mapping):
__slots__ = ("_ctx", "_succ", "_multi")
def __init__(self, ctx, succ=True, multi=True):
self._ctx = ctx
self._succ = succ
self._multi = multi
def __getitem__(self, key):
if key not in self:
raise KeyError
return MemgraphAdjlistInnerDict(key, succ=self._succ, multi=self._multi)
def __iter__(self):
return iter(self._ctx.graph.vertices)
def __len__(self):
return len(self._ctx.graph.vertices)
def __contains__(self, key):
if not isinstance(key, mgp.Vertex):
raise TypeError
return key in self._ctx.graph.vertices
class MemgraphAdjlistInnerDict(collections.abc.Mapping):
__slots__ = ("_node", "_succ", "_multi", "_neighbors")
def __init__(self, node, succ=True, multi=True):
self._node = node
self._succ = succ
self._multi = multi
self._neighbors = None
def __getitem__(self, key):
# NOTE: NetworkX 2.4, classes/coreviews.py:143. UnionAtlas expects a
# KeyError when indexing with a vertex that is not a neighbor.
if key not in self:
raise KeyError
if not self._multi:
return UnhashableProperties(self._get_edge(key).properties)
return MemgraphEdgeKeyDict(self._node, key, self._succ)
def __iter__(self):
yield from self._get_neighbors()
def __len__(self):
return len(self._get_neighbors())
def __contains__(self, key):
if not isinstance(key, mgp.Vertex):
raise TypeError
return key in self._get_neighbors()
def _get_neighbors(self):
if not self._neighbors:
if self._succ:
self._neighbors = set(e.to_vertex for e in self._node.out_edges)
else:
self._neighbors = set(e.from_vertex for e in self._node.in_edges)
return self._neighbors
def _get_edge(self, neighbor):
if self._succ:
edge = list(filter(lambda e: e.to_vertex == neighbor, self._node.out_edges))
else:
edge = list(filter(lambda e: e.from_vertex == neighbor, self._node.in_edges))
assert len(edge) >= 1
if len(edge) > 1:
raise RuntimeError("Graph contains multiedges but " "is of non-multigraph type: {}".format(edge))
return edge[0]
class MemgraphEdgeKeyDict(collections.abc.Mapping):
__slots__ = ("_node", "_neighbor", "_succ", "_edges")
def __init__(self, node, neighbor, succ=True):
self._node = node
self._neighbor = neighbor
self._succ = succ
self._edges = None
def __getitem__(self, key):
if key not in self:
raise KeyError
return UnhashableProperties(key.properties)
def __iter__(self):
yield from self._get_edges()
def __len__(self):
return len(self._get_edges())
def __contains__(self, key):
if not isinstance(key, mgp.Edge):
raise TypeError
return key in self._get_edges()
def _get_edges(self):
if not self._edges:
if self._succ:
self._edges = list(filter(lambda e: e.to_vertex == self._neighbor, self._node.out_edges))
else:
self._edges = list(filter(lambda e: e.from_vertex == self._neighbor, self._node.in_edges))
return self._edges
class UnhashableProperties(collections.abc.Mapping):
__slots__ = "_properties"
def __init__(self, properties):
self._properties = properties
def __getitem__(self, key):
return self._properties[key]
def __iter__(self):
yield from self._properties
def __len__(self):
return len(self._properties)
def __contains__(self, key):
return key in self._properties
# NOTE: Explicitly disable hashing. See the comment in MemgraphNodeDict.
__hash__ = None
class MemgraphNodeDict(collections.abc.Mapping):
__slots__ = ("_ctx",)
def __init__(self, ctx):
self._ctx = ctx
def __getitem__(self, key):
if key not in self:
raise KeyError
# NOTE: NetworkX 2.4, classes/digraph.py:484. NetworkX expects the
# tuples provided to add_nodes_from to be unhashable and cause a
# TypeError when trying to index the node dictionary. This happens
# because the data dictionary element of the tuple is unhashable. We do
# the same thing by returning an unhashable data dictionary.
return UnhashableProperties(key.properties)
def __iter__(self):
return iter(self._ctx.graph.vertices)
def __len__(self):
return len(self._ctx.graph.vertices)
def __contains__(self, key):
# NOTE: NetworkX 2.4, graph.py:425. Graph.__contains__ relies on
# self._node's (i.e. the dictionary produced by node_dict_factory)
# __contains__ to raise a TypeError when indexing with something weird,
# e.g. with sets. This is the behavior of dict.
if not isinstance(key, mgp.Vertex):
raise TypeError
return key in self._ctx.graph.vertices
class MemgraphDiGraphBase:
def __init__(self, incoming_graph_data=None, ctx=None, multi=True, **kwargs):
# NOTE: We assume that our graph will never be given any initial data
# because we already pull our data from the Memgraph database. This
# assert is triggered by certain NetworkX procedures because they
# create a new instance of our class and try to populate it with their
# own data.
assert incoming_graph_data is None
# NOTE: We allow for ctx to be None in order to allow certain NetworkX
# procedures (such as `subgraph`) to work. Such procedures directly
# modify the graph's internal attributes and don't try to populate it
# with initial data or modify it.
self.node_dict_factory = lambda: MemgraphNodeDict(ctx) if ctx else self._error
self.node_attr_dict_factory = self._error
self.adjlist_outer_dict_factory = lambda: MemgraphAdjlistOuterDict(ctx, multi=multi) if ctx else self._error
self.adjlist_inner_dict_factory = self._error
self.edge_key_dict_factory = self._error
self.edge_attr_dict_factory = self._error
# NOTE: We forbid any mutating operations because our graph is
# immutable and pulls its data from the Memgraph database.
for f in [
"add_node",
"add_nodes_from",
"remove_node",
"remove_nodes_from",
"add_edge",
"add_edges_from",
"add_weighted_edges_from",
"new_edge_key",
"remove_edge",
"remove_edges_from",
"update",
"clear",
]:
setattr(self, f, lambda *args, **kwargs: self._error())
super().__init__(None, **kwargs)
# NOTE: This is a necessary hack because NetworkX assumes that the
# customizable factory functions will only ever return *empty*
# dictionaries. In our case, the factory functions return our custom,
# already populated, dictionaries. Because self._pred and self._end are
# initialized by the same factory function, they end up storing the
# same adjacency lists which is not good. We correct that here.
self._pred = MemgraphAdjlistOuterDict(ctx, succ=False, multi=multi)
def _error(self):
raise RuntimeError("Modification operations are not supported")
class MemgraphMultiDiGraph(MemgraphDiGraphBase, nx.MultiDiGraph):
def __init__(self, incoming_graph_data=None, ctx=None, **kwargs):
super().__init__(incoming_graph_data=incoming_graph_data, ctx=ctx, multi=True, **kwargs)
def MemgraphMultiGraph(incoming_graph_data=None, ctx=None, **kwargs):
return MemgraphMultiDiGraph(incoming_graph_data=incoming_graph_data, ctx=ctx, **kwargs).to_undirected(as_view=True)
class MemgraphDiGraph(MemgraphDiGraphBase, nx.DiGraph):
def __init__(self, incoming_graph_data=None, ctx=None, **kwargs):
super().__init__(incoming_graph_data=incoming_graph_data, ctx=ctx, multi=False, **kwargs)
def MemgraphGraph(incoming_graph_data=None, ctx=None, **kwargs):
return MemgraphDiGraph(incoming_graph_data=incoming_graph_data, ctx=ctx, **kwargs).to_undirected(as_view=True)
class PropertiesDictionary(collections.abc.Mapping):
__slots__ = ("_ctx", "_prop", "_len")
def __init__(self, ctx, prop):
self._ctx = ctx
self._prop = prop
self._len = None
def __getitem__(self, vertex):
if vertex not in self:
raise KeyError
try:
return vertex.properties[self._prop]
except KeyError:
raise KeyError(("{} doesn\t have the required " + "property '{}'").format(vertex, self._prop))
def __iter__(self):
for v in self._ctx.graph.vertices:
if self._prop in v.properties:
yield v
def __len__(self):
if not self._len:
self._len = sum(1 for _ in self)
return self._len
def __contains__(self, vertex):
if not isinstance(vertex, mgp.Vertex):
raise TypeError
return self._prop in vertex.properties