from google.protobuf.field_mask_pb2 import FieldMask
from swh.graph.grpc.swhgraph_pb2 import (
    NodeFilter,
    StatsRequest,
    TraversalRequest,
)
from swh.graph.grpc.swhgraph_pb2_grpc import TraversalServiceStub


channel = grpc.aio.insecure_channel('localhost:50091')
with channel:
    client = TraversalServiceStub(channel)

    # Stats request
    print(client.Stats(StatsRequest()))

    # Traversal request
    request = TraversalRequest(
        src=["swh:1:cnt:cafedeadbeef"],
        edges="rev:rev,rev:rel",
        direction="BACKWARD",
        return_nodes=NodeFilter("rel"),
        mask=FieldMask(paths=["swhid", "rel.message", "rel.author"]),
    )
    for node in client.Traverse(request):
        print(node)