Changeset View
Changeset View
Standalone View
Standalone View
swh/graph/client.py
# Copyright (C) 2019-2020 The Software Heritage developers | # Copyright (C) 2019-2020 The Software Heritage developers | ||||
# See the AUTHORS file at the top-level directory of this distribution | # See the AUTHORS file at the top-level directory of this distribution | ||||
# License: GNU General Public License version 3, or any later version | # License: GNU General Public License version 3, or any later version | ||||
# See top-level LICENSE file for more information | # See top-level LICENSE file for more information | ||||
import json | import json | ||||
from swh.core.api import RPCClient | from swh.core.api import RPCClient | ||||
class GraphAPIError(Exception): | class GraphAPIError(Exception): | ||||
"""Graph API Error""" | """Graph API Error""" | ||||
def __str__(self): | def __str__(self): | ||||
return "An unexpected error occurred in the Graph backend: {}".format(self.args) | return """An unexpected error occurred | ||||
in the Graph backend: {}""".format( | |||||
self.args | |||||
) | |||||
class RemoteGraphClient(RPCClient): | class RemoteGraphClient(RPCClient): | ||||
"""Client to the Software Heritage Graph.""" | """Client to the Software Heritage Graph.""" | ||||
def __init__(self, url, timeout=None): | def __init__(self, url, timeout=None): | ||||
super().__init__(api_exception=GraphAPIError, url=url, timeout=timeout) | super().__init__(api_exception=GraphAPIError, url=url, timeout=timeout) | ||||
def raw_verb_lines(self, verb, endpoint, **kwargs): | def raw_verb_lines(self, verb, endpoint, **kwargs): | ||||
response = self.raw_verb(verb, endpoint, stream=True, **kwargs) | response = self.raw_verb(verb, endpoint, stream=True, **kwargs) | ||||
self.raise_for_status(response) | self.raise_for_status(response) | ||||
for line in response.iter_lines(): | for line in response.iter_lines(): | ||||
yield line.decode().lstrip("\n") | yield line.decode().lstrip("\n") | ||||
def get_lines(self, endpoint, **kwargs): | def get_lines(self, endpoint, **kwargs): | ||||
yield from self.raw_verb_lines("get", endpoint, **kwargs) | yield from self.raw_verb_lines("get", endpoint, **kwargs) | ||||
# Web API endpoints | # Web API endpoints | ||||
def stats(self): | def stats(self): | ||||
return self.get("stats") | return self.get("stats") | ||||
def leaves(self, src, edges="*", direction="forward", max_edges=0): | def leaves( | ||||
self, src, edges="*", direction="forward", max_edges=0, return_types="*" | |||||
): | |||||
return self.get_lines( | return self.get_lines( | ||||
"leaves/{}".format(src), | "leaves/{}".format(src), | ||||
params={"edges": edges, "direction": direction, "max_edges": max_edges}, | params={ | ||||
"edges": edges, | |||||
"direction": direction, | |||||
"max_edges": max_edges, | |||||
"return_types": return_types, | |||||
}, | |||||
) | ) | ||||
def neighbors(self, src, edges="*", direction="forward", max_edges=0): | def neighbors( | ||||
self, src, edges="*", direction="forward", max_edges=0, return_types="*" | |||||
): | |||||
return self.get_lines( | return self.get_lines( | ||||
"neighbors/{}".format(src), | "neighbors/{}".format(src), | ||||
params={"edges": edges, "direction": direction, "max_edges": max_edges}, | params={ | ||||
"edges": edges, | |||||
"direction": direction, | |||||
"max_edges": max_edges, | |||||
"return_types": return_types, | |||||
}, | |||||
) | ) | ||||
def visit_nodes(self, src, edges="*", direction="forward", max_edges=0): | def visit_nodes( | ||||
self, src, edges="*", direction="forward", max_edges=0, return_types="*" | |||||
): | |||||
return self.get_lines( | return self.get_lines( | ||||
"visit/nodes/{}".format(src), | "visit/nodes/{}".format(src), | ||||
params={"edges": edges, "direction": direction, "max_edges": max_edges}, | params={ | ||||
"edges": edges, | |||||
"direction": direction, | |||||
"max_edges": max_edges, | |||||
"return_types": return_types, | |||||
}, | |||||
) | ) | ||||
def visit_edges(self, src, edges="*", direction="forward", max_edges=0): | def visit_edges(self, src, edges="*", direction="forward", max_edges=0): | ||||
for edge in self.get_lines( | for edge in self.get_lines( | ||||
"visit/edges/{}".format(src), | "visit/edges/{}".format(src), | ||||
params={"edges": edges, "direction": direction, "max_edges": max_edges}, | params={"edges": edges, "direction": direction, "max_edges": max_edges}, | ||||
): | ): | ||||
yield tuple(edge.split()) | yield tuple(edge.split()) | ||||
Show All 19 Lines | ): | ||||
params={ | params={ | ||||
"edges": edges, | "edges": edges, | ||||
"traversal": traversal, | "traversal": traversal, | ||||
"direction": direction, | "direction": direction, | ||||
"limit": limit, | "limit": limit, | ||||
}, | }, | ||||
) | ) | ||||
def random_walk(self, src, dst, edges="*", direction="forward", limit=None): | def random_walk( | ||||
self, src, dst, edges="*", direction="forward", limit=None, return_types="*" | |||||
): | |||||
endpoint = "randomwalk/{}/{}" | endpoint = "randomwalk/{}/{}" | ||||
return self.get_lines( | return self.get_lines( | ||||
endpoint.format(src, dst), | endpoint.format(src, dst), | ||||
params={"edges": edges, "direction": direction, "limit": limit}, | params={ | ||||
"edges": edges, | |||||
"direction": direction, | |||||
"limit": limit, | |||||
"return_types": return_types, | |||||
}, | |||||
) | ) | ||||
def count_leaves(self, src, edges="*", direction="forward"): | def count_leaves(self, src, edges="*", direction="forward"): | ||||
return self.get( | return self.get( | ||||
"leaves/count/{}".format(src), | "leaves/count/{}".format(src), | ||||
params={"edges": edges, "direction": direction}, | params={"edges": edges, "direction": direction}, | ||||
) | ) | ||||
Show All 11 Lines |