Python Graph implented by Adjacency List

graph_adj_list.py

""" Graph Data Structure: Adjacency List implementation
    that matches the Abstract Data Type as defined in the eBook
    https://runestone.academy/runestone/static/pythonds/index.html
    and goes with the video course at http://youtube.com/gjenkinslbcc channel
    This code is attributed to the book
    "Problem Solving with Algorithms and Data Structures using Python"
    by the Authors Bradley N. Miller and David Ranum
    under the Creative Commons
    Attribution-NonCommercial-ShareAlike 4.0 International License (CC-BY-NC-SA)
    License link: https://goo.gl/m2m1ww
    with some slight modifications
    """


class Vertex:
    def __init__(self, key, payload=None):
        self.id = key
        self.connectedTo = {}  # key is Vertex object, value is cost
        self.payload = payload

    def addNeighbor(self, vertex_object, cost=None):
        self.connectedTo[vertex_object] = cost

    def __str__(self):
        return str(self.id) + ' connectedTo: ' + str([x.id for x in self.connectedTo])

    def getConnections(self):
        return self.connectedTo.keys()

    def getId(self):
        return self.id

    def getWeight(self, vertex_obj):
        return self.connectedTo[vertex_obj]


class Graph:
    def __init__(self):
        self.vertList = {}
        self.numVertices = 0

    def addVertex(self, key):
        self.numVertices = self.numVertices + 1
        newVertex = Vertex(key)
        self.vertList[key] = newVertex
        return newVertex

    def getVertex(self, id):
        if id in self.vertList:
            return self.vertList[id]
        else:
            return None

    def __contains__(self, id):
        return id in self.vertList

    def addEdge(self, from_id, to_id, cost=None):
        if from_id not in self.vertList:
            nv = self.addVertex(from_id)
        if to_id not in self.vertList:
            nv = self.addVertex(to_id)
        self.vertList[from_id].addNeighbor(self.vertList[to_id], cost)

    def getVertices(self):
        """ return: dict_keys( of Vertex Objects) """
        return self.vertList.keys()

    def __iter__(self):
        return iter(self.vertList.values())
    
if __name__ == '__main__':
    # test case figure 2 @
    # https://runestone.academy/runestone/static/pythonds/Graphs/VocabularyandDefinitions.html#fig-dgsimple

    g = Graph()
    g.addEdge('v0', 'v1', 5)
    g.addEdge('v1', 'v2', 4)
    g.addEdge('v2', 'v3', 9)
    g.addEdge('v3', 'v4', 7)
    g.addEdge('v4', 'v0', 1)
    g.addEdge('v0', 'v5', 2)
    g.addEdge('v5', 'v4', 8)
    g.addEdge('v3', 'v5', 3)
    g.addEdge('v5', 'v2', 1)
    g.addVertex('v6')  # extra vertex with no connections

    print("The Adjacency List for this Graph:")
    vertex_keys = g.getVertices()
    for vid in vertex_keys:
        vertex = g.getVertex(vid)
        row = [ v.getId() for v in vertex.getConnections() ]

        print(f"{vertex.getId()}: {list(row)}")

    print('\ng.getConnections("v5")', [v.getId() for v in g.getVertex('v5').getConnections()])
    print('"v5" in g', "v5" in g)
    print('"v6" in g', "v6" in g)
    print('"v7" in g', "v7" in g)