Build a Titan graph for testing the behavior of "CONTAINS" in Gremlin.

build_graph.py

__author__ = "zxk"

from bulbs.titan import Graph, Config, TITAN_URI
from bulbs.model import Node, Relationship
from bulbs.property import String, Integer
from bulbs.config import DEBUG

import math
import random
from termcolor import colored

class TestNode(Node):
    content = String(nullable = False, indexed = True)
    count = Integer(nullable = False, indexed = False)

KEYWORD = "hello"
OUTLIERS = ["world", "kugou", "moto", "python", "titan"]
LENGTH = 100
SIZE = 100

rand = lambda(x) : random.randint(0, x - 1)

class TestClient(object):
    def __init__(self, host, port, db, username = None, passwd = None):
        self.url_base = "http://%s:%s/graphs/%s" % (host, port, db)
        self.username = username
        self.passwd = passwd
        self.config = Config(self.url_base, self.username, self.passwd)
        self.config.autoindex = True
        self.g = Graph(self.config)
        self.g.add_proxy("test_node", TestNode)
        self.init_index()

    def init_index(self):
        self.g.test_node.index.create_key("content")

    def add_test_node(self, _count, _content):
        self.g.test_node.create(count = _count, content = _content)

    def generate_node(self, x):
        tmp_list = [1] * x + [0] * (LENGTH - x)
        random.shuffle(tmp_list)
        fin_list = []
        for each in tmp_list:
            if each:
                fin_list.append(KEYWORD)
            else:
                fin_list.append(OUTLIERS[rand(len(OUTLIERS))])
        fin_str = ' '.join(fin_list)
        self.add_test_node(x, fin_str)

    def build_graph(self):
        graph_size = SIZE
        for g_it in range(graph_size):
            self.generate_node(rand(min(SIZE, LENGTH)))
        print colored("Graph ready", "green")

if __name__ == "__main__":
    t = TestClient("0.0.0.0", 8182, "t1")
    t.build_graph()