visualiser.py

import sys

import networkx as nx
import matplotlib.pyplot as plt


from arango import create
from arango.aql import F, V


def edges(db, key, graph, depth=2, cur_depth=0, counter=None):
    if cur_depth > depth:
        return

    query = db.tweets_edge.query.over(
        F.EDGES("tweets_edges", ~V(key), ~V("outbound")))
    for n, obj in enumerate(query.execute()):
        if obj["_to"] in counter:
            continue
        graph.add_node("{}/{}".format(cur_depth, n))
        graph.add_edge(key, obj["_to"])
        counter.add(obj["_to"])
        edges(
            db, obj["_to"], graph, depth=depth,
            cur_depth=cur_depth + 1, counter=counter)


def visualize(db, name, depth=1):
    print("Visualizing {}".format(name))
    graph = nx.Graph()
    query = db.tweets_edge.query.over(
        F.EDGES(
            "tweets_edges",
            ~V("tweets/196297127"), ~V("outbound")))

    counter = set()
    for n, q in enumerate(query.execute()):
        graph.add_node(q["_to"])
        counter.add(q["_to"])
        edges(db, q["_to"], graph, depth=depth, counter=counter)

    nx.draw(
        graph,
        edge_color='g',
        node_color='y',
        width=0.4,
        with_labels=False,
        alpha=0.5,
        node_size=25)
    plt.savefig("out/{}.png".format(name))

if __name__ == "__main__":
    name = sys.argv[1]
    depth = int(sys.argv[2])
    arango = create(db="tweets_{}".format(name))
    visualize(arango, name, depth=depth)