extenting_gremlin_traversal.py

from gremlin_python.driver.driver_remote_connection import DriverRemoteConnection
from gremlin_python.process.anonymous_traversal import traversal
from gremlin_python.process.traversal import Order


class GraphConnection:

    def __init__(self, gremlin_url):
        self.connection = DriverRemoteConnection(
            gremlin_url
        )
        self.g = traversal().withRemote(self.connection)


class TraversalGraph:

    def __init__(self, connection):
        self.connection = connection
        self.query = self.connection.g

    def get_graph_traversal(self):
        return self.query

    def toList(self):
        return self.query.toList()

    def next(self):
        return self.query.next()

    def values(self, *args):
        self.query = self.query.elementMap(*args)
        return self

    def V(self):
        self.query = self.query.V()
        return self

    def E(self):
        self.query = self.query.E()
        return self

    def find(self, **kwargs):
        self.query = self.query.hasLabel(kwargs['has__label']).limit(kwargs['pagination__limit'])
        return self

    def count(self):
        return self.query.count()

    def order_by(self, property_name, order):
        self.query = self.query.order().by(property_name, order)
        return self

    def outE(self, **kwargs):
        self.query = self.query.outE(**kwargs)
        return self

    def inV(self, **kwargs):
        self.query = self.query.inV(**kwargs)
        return self

    def inE(self, **kwargs):
        self.query = self.query.inE(**kwargs)
        return self

    def outV(self, **kwargs):
        self.query = self.query.outV(**kwargs)
        return self


graph_connection = GraphConnection("ws://megamind-ws:8182/gremlin")
graph_traversal = TraversalGraph(graph_connection)
c = graph_traversal.V().find(has__label="TestLabel", pagination__limit=4) \
    .order_by("name", Order.asc) \
    .values("name").toList()
print("count", c)
graph_connection.connection.close()