Ejemplo n.º 1
0
 def test_filter_for_labels_with_no_props(self):
     marko = Vertex("person", name="Marko")
     josh = Vertex("person", name="Marko")
     edge = Edge(marko, "father", josh)
     container = EntitySet()
     container.add(edge)
     self.assertEqual(container.filter("father").sorted(), sorted([edge]))
Ejemplo n.º 2
0
 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()
Ejemplo n.º 3
0
 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()
Ejemplo n.º 4
0
 def test_filter_for_labels_with_no_props(self):
     marko = Vertex(
         "person",
         name="Marko",
     )
     josh = Vertex(
         "person",
         name="Marko",
     )
     edge = Edge(marko, "father", josh)
     container = EntitySet()
     container.add(edge)
     self.assertEqual(container.filter("father").sorted(), sorted([edge]))
Ejemplo n.º 5
0
    def test_union_filter_for_labels_with_no_props(self):
        id_generator = IDGenerator()
        marko = Vertex(
            "person",
            name="Marko",
        )
        josh = Vertex(
            "person",
            name="Marko",
        )
        edge = Edge(marko, "father", josh)
        edge2 = Edge(josh, "son", marko)

        # attach ID's to the vertices and edges.
        marko.ident = id_generator.get_vertex_id()
        josh.ident = id_generator.get_vertex_id()
        edge.ident = id_generator.get_edge_id()
        edge2.ident = id_generator.get_edge_id()

        container1 = EntitySet()
        container2 = EntitySet()

        container1.add(edge)
        container2.add(edge2)
        container3 = container1 | container2

        self.assertEqual(container3.filter("father").sorted(), sorted([edge]))
        self.assertEqual(container3.filter("son").sorted(), sorted([edge2]))
Ejemplo n.º 6
0
    def setUp(self):
        id_generator = IDGenerator()
        self.marko = Vertex(
            "Father",
            name="Marko",
            surname="Jones",
            age=30,
        )
        self.john = Vertex(
            "Brother",
            name="John",
            surname="Jones",
            age=30
        )
        self.peter = Vertex(
            "Uncle",
            name="Peter",
            surname="Doe",
            age=20
        )

        # add id's to the unbound vertices
        self.marko.ident = id_generator.get_vertex_id()
        self.john.ident = id_generator.get_vertex_id()
        self.peter.ident = id_generator.get_vertex_id()
        self.container = EntitySet([self.marko, self.john, self.peter])
Ejemplo n.º 7
0
 def test_intersection(self):
     sue = Vertex("Sister", name="Sue")
     intersection = self.container & EntitySet([self.marko, sue])
     self.assertEqual(
         intersection.sorted(),  # pylint: disable=no-member
         sorted([
             self.marko,
         ]))
Ejemplo n.º 8
0
 def test_union_with_duplicate_id(self):
     union = self.container | EntitySet([self.marko])
     self.assertEqual(
         union.sorted(),  # pylint: disable=no-member
         sorted([
             self.marko,
             self.john,
             self.peter,
         ]))
Ejemplo n.º 9
0
 def test_difference(self):
     sue = Vertex("Sister", name="Sue")
     diff = self.container - EntitySet([self.marko, sue])
     self.assertEqual(
         diff.sorted(),  # pylint: disable=no-member
         sorted([
             self.john,
             self.peter,
         ]))
Ejemplo n.º 10
0
    def test_union_filter_for_labels_with_no_props(self):
        id_generator = IDGenerator()
        marko = Vertex(
            "person",
            name="Marko",
        )
        josh = Vertex(
            "person",
            name="Marko",
        )
        edge = Edge(marko, "father", josh)
        edge2 = Edge(josh, "son", marko)

        # attach ID's to the vertices and edges.
        marko.ident = id_generator.get_vertex_id()
        josh.ident = id_generator.get_vertex_id()
        edge.ident = id_generator.get_edge_id()
        edge2.ident = id_generator.get_edge_id()

        container1 = EntitySet()
        container2 = EntitySet()

        container1.add(edge)
        container2.add(edge2)
        container3 = container1 | container2

        self.assertEqual(
            container3.filter("father").sorted(),
            sorted([edge])
        )
        self.assertEqual(
            container3.filter("son").sorted(),
            sorted([edge2])
        )
Ejemplo n.º 11
0
    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
Ejemplo n.º 12
0
 def test_union(self):
     sue = Vertex(
         "Sister",
         name="Sue"
     )
     union = self.container | EntitySet([sue, self.marko])
     self.assertEqual(
         union.sorted(),  # pylint: disable=no-member
         sorted(
             [
                 self.marko,
                 self.john,
                 self.peter,
                 sue
             ]
         )
     )
Ejemplo n.º 13
0
    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
Ejemplo n.º 14
0
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
Ejemplo n.º 15
0
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"],
            ))
Ejemplo n.º 16
0
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
Ejemplo n.º 17
0
 def setUp(self):
     super(TestEntitySet, self).setUp()
     self.container = EntitySet([self.marko, self.josh, self.peter])
Ejemplo n.º 18
0
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"]),
        )
Ejemplo n.º 19
0
 def setUp(self):
     super(TestEntitySet, self).setUp()
     self.container = EntitySet([self.marko, self.josh, self.peter])