class Graph(interfaces.IGraph):
"""
In-memory graph database.
See :class:`~.IGraph` for doco.
"""
def __init__(self):
self._vclass = Vertex
self._eclass = Edge
self._id_tracker = IDGenerator()
self._vconstraints = defaultdict(dict)
self._econstraints = defaultdict()
self.vertices = EntitySet()
self.edges = EntitySet()
def load(self, file_handler):
vertex_id_mapping = {}
data = json.load(file_handler)
constraints = data.get("constraints", [])
for constraint_dict in constraints:
self.add_vertex_constraint(
constraint_dict["label"],
constraint_dict["key"],
)
vertices = sorted(data.get("vertices", []), key=lambda x: x["id"])
for vertex_dict in vertices:
vertex = self.get_or_create_vertex(
vertex_dict["label"],
**vertex_dict["properties"]
)
vertex_id_mapping[vertex_dict["id"]] = vertex
edges = sorted(data.get("edges", []), key=lambda x: x["id"])
for edge_dict in edges:
head = vertex_id_mapping[edge_dict["head_id"]]
tail = vertex_id_mapping[edge_dict["tail_id"]]
self.get_or_create_edge(
head,
edge_dict["label"],
tail,
**edge_dict["properties"]
)
def dump(self, file_handler):
data = {
"vertices": [],
"edges": [],
"constraints": [],
}
for vertex in self.vertices:
data["vertices"].append(vertex.as_dict())
for edge in self.edges:
data["edges"].append(edge.as_dict())
for label, key in self.get_vertex_constraints():
data["constraints"].append(
{
"label": label,
"key": key
}
)
json.dump(data, file_handler, indent=4, sort_keys=True)
def add_vertex_constraint(self, label, key):
self._vconstraints[label][key] = set()
def get_vertex_constraints(self):
constraints = []
for label in self._vconstraints:
for key in self._vconstraints[label]:
constraints.append((label, key))
return constraints
def bind_to_graph(self, entity):
if isinstance(entity, interfaces.IVertex):
entity.ident = self._id_tracker.get_vertex_id()
elif isinstance(entity, interfaces.IEdge):
entity.ident = self._id_tracker.get_edge_id()
else:
raise interfaces.UnknownEntityError(
"Unknown entity {0!r}".format(entity)
)
entity.graph = self
def get_or_create_vertex(self, label=None, **kwargs):
if not label and not kwargs:
return None
# first check constraints.
if label in self._vconstraints:
for key, collection in self._vconstraints[label].items():
if key not in kwargs:
continue
for vertex in collection:
if vertex.properties[key] == kwargs[key]:
return vertex
# no matches in constraints, so do a EntitySet filter
vertices = self.vertices.filter(label, **kwargs)
if len(vertices) > 1:
raise interfaces.MultipleFoundExpectedOne(
"Multiple vertices found when one expected."
)
elif len(vertices) == 1:
return vertices.all()[0]
return self.add_vertex(label, **kwargs)
def get_or_create_edge(self, head, label, tail, **kwargs):
if isinstance(head, tuple):
head = self.get_or_create_vertex(head[0], **head[1])
if isinstance(tail, tuple):
tail = self.get_or_create_vertex(tail[0], **tail[1])
# There can only a single edge between head and tail with a
# particular label. So there is not point filtering for
# properties.
indexed_edge = self._econstraints.get((head, label, tail))
if indexed_edge:
return indexed_edge
return self.add_edge(head, label, tail, **kwargs)
def append_edge(self, edge):
head = edge.head
tail = edge.tail
self.append_vertex(head)
self.append_vertex(tail)
if edge.graph is not None and edge.graph is not self:
raise interfaces.DatabaseException(
"Can not append edge {} which is already bound to "
"anther graph instance.".format(edge)
)
if edge in self:
return edge
if edge.ident is not None:
raise interfaces.EntityIDError(
"Edge {} already has it identity number set.".format(edge)
)
self._edge_constraint_violated(edge)
self._econstraints[(head, edge.label, tail)] = edge
self.bind_to_graph(edge)
self.edges.add(edge)
head.out_edges.add(edge)
tail.in_edges.add(edge)
return edge
def append_vertex(self, vertex):
if vertex.graph is not None and vertex.graph is not self:
raise interfaces.DatabaseException(
"Can not append vertex {} which is already bound to "
"anther graph instance.".format(vertex)
)
if vertex in self:
return vertex
if vertex.ident is not None:
raise interfaces.EntityIDError(
"Vertex {} already has it identity number set.".format(vertex)
)
self._vertex_constraint_violated(vertex)
if vertex.label in self._vconstraints:
for key in self._vconstraints[vertex.label]:
if key in vertex.properties:
self._vconstraints[vertex.label][key].add(vertex)
self.bind_to_graph(vertex)
self.vertices.add(vertex)
return vertex
def add_edge(self, head, label, tail, **kwargs):
edge = self._eclass(head, label, tail, **kwargs)
return self.append_edge(edge)
def add_vertex(self, label=None, **kwargs):
vertex = self._vclass(label=label, **kwargs)
return self.append_vertex(vertex)
def _vertex_constraint_violated(self, vertex, **kwargs):
"""
Check if the given vertex violates any of the constraints.
:param vertex: vertex that you are checking for constraint violations.
:type vertex: :class:`~.IVertex`
:param kwargs: Additional properties.
:type kwargs: :class:`dict`
:raises ConstraintViolation: Raised if you a constraint violation has
been found.
"""
key_index = self._vconstraints.get(vertex.label, {})
# first check the entity properties for constraint violations
# Then check any additional properties for constraint violations.
# Additional properties are for cases like `.set_property`
for props in [vertex.properties, kwargs]:
for key, value in props.items():
if key not in key_index:
continue
for indexed_entity in key_index[key]:
if indexed_entity != vertex:
if indexed_entity.properties[key] == value:
raise interfaces.ConstraintViolation(
"{!r} violated constraint {!r}".format(
vertex, key
)
)
# todo: add in property constraint violation checks for edges
def _edge_constraint_violated(self, edge):
"""
Check if the given edge violates any of the constraints.
:param edge: Edge that you are checking for constraint violations.
:type edge: :class:`~.IEdge`
:raises ConstraintViolation: Raised if you a constraint violation has
been found.
"""
if (edge.head, edge.label, edge.tail) in self._econstraints:
raise interfaces.ConstraintViolation(
"Duplicate {0!r} edges between head {1!r} and tail {2!r} "
"is not allowed".format(
edge.label,
edge.head,
edge.tail,
)
)
def set_property(self, entity, **kwargs):
if entity not in self:
raise interfaces.UnknownEntityError(
"Unknown entity {0!r}".format(entity)
)
if isinstance(entity, interfaces.IVertex):
self._vertex_constraint_violated(entity, **kwargs)
self.vertices.update_index(entity, **kwargs)
if isinstance(entity, interfaces.IEdge):
# edge property constraints are not supported at this stage.
# todo: enable once edge property constraints are added.
# self._edge_constraint_violated(entity)
self.edges.update_index(entity, **kwargs)
entity._update_properties(kwargs) # pylint: disable=protected-access
def get_edge(self, id_num):
return self.edges.get(id_num)
def get_vertex(self, id_num):
return self.vertices.get(id_num)
def get_edges(self, head=None, label=None, tail=None, **kwargs):
if head is None and tail is None:
return self.edges.filter(label, **kwargs)
container = EntitySet()
for edge in self.edges.filter(label, **kwargs):
if head and tail is None:
if edge.head == head:
container.add(edge)
elif tail and head is None:
if edge.tail == tail:
container.add(edge)
else:
if edge.head == head and edge.tail == tail:
container.add(edge)
return container
def get_vertices(self, label=None, **kwargs):
return self.vertices.filter(label, **kwargs)
def remove_edge(self, edge):
edge.head.remove_edge(edge)
edge.tail.remove_edge(edge)
self.edges.remove(edge)
def remove_vertex(self, vertex):
if len(vertex.get_both_edges()) > 0:
raise interfaces.VertexBoundByEdges(
"Vertex {0!r} is still bound to another vertex "
"by an edge. First remove all the edges on the vertex and "
"then remove it again.".format(vertex)
)
self.vertices.remove(vertex)
def close(self): # pragma: no cover
# Nothing to do for the close at this stage.
return
def __contains__(self, entity):
is_vertex = isinstance(entity, interfaces.IVertex)
is_edge = isinstance(entity, interfaces.IEdge)
if not is_vertex and not is_edge:
raise TypeError(
"Unsupported entity type {0}".format(type(entity))
)
return entity in self.vertices or entity in self.edges
class TestEntitySet(base.TestBase):
def setUp(self):
super(TestEntitySet, self).setUp()
self.container = EntitySet([self.marko, self.josh, self.peter])
def test_add(self):
self.container.add(self.lop)
self.assertEqual(
self.container.sorted(key=lambda x: x.ident),
sorted([self.marko, self.josh, self.peter, self.lop], key=lambda x: x.ident),
)
def test_add_dup_id(self):
sue = Vertex("person", name="dup_vertex_id")
sue.ident = 0
self.assertRaises(KeyError, self.container.add, sue)
def test_add_same_vertex(self):
self.container.add(self.marko)
self.assertEqual(self.container.sorted(), sorted([self.marko, self.josh, self.peter]))
def test_get_labels(self):
self.container.add(self.lop)
self.assertEqual(sorted(self.container.get_labels()), sorted(["person", "app"]))
def test_get_indexes(self):
self.assertEqual(sorted(self.container.get_indexes()), sorted([("person", "age"), ("person", "name")]))
def test_update_index(self):
self.container.update_index(self.marko, surname="Foo")
self.assertEqual(
sorted(self.container.get_indexes()), sorted([("person", "age"), ("person", "name"), ("person", "surname")])
)
def test_remove(self):
self.container.remove(self.marko)
self.assertEqual(self.container.sorted(), sorted([self.josh, self.peter]))
def test_remove_with_unindexed_property(self):
self.marko.properties["unindexed"] = "say what"
self.container.remove(self.marko)
self.assertEqual(self.container.sorted(), sorted([self.josh, self.peter]))
def test_remove_unknown_entity_id(self):
sue = Vertex(100, name="sue")
self.assertRaises(KeyError, self.container.remove, sue)
def test_contains(self):
self.assertIn(self.marko, self.container)
def test_iter(self):
self.assertEqual(self.container.sorted(), sorted([self.marko, self.josh, self.peter]))
def test_get(self):
self.assertEqual(self.container.get(0), self.marko)
def test_get_no_such_entity_with_id(self):
self.assertRaises(KeyError, self.container.get, 100)
def test_len(self):
self.assertEqual(3, len(self.container))
def test_all(self):
self.assertEqual(
sorted(self.container.all(), key=lambda x: x.ident),
sorted([self.marko, self.josh, self.peter], key=lambda x: x.ident),
)
def test_sorted(self):
self.assertEqual(self.container.sorted(), sorted([self.marko, self.josh, self.peter]))
def test_sorted_cmp_key(self):
self.assertEqual(
self.container.sorted(key=lambda x: x.properties["name"]),
sorted([self.marko, self.josh, self.peter], key=lambda x: x.properties["name"]),
)
class Graph(interfaces.IGraph):
"""
In-memory graph database.
See :class:`~.IGraph` for doco.
"""
def __init__(self):
self._vclass = Vertex
self._eclass = Edge
self._id_tracker = IDGenerator()
self._vconstraints = defaultdict(dict)
self._econstraints = defaultdict()
self.vertices = EntitySet()
self.edges = EntitySet()
def load(self, file_handler):
vertex_id_mapping = {}
data = json.load(file_handler)
constraints = data.get("constraints", [])
for constraint_dict in constraints:
self.add_vertex_constraint(
constraint_dict["label"],
constraint_dict["key"],
)
vertices = sorted(data.get("vertices", []), key=lambda x: x["id"])
for vertex_dict in vertices:
vertex = self.get_or_create_vertex(vertex_dict["label"],
**vertex_dict["properties"])
vertex_id_mapping[vertex_dict["id"]] = vertex
edges = sorted(data.get("edges", []), key=lambda x: x["id"])
for edge_dict in edges:
head = vertex_id_mapping[edge_dict["head_id"]]
tail = vertex_id_mapping[edge_dict["tail_id"]]
self.get_or_create_edge(head, edge_dict["label"], tail,
**edge_dict["properties"])
def dump(self, file_handler):
data = {
"vertices": [],
"edges": [],
"constraints": [],
}
for vertex in self.vertices:
data["vertices"].append(vertex.as_dict())
for edge in self.edges:
data["edges"].append(edge.as_dict())
for label, key in self.get_vertex_constraints():
data["constraints"].append({"label": label, "key": key})
json.dump(data, file_handler, indent=4, sort_keys=True)
def add_vertex_constraint(self, label, key):
self._vconstraints[label][key] = set()
def get_vertex_constraints(self):
constraints = []
for label in self._vconstraints:
for key in self._vconstraints[label]:
constraints.append((label, key))
return constraints
def bind_to_graph(self, entity):
if isinstance(entity, interfaces.IVertex):
entity.ident = self._id_tracker.get_vertex_id()
elif isinstance(entity, interfaces.IEdge):
entity.ident = self._id_tracker.get_edge_id()
else:
raise interfaces.UnknownEntityError(
"Unknown entity {0!r}".format(entity))
entity.graph = self
def get_or_create_vertex(self, label=None, **kwargs):
if not label and not kwargs:
return None
# first check constraints.
if label in self._vconstraints:
for key, collection in self._vconstraints[label].items():
if key not in kwargs:
continue
for vertex in collection:
if vertex.properties[key] == kwargs[key]:
return vertex
# no matches in constraints, so do a EntitySet filter
vertices = self.vertices.filter(label, **kwargs)
if len(vertices) > 1:
raise interfaces.MultipleFoundExpectedOne(
"Multiple vertices found when one expected.")
elif len(vertices) == 1:
return vertices.all()[0]
return self.add_vertex(label, **kwargs)
def get_or_create_edge(self, head, label, tail, **kwargs):
if isinstance(head, tuple):
head = self.get_or_create_vertex(head[0], **head[1])
if isinstance(tail, tuple):
tail = self.get_or_create_vertex(tail[0], **tail[1])
# There can only a single edge between head and tail with a
# particular label. So there is not point filtering for
# properties.
indexed_edge = self._econstraints.get((head, label, tail))
if indexed_edge:
return indexed_edge
return self.add_edge(head, label, tail, **kwargs)
def append_edge(self, edge):
head = edge.head
tail = edge.tail
self.append_vertex(head)
self.append_vertex(tail)
if edge.graph is not None and edge.graph is not self:
raise interfaces.DatabaseException(
"Can not append edge {} which is already bound to "
"anther graph instance.".format(edge))
if edge in self:
return edge
if edge.ident is not None:
raise interfaces.EntityIDError(
"Edge {} already has it identity number set.".format(edge))
self._edge_constraint_violated(edge)
self._econstraints[(head, edge.label, tail)] = edge
self.bind_to_graph(edge)
self.edges.add(edge)
head.out_edges.add(edge)
tail.in_edges.add(edge)
return edge
def append_vertex(self, vertex):
if vertex.graph is not None and vertex.graph is not self:
raise interfaces.DatabaseException(
"Can not append vertex {} which is already bound to "
"anther graph instance.".format(vertex))
if vertex in self:
return vertex
if vertex.ident is not None:
raise interfaces.EntityIDError(
"Vertex {} already has it identity number set.".format(vertex))
self._vertex_constraint_violated(vertex)
if vertex.label in self._vconstraints:
for key in self._vconstraints[vertex.label]:
if key in vertex.properties:
self._vconstraints[vertex.label][key].add(vertex)
self.bind_to_graph(vertex)
self.vertices.add(vertex)
return vertex
def add_edge(self, head, label, tail, **kwargs):
edge = self._eclass(head, label, tail, **kwargs)
return self.append_edge(edge)
def add_vertex(self, label=None, **kwargs):
vertex = self._vclass(label=label, **kwargs)
return self.append_vertex(vertex)
def _vertex_constraint_violated(self, vertex, **kwargs):
"""
Check if the given vertex violates any of the constraints.
:param vertex: vertex that you are checking for constraint violations.
:type vertex: :class:`~.IVertex`
:param kwargs: Additional properties.
:type kwargs: :class:`dict`
:raises ConstraintViolation: Raised if you a constraint violation has
been found.
"""
key_index = self._vconstraints.get(vertex.label, {})
# first check the entity properties for constraint violations
# Then check any additional properties for constraint violations.
# Additional properties are for cases like `.set_property`
for props in [vertex.properties, kwargs]:
for key, value in props.items():
if key not in key_index:
continue
for indexed_entity in key_index[key]:
if indexed_entity != vertex:
if indexed_entity.properties[key] == value:
raise interfaces.ConstraintViolation(
"{!r} violated constraint {!r}".format(
vertex, key))
# todo: add in property constraint violation checks for edges
def _edge_constraint_violated(self, edge):
"""
Check if the given edge violates any of the constraints.
:param edge: Edge that you are checking for constraint violations.
:type edge: :class:`~.IEdge`
:raises ConstraintViolation: Raised if you a constraint violation has
been found.
"""
if (edge.head, edge.label, edge.tail) in self._econstraints:
raise interfaces.ConstraintViolation(
"Duplicate {0!r} edges between head {1!r} and tail {2!r} "
"is not allowed".format(
edge.label,
edge.head,
edge.tail,
))
def set_property(self, entity, **kwargs):
if entity not in self:
raise interfaces.UnknownEntityError(
"Unknown entity {0!r}".format(entity))
if isinstance(entity, interfaces.IVertex):
self._vertex_constraint_violated(entity, **kwargs)
self.vertices.update_index(entity, **kwargs)
if isinstance(entity, interfaces.IEdge):
# edge property constraints are not supported at this stage.
# todo: enable once edge property constraints are added.
# self._edge_constraint_violated(entity)
self.edges.update_index(entity, **kwargs)
entity._update_properties(kwargs) # pylint: disable=protected-access
def get_edge(self, id_num):
return self.edges.get(id_num)
def get_vertex(self, id_num):
return self.vertices.get(id_num)
def get_edges(self, head=None, label=None, tail=None, **kwargs):
if head is None and tail is None:
return self.edges.filter(label, **kwargs)
container = EntitySet()
for edge in self.edges.filter(label, **kwargs):
if head and tail is None:
if edge.head == head:
container.add(edge)
elif tail and head is None:
if edge.tail == tail:
container.add(edge)
else:
if edge.head == head and edge.tail == tail:
container.add(edge)
return container
def get_vertices(self, label=None, **kwargs):
return self.vertices.filter(label, **kwargs)
def remove_edge(self, edge):
edge.head.remove_edge(edge)
edge.tail.remove_edge(edge)
self.edges.remove(edge)
def remove_vertex(self, vertex):
if len(vertex.get_both_edges()) > 0:
raise interfaces.VertexBoundByEdges(
"Vertex {0!r} is still bound to another vertex "
"by an edge. First remove all the edges on the vertex and "
"then remove it again.".format(vertex))
self.vertices.remove(vertex)
def close(self): # pragma: no cover
# Nothing to do for the close at this stage.
return
def __contains__(self, entity):
is_vertex = isinstance(entity, interfaces.IVertex)
is_edge = isinstance(entity, interfaces.IEdge)
if not is_vertex and not is_edge:
raise TypeError("Unsupported entity type {0}".format(type(entity)))
return entity in self.vertices or entity in self.edges
class TestEntitySet(base.TestBase):
def setUp(self):
super(TestEntitySet, self).setUp()
self.container = EntitySet([self.marko, self.josh, self.peter])
def test_add(self):
self.container.add(self.lop)
self.assertEqual(
self.container.sorted(key=lambda x: x.ident),
sorted(
[self.marko, self.josh, self.peter, self.lop],
key=lambda x: x.ident,
))
def test_add_dup_id(self):
sue = Vertex("person", name="dup_vertex_id")
sue.ident = 0
self.assertRaises(KeyError, self.container.add, sue)
def test_add_same_vertex(self):
self.container.add(self.marko)
self.assertEqual(
self.container.sorted(),
sorted([self.marko, self.josh, self.peter]),
)
def test_get_labels(self):
self.container.add(self.lop)
self.assertEqual(
sorted(self.container.get_labels()),
sorted(["person", "app"]),
)
def test_get_indexes(self):
self.assertEqual(
sorted(self.container.get_indexes()),
sorted([
("person", "age"),
("person", "name"),
]),
)
def test_update_index(self):
self.container.update_index(self.marko, surname="Foo")
self.assertEqual(
sorted(self.container.get_indexes()),
sorted([
("person", "age"),
("person", "name"),
("person", "surname"),
]),
)
def test_remove(self):
self.container.remove(self.marko)
self.assertEqual(self.container.sorted(),
sorted([self.josh, self.peter]))
def test_remove_with_unindexed_property(self):
self.marko.properties["unindexed"] = "say what"
self.container.remove(self.marko)
self.assertEqual(self.container.sorted(),
sorted([self.josh, self.peter]))
def test_remove_unknown_entity_id(self):
sue = Vertex(100, name="sue")
self.assertRaises(
KeyError,
self.container.remove,
sue,
)
def test_contains(self):
self.assertIn(self.marko, self.container)
def test_iter(self):
self.assertEqual(self.container.sorted(),
sorted([self.marko, self.josh, self.peter]))
def test_get(self):
self.assertEqual(self.container.get(0), self.marko)
def test_get_no_such_entity_with_id(self):
self.assertRaises(
KeyError,
self.container.get,
100,
)
def test_len(self):
self.assertEqual(3, len(self.container))
def test_all(self):
self.assertEqual(
sorted(self.container.all(), key=lambda x: x.ident),
sorted([self.marko, self.josh, self.peter], key=lambda x: x.ident))
def test_sorted(self):
self.assertEqual(self.container.sorted(),
sorted([self.marko, self.josh, self.peter]))
def test_sorted_cmp_key(self):
self.assertEqual(
self.container.sorted(key=lambda x: x.properties["name"]),
sorted(
[self.marko, self.josh, self.peter],
key=lambda x: x.properties["name"],
))