def age_to_networkx()

in drivers/python/age/networkx/age_to_networkx.py [0:0]


def age_to_networkx(connection: psycopg.connect,
                    graphName: str,
                    G: None | nx.DiGraph = None,
                    query: str | None = None
                    ) -> nx.DiGraph:
    """
    @params
    ---------------------
    connection - (psycopg.connect) Connection object
    graphName - (str) Name of the graph
    G - (networkx.DiGraph) Networkx directed Graph [optional]
    query - (str) Cypher query [optional]

        @returns
    ------------
    Networkx directed Graph

    """

    # Check if the age graph exists
    checkIfGraphNameExistInAGE(connection, graphName)

    # Create an empty directed graph
    if G == None:
        G = nx.DiGraph()

    def addNodeToNetworkx(node):
        """Add Nodes in Networkx"""
        G.add_node(node.id,
                   label=node.label,
                   properties=node.properties)

    def addEdgeToNetworkx(edge):
        """Add Edge in Networkx"""
        G.add_edge(edge.start_id,
                   edge.end_id,
                   label=edge.label,
                   properties=edge.properties)

    def addPath(path):
        """Add Edge in Networkx"""
        for x in path:
            if (type(x) == Path):
                addPath(x)
        for x in path:
            if (type(x) == Vertex):
                addNodeToNetworkx(x)
        for x in path:
            if (type(x) == Edge):
                addEdgeToNetworkx(x)

    # Setting up connection to work with Graph
    age.setUpAge(connection, graphName)

    if (query == None):
        addAllNodesIntoNetworkx(connection, graphName, G)
        addAllEdgesIntoNetworkx(connection, graphName, G)
    else:
        with connection.cursor() as cursor:
            cursor.execute(query)
            rows = cursor.fetchall()
            for row in rows:
                for x in row:
                    if type(x) == Path:
                        addPath(x)
                    elif type(x) == Edge:
                        addEdgeToNetworkx(x)
                    elif type(x) == Vertex:
                        addNodeToNetworkx(x)
    return G