graphs.py

from collections import deque

class Edge(object):
    def __init__(self, _id, data):
        self.id = _id
        self.data = data

    def __str__(self):
        return "Edge {}".format(self.id)

class DirectedEdge(Edge):
    def __init__(self, _id, data, f, t):
        super(DirectedEdge, self).__init__(_id, data)
        self.f = f
        self.t = t

    def __str__(self):
        return "Edge id: {} -> from {} to {}".format(self.id, self.f, self.t)

class Vertex(object):
    def __init__(self, _id, data):
        self.id = _id
        self.data = data
        self.edges = []
        self.visited = False

    def __str__(self):
        return "V {}".format(self.id)

    def add_from(self, v):
        pass

    def __eq__(self, other):
        return self.id == other.id

class Graph(object):
    def __init__(self):
        self.vertexes = []
        self.edges = []

    def get_vertex(self, _id):
        vertex = list(filter(lambda x: x.id == _id, self.vertexes))
        if not vertex:
            return None
        if len(vertex) > 1:
            raise Exception("Two vertexes with same ID {}".format(_id))
        return vertex[0]

    def add_edge(self, f_id, t_id, data=None, directed=True):
        f = self.get_vertex(f_id)
        if not f:
            f = Vertex(f_id, data)
        t = self.get_vertex(t_id)
        if not t:
            t = Vertex(t_id, data)
        edge = None
        e_id = (f.id, t.id)
        if directed:
            edge = DirectedEdge(e_id, data, f, t)
            if not f in self.vertexes:
                self.vertexes.append(f)
            if not t in self.vertexes:
                self.vertexes.append(t)

            if not t in f.edges:
                f.edges.append(t)
            self.edges.append(edge)
        else:
            raise Exception("Use directed edges only")

def clear_visited(G):
    for v in G.vertexes:
        v.visited = False

def BFS(G, s=None):
    if not s:
        s = G.vertexes[0]
    s.visited = True
    queue = deque([])
    queue.append(s)
    while queue:
        v = queue.popleft()
        print("Visited {}".format(v))
        for e in v.edges:
            if not e.visited:
                e.visited = True
                queue.append(e)

    clear_visited(G)

def DFS(G, s=None):
    if not s:
        s = G.vertexes[0]
    s.visited = True
    stack = []
    stack.append(s)
    while stack:
        v = stack.pop()
        print("Visited {}".format(v))
        for e in v.edges:
            if not e.visited:
                e.visited = True
                stack.append(e)

G = Graph()
G.add_edge(1, 2)
G.add_edge(1, 3)
G.add_edge(2, 4)
G.add_edge(3, 5)
G.add_edge(4, 6)
G.add_edge(5, 6)

print("Edges")
for edge in G.edges:
    print(edge)

print("Vertexes")
for vertex in G.vertexes:
    print(vertex)

print("Edges by vertex")
for vertex in G.vertexes:
    for edge in vertex.edges:
        print("Vertex {} -> {}".format(vertex.id, edge.id))

print("BFS")
BFS(G)

print("DSF")
DFS(G)