diff --git a/swh/graph/backend.py b/swh/graph/backend.py
index 22d4036..de54810 100644
--- a/swh/graph/backend.py
+++ b/swh/graph/backend.py
@@ -1,206 +1,206 @@
# Copyright (C) 2019-2020 The Software Heritage developers
# See the AUTHORS file at the top-level directory of this distribution
# License: GNU General Public License version 3, or any later version
# See top-level LICENSE file for more information
import asyncio
import contextlib
import io
import os
import struct
import subprocess
import sys
import tempfile
from py4j.java_gateway import JavaGateway
from swh.graph.config import check_config
from swh.graph.swhid import NodeToSwhidMap, SwhidToNodeMap
-from swh.model.identifiers import EXTENDED_SWHID_TYPES
+from swh.model.swhids import EXTENDED_SWHID_TYPES
BUF_SIZE = 64 * 1024
BIN_FMT = ">q" # 64 bit integer, big endian
PATH_SEPARATOR_ID = -1
NODE2SWHID_EXT = "node2swhid.bin"
SWHID2NODE_EXT = "swhid2node.bin"
def _get_pipe_stderr():
# Get stderr if possible, or pipe to stdout if running with Jupyter.
try:
sys.stderr.fileno()
except io.UnsupportedOperation:
return subprocess.STDOUT
else:
return sys.stderr
class Backend:
def __init__(self, graph_path, config=None):
self.gateway = None
self.entry = None
self.graph_path = graph_path
self.config = check_config(config or {})
def start_gateway(self):
self.gateway = JavaGateway.launch_gateway(
java_path=None,
javaopts=self.config["java_tool_options"].split(),
classpath=self.config["classpath"],
die_on_exit=True,
redirect_stdout=sys.stdout,
redirect_stderr=_get_pipe_stderr(),
)
self.entry = self.gateway.jvm.org.softwareheritage.graph.Entry()
self.entry.load_graph(self.graph_path)
self.node2swhid = NodeToSwhidMap(self.graph_path + "." + NODE2SWHID_EXT)
self.swhid2node = SwhidToNodeMap(self.graph_path + "." + SWHID2NODE_EXT)
self.stream_proxy = JavaStreamProxy(self.entry)
def stop_gateway(self):
self.gateway.shutdown()
def __enter__(self):
self.start_gateway()
return self
def __exit__(self, exc_type, exc_value, tb):
self.stop_gateway()
def stats(self):
return self.entry.stats()
def count(self, ttype, direction, edges_fmt, src):
method = getattr(self.entry, "count_" + ttype)
return method(direction, edges_fmt, src)
async def simple_traversal(
self, ttype, direction, edges_fmt, src, max_edges, return_types
):
assert ttype in ("leaves", "neighbors", "visit_nodes")
method = getattr(self.stream_proxy, ttype)
async for node_id in method(direction, edges_fmt, src, max_edges, return_types):
yield node_id
async def walk(self, direction, edges_fmt, algo, src, dst):
if dst in EXTENDED_SWHID_TYPES:
it = self.stream_proxy.walk_type(direction, edges_fmt, algo, src, dst)
else:
it = self.stream_proxy.walk(direction, edges_fmt, algo, src, dst)
async for node_id in it:
yield node_id
async def random_walk(self, direction, edges_fmt, retries, src, dst, return_types):
if dst in EXTENDED_SWHID_TYPES:
it = self.stream_proxy.random_walk_type(
direction, edges_fmt, retries, src, dst, return_types
)
else:
it = self.stream_proxy.random_walk(
direction, edges_fmt, retries, src, dst, return_types
)
async for node_id in it: # TODO return 404 if path is empty
yield node_id
async def visit_edges(self, direction, edges_fmt, src, max_edges):
it = self.stream_proxy.visit_edges(direction, edges_fmt, src, max_edges)
# convert stream a, b, c, d -> (a, b), (c, d)
prevNode = None
async for node in it:
if prevNode is not None:
yield (prevNode, node)
prevNode = None
else:
prevNode = node
async def visit_paths(self, direction, edges_fmt, src, max_edges):
path = []
async for node in self.stream_proxy.visit_paths(
direction, edges_fmt, src, max_edges
):
if node == PATH_SEPARATOR_ID:
yield path
path = []
else:
path.append(node)
class JavaStreamProxy:
"""A proxy class for the org.softwareheritage.graph.Entry Java class that
takes care of the setup and teardown of the named-pipe FIFO communication
between Python and Java.
Initialize JavaStreamProxy using:
proxy = JavaStreamProxy(swh_entry_class_instance)
Then you can call an Entry method and iterate on the FIFO results like
this:
async for value in proxy.java_method(arg1, arg2):
print(value)
"""
def __init__(self, entry):
self.entry = entry
async def read_node_ids(self, fname):
loop = asyncio.get_event_loop()
open_thread = loop.run_in_executor(None, open, fname, "rb")
# Since the open() call on the FIFO is blocking until it is also opened
# on the Java side, we await it with a timeout in case there is an
# exception that prevents the write-side open().
with (await asyncio.wait_for(open_thread, timeout=2)) as f:
while True:
data = await loop.run_in_executor(None, f.read, BUF_SIZE)
if not data:
break
for data in struct.iter_unpack(BIN_FMT, data):
yield data[0]
class _HandlerWrapper:
def __init__(self, handler):
self._handler = handler
def __getattr__(self, name):
func = getattr(self._handler, name)
async def java_call(*args, **kwargs):
loop = asyncio.get_event_loop()
await loop.run_in_executor(None, lambda: func(*args, **kwargs))
def java_task(*args, **kwargs):
return asyncio.create_task(java_call(*args, **kwargs))
return java_task
@contextlib.contextmanager
def get_handler(self):
with tempfile.TemporaryDirectory(prefix="swh-graph-") as tmpdirname:
cli_fifo = os.path.join(tmpdirname, "swh-graph.fifo")
os.mkfifo(cli_fifo)
reader = self.read_node_ids(cli_fifo)
query_handler = self.entry.get_handler(cli_fifo)
handler = self._HandlerWrapper(query_handler)
yield (handler, reader)
def __getattr__(self, name):
async def java_call_iterator(*args, **kwargs):
with self.get_handler() as (handler, reader):
java_task = getattr(handler, name)(*args, **kwargs)
try:
async for value in reader:
yield value
except asyncio.TimeoutError:
# If the read-side open() timeouts, an exception on the
# Java side probably happened that prevented the
# write-side open(). We propagate this exception here if
# that is the case.
task_exc = java_task.exception()
if task_exc:
raise task_exc
raise
await java_task
return java_call_iterator
diff --git a/swh/graph/cli.py b/swh/graph/cli.py
index 3b224b3..8bd11d3 100644
--- a/swh/graph/cli.py
+++ b/swh/graph/cli.py
@@ -1,446 +1,446 @@
# Copyright (C) 2019-2020 The Software Heritage developers
# See the AUTHORS file at the top-level directory of this distribution
# License: GNU General Public License version 3, or any later version
# See top-level LICENSE file for more information
import logging
from pathlib import Path
import sys
from typing import TYPE_CHECKING, Any, Dict, Set, Tuple
# WARNING: do not import unnecessary things here to keep cli startup time under
# control
import click
from swh.core.cli import CONTEXT_SETTINGS, AliasedGroup
from swh.core.cli import swh as swh_cli_group
if TYPE_CHECKING:
from swh.graph.webgraph import CompressionStep # noqa
class StepOption(click.ParamType):
"""click type for specifying a compression step on the CLI
parse either individual steps, specified as step names or integers, or step
ranges
"""
name = "compression step"
def convert(self, value, param, ctx): # type: (...) -> Set[CompressionStep]
from swh.graph.webgraph import COMP_SEQ, CompressionStep # noqa
steps: Set[CompressionStep] = set()
specs = value.split(",")
for spec in specs:
if "-" in spec: # step range
(raw_l, raw_r) = spec.split("-", maxsplit=1)
if raw_l == "": # no left endpoint
raw_l = COMP_SEQ[0].name
if raw_r == "": # no right endpoint
raw_r = COMP_SEQ[-1].name
l_step = self.convert(raw_l, param, ctx)
r_step = self.convert(raw_r, param, ctx)
if len(l_step) != 1 or len(r_step) != 1:
self.fail(f"invalid step specification: {value}, " f"see --help")
l_idx = l_step.pop()
r_idx = r_step.pop()
steps = steps.union(
set(CompressionStep(i) for i in range(l_idx.value, r_idx.value + 1))
)
else: # singleton step
try:
steps.add(CompressionStep(int(spec))) # integer step
except ValueError:
try:
steps.add(CompressionStep[spec.upper()]) # step name
except KeyError:
self.fail(
f"invalid step specification: {value}, " f"see --help"
)
return steps
class PathlibPath(click.Path):
"""A Click path argument that returns a pathlib Path, not a string"""
def convert(self, value, param, ctx):
return Path(super().convert(value, param, ctx))
DEFAULT_CONFIG: Dict[str, Tuple[str, Any]] = {"graph": ("dict", {})}
@swh_cli_group.group(name="graph", context_settings=CONTEXT_SETTINGS, cls=AliasedGroup)
@click.option(
"--config-file",
"-C",
default=None,
type=click.Path(exists=True, dir_okay=False,),
help="YAML configuration file",
)
@click.pass_context
def graph_cli_group(ctx, config_file):
"""Software Heritage graph tools."""
from swh.core import config
ctx.ensure_object(dict)
conf = config.read(config_file, DEFAULT_CONFIG)
if "graph" not in conf:
raise ValueError(
'no "graph" stanza found in configuration file %s' % config_file
)
ctx.obj["config"] = conf
@graph_cli_group.command("api-client")
@click.option("--host", default="localhost", help="Graph server host")
@click.option("--port", default="5009", help="Graph server port")
@click.pass_context
def api_client(ctx, host, port):
"""client for the graph RPC service"""
from swh.graph import client
url = "http://{}:{}".format(host, port)
app = client.RemoteGraphClient(url)
# TODO: run web app
print(app.stats())
@graph_cli_group.group("map")
@click.pass_context
def map(ctx):
"""Manage swh-graph on-disk maps"""
pass
def dump_swhid2node(filename):
from swh.graph.swhid import SwhidToNodeMap
for (swhid, int) in SwhidToNodeMap(filename):
print("{}\t{}".format(swhid, int))
def dump_node2swhid(filename):
from swh.graph.swhid import NodeToSwhidMap
for (int, swhid) in NodeToSwhidMap(filename):
print("{}\t{}".format(int, swhid))
def restore_swhid2node(filename):
"""read a textual SWHID->int map from stdin and write its binary version to
filename
"""
from swh.graph.swhid import SwhidToNodeMap
with open(filename, "wb") as dst:
for line in sys.stdin:
(str_swhid, str_int) = line.split()
SwhidToNodeMap.write_record(dst, str_swhid, int(str_int))
def restore_node2swhid(filename, length):
"""read a textual int->SWHID map from stdin and write its binary version to
filename
"""
from swh.graph.swhid import NodeToSwhidMap
node2swhid = NodeToSwhidMap(filename, mode="wb", length=length)
for line in sys.stdin:
(str_int, str_swhid) = line.split()
node2swhid[int(str_int)] = str_swhid
node2swhid.close()
@map.command("dump")
@click.option(
"--type",
"-t",
"map_type",
required=True,
type=click.Choice(["swhid2node", "node2swhid"]),
help="type of map to dump",
)
@click.argument("filename", required=True, type=click.Path(exists=True))
@click.pass_context
def dump_map(ctx, map_type, filename):
"""Dump a binary SWHID<->node map to textual format."""
if map_type == "swhid2node":
dump_swhid2node(filename)
elif map_type == "node2swhid":
dump_node2swhid(filename)
else:
raise ValueError("invalid map type: " + map_type)
pass
@map.command("restore")
@click.option(
"--type",
"-t",
"map_type",
required=True,
type=click.Choice(["swhid2node", "node2swhid"]),
help="type of map to dump",
)
@click.option(
"--length",
"-l",
type=int,
help="""map size in number of logical records
(required for node2swhid maps)""",
)
@click.argument("filename", required=True, type=click.Path())
@click.pass_context
def restore_map(ctx, map_type, length, filename):
"""Restore a binary SWHID<->node map from textual format."""
if map_type == "swhid2node":
restore_swhid2node(filename)
elif map_type == "node2swhid":
if length is None:
raise click.UsageError(
"map length is required when restoring {} maps".format(map_type), ctx
)
restore_node2swhid(filename, length)
else:
raise ValueError("invalid map type: " + map_type)
@map.command("write")
@click.option(
"--type",
"-t",
"map_type",
required=True,
type=click.Choice(["swhid2node", "node2swhid"]),
help="type of map to write",
)
@click.argument("filename", required=True, type=click.Path())
@click.pass_context
def write(ctx, map_type, filename):
"""Write a map to disk sequentially.
read from stdin a textual SWHID->node mapping (for swhid2node, or a simple
sequence of SWHIDs for node2swhid) and write it to disk in the requested binary
map format
note that no sorting is applied, so the input should already be sorted as
required by the chosen map type (by SWHID for swhid2node, by int for node2swhid)
"""
from swh.graph.swhid import NodeToSwhidMap, SwhidToNodeMap
with open(filename, "wb") as f:
if map_type == "swhid2node":
for line in sys.stdin:
(swhid, int_str) = line.rstrip().split(maxsplit=1)
SwhidToNodeMap.write_record(f, swhid, int(int_str))
elif map_type == "node2swhid":
for line in sys.stdin:
swhid = line.rstrip()
NodeToSwhidMap.write_record(f, swhid)
else:
raise ValueError("invalid map type: " + map_type)
@map.command("lookup")
@click.option(
"--graph", "-g", required=True, metavar="GRAPH", help="compressed graph basename"
)
@click.argument("identifiers", nargs=-1)
def map_lookup(graph, identifiers):
"""Lookup identifiers using on-disk maps.
Depending on the identifier type lookup either a SWHID into a SWHID->node (and
return the node integer identifier) or, vice-versa, lookup a node integer
identifier into a node->SWHID (and return the SWHID). The desired behavior is
chosen depending on the syntax of each given identifier.
Identifiers can be passed either directly on the command line or on
standard input, separate by blanks. Logical lines (as returned by
readline()) in stdin will be preserved in stdout.
"""
from swh.graph.backend import NODE2SWHID_EXT, SWHID2NODE_EXT
from swh.graph.swhid import NodeToSwhidMap, SwhidToNodeMap
import swh.model.exceptions
- from swh.model.identifiers import ExtendedSWHID
+ from swh.model.swhids import ExtendedSWHID
success = True # no identifiers failed to be looked up
swhid2node = SwhidToNodeMap(f"{graph}.{SWHID2NODE_EXT}")
node2swhid = NodeToSwhidMap(f"{graph}.{NODE2SWHID_EXT}")
def lookup(identifier):
nonlocal success, swhid2node, node2swhid
is_swhid = None
try:
int(identifier)
is_swhid = False
except ValueError:
try:
ExtendedSWHID.from_string(identifier)
is_swhid = True
except swh.model.exceptions.ValidationError:
success = False
logging.error(f'invalid identifier: "{identifier}", skipping')
try:
if is_swhid:
return str(swhid2node[identifier])
else:
return node2swhid[int(identifier)]
except KeyError:
success = False
logging.error(f'identifier not found: "{identifier}", skipping')
if identifiers: # lookup identifiers passed via CLI
for identifier in identifiers:
print(lookup(identifier))
else: # lookup identifiers passed via stdin, preserving logical lines
for line in sys.stdin:
results = [lookup(id) for id in line.rstrip().split()]
if results: # might be empty if all IDs on the same line failed
print(" ".join(results))
sys.exit(0 if success else 1)
@graph_cli_group.command(name="rpc-serve")
@click.option(
"--host",
"-h",
default="0.0.0.0",
metavar="IP",
show_default=True,
help="host IP address to bind the server on",
)
@click.option(
"--port",
"-p",
default=5009,
type=click.INT,
metavar="PORT",
show_default=True,
help="port to bind the server on",
)
@click.option(
"--graph", "-g", required=True, metavar="GRAPH", help="compressed graph basename"
)
@click.pass_context
def serve(ctx, host, port, graph):
"""run the graph RPC service"""
import aiohttp
from swh.graph.server.app import make_app
config = ctx.obj["config"]
config.setdefault("graph", {})
config["graph"]["path"] = graph
app = make_app(config=config)
aiohttp.web.run_app(app, host=host, port=port)
@graph_cli_group.command()
@click.option(
"--graph",
"-g",
required=True,
metavar="GRAPH",
type=PathlibPath(),
help="input graph basename",
)
@click.option(
"--outdir",
"-o",
"out_dir",
required=True,
metavar="DIR",
type=PathlibPath(),
help="directory where to store compressed graph",
)
@click.option(
"--steps",
"-s",
metavar="STEPS",
type=StepOption(),
help="run only these compression steps (default: all steps)",
)
@click.pass_context
def compress(ctx, graph, out_dir, steps):
"""Compress a graph using WebGraph
Input: a pair of files g.nodes.csv.gz, g.edges.csv.gz
Output: a directory containing a WebGraph compressed graph
Compression steps are: (1) mph, (2) bv, (3) bv_obl, (4) bfs, (5) permute,
(6) permute_obl, (7) stats, (8) transpose, (9) transpose_obl, (10) maps,
(11) clean_tmp. Compression steps can be selected by name or number using
--steps, separating them with commas; step ranges (e.g., 3-9, 6-, etc.) are
also supported.
"""
from swh.graph import webgraph
graph_name = graph.name
in_dir = graph.parent
try:
conf = ctx.obj["config"]["graph"]["compress"]
except KeyError:
conf = {} # use defaults
webgraph.compress(graph_name, in_dir, out_dir, steps, conf)
@graph_cli_group.command(name="cachemount")
@click.option(
"--graph", "-g", required=True, metavar="GRAPH", help="compressed graph basename"
)
@click.option(
"--cache",
"-c",
default="/dev/shm/swh-graph/default",
metavar="CACHE",
type=PathlibPath(),
help="Memory cache path (defaults to /dev/shm/swh-graph/default)",
)
@click.pass_context
def cachemount(ctx, graph, cache):
"""
Cache the mmapped files of the compressed graph in a tmpfs.
This command creates a new directory at the path given by CACHE that has
the same structure as the compressed graph basename, except it copies the
files that require mmap access (:file:`{*}.graph`) but uses symlinks from the source
for all the other files (:file:`{*}.map`, :file:`{*}.bin`, ...).
The command outputs the path to the memory cache directory (particularly
useful when relying on the default value).
"""
import shutil
cache.mkdir(parents=True)
for src in Path(graph).parent.glob("*"):
dst = cache / src.name
if src.suffix == ".graph":
shutil.copy2(src, dst)
else:
dst.symlink_to(src.resolve())
print(cache)
def main():
return graph_cli_group(auto_envvar_prefix="SWH_GRAPH")
if __name__ == "__main__":
main()
diff --git a/swh/graph/naive_client.py b/swh/graph/naive_client.py
index 8191311..9e08d65 100644
--- a/swh/graph/naive_client.py
+++ b/swh/graph/naive_client.py
@@ -1,369 +1,369 @@
# Copyright (C) 2021 The Software Heritage developers
# See the AUTHORS file at the top-level directory of this distribution
# License: GNU General Public License version 3, or any later version
# See top-level LICENSE file for more information
import functools
import inspect
import re
import statistics
from typing import (
Callable,
Dict,
Iterable,
Iterator,
List,
Optional,
Set,
Tuple,
TypeVar,
)
-from swh.model.identifiers import ExtendedSWHID, ValidationError
+from swh.model.swhids import ExtendedSWHID, ValidationError
from .client import GraphArgumentException
_NODE_TYPES = "ori|snp|rel|rev|dir|cnt"
NODES_RE = re.compile(fr"(\*|{_NODE_TYPES})")
EDGES_RE = re.compile(fr"(\*|{_NODE_TYPES}):(\*|{_NODE_TYPES})")
T = TypeVar("T", bound=Callable)
def check_arguments(f: T) -> T:
"""Decorator for generic argument checking for methods of NaiveClient.
Checks ``src`` is a valid and known SWHID, and ``edges`` has the right format."""
signature = inspect.signature(f)
@functools.wraps(f)
def newf(*args, **kwargs):
__tracebackhide__ = True # for pytest
try:
bound_args = signature.bind(*args, **kwargs)
except TypeError as e:
# rethrow the exception from here so pytest doesn't flood the terminal
# with signature.bind's call stack.
raise TypeError(*e.args) from None
self = bound_args.arguments["self"]
src = bound_args.arguments.get("src")
if src:
self._check_swhid(src)
edges = bound_args.arguments.get("edges")
if edges:
if edges != "*" and not EDGES_RE.match(edges):
raise GraphArgumentException(f"invalid edge restriction: {edges}")
return_types = bound_args.arguments.get("return_types")
if return_types:
if not NODES_RE.match(return_types):
raise GraphArgumentException(
f"invalid return_types restriction: {return_types}"
)
return f(*args, **kwargs)
return newf # type: ignore
def filter_node_types(node_types: str, nodes: Iterable[str]) -> Iterator[str]:
if node_types == "*":
yield from nodes
else:
prefixes = tuple(f"swh:1:{type_}:" for type_ in node_types.split(","))
for node in nodes:
if node.startswith(prefixes):
yield node
class NaiveClient:
"""An alternative implementation of :class:`swh.graph.backend.Backend`,
written in pure-python and meant for simulating it in other components' test
cases.
It is NOT meant to be efficient in any way; only to be a very simple
implementation that provides the same behavior."""
def __init__(self, *, nodes: List[str], edges: List[Tuple[str, str]]):
self.graph = Graph(nodes, edges)
def _check_swhid(self, swhid):
try:
ExtendedSWHID.from_string(swhid)
except ValidationError as e:
raise GraphArgumentException(*e.args) from None
if swhid not in self.graph.nodes:
raise GraphArgumentException(f"SWHID not found: {swhid}")
def stats(self) -> Dict:
return {
"counts": {
"nodes": len(self.graph.nodes),
"edges": sum(map(len, self.graph.forward_edges.values())),
},
"ratios": {
"compression": 1.0,
"bits_per_edge": 100.0,
"bits_per_node": 100.0,
"avg_locality": 0.0,
},
"indegree": {
"min": min(map(len, self.graph.backward_edges.values())),
"max": max(map(len, self.graph.backward_edges.values())),
"avg": statistics.mean(map(len, self.graph.backward_edges.values())),
},
"outdegree": {
"min": min(map(len, self.graph.forward_edges.values())),
"max": max(map(len, self.graph.forward_edges.values())),
"avg": statistics.mean(map(len, self.graph.forward_edges.values())),
},
}
@check_arguments
def leaves(
self,
src: str,
edges: str = "*",
direction: str = "forward",
max_edges: int = 0,
return_types: str = "*",
) -> Iterator[str]:
# TODO: max_edges
yield from filter_node_types(
return_types,
[
node
for node in self.graph.get_subgraph(src, edges, direction)
if not self.graph.get_filtered_neighbors(node, edges, direction)
],
)
@check_arguments
def neighbors(
self,
src: str,
edges: str = "*",
direction: str = "forward",
max_edges: int = 0,
return_types: str = "*",
) -> Iterator[str]:
# TODO: max_edges
yield from filter_node_types(
return_types, self.graph.get_filtered_neighbors(src, edges, direction)
)
@check_arguments
def visit_nodes(
self,
src: str,
edges: str = "*",
direction: str = "forward",
max_edges: int = 0,
return_types: str = "*",
) -> Iterator[str]:
# TODO: max_edges
yield from filter_node_types(
return_types, self.graph.get_subgraph(src, edges, direction)
)
@check_arguments
def visit_edges(
self, src: str, edges: str = "*", direction: str = "forward", max_edges: int = 0
) -> Iterator[Tuple[str, str]]:
if max_edges == 0:
max_edges = None # type: ignore
else:
max_edges -= 1
yield from list(self.graph.iter_edges_dfs(direction, edges, src))[:max_edges]
@check_arguments
def visit_paths(
self, src: str, edges: str = "*", direction: str = "forward", max_edges: int = 0
) -> Iterator[List[str]]:
# TODO: max_edges
for path in self.graph.iter_paths_dfs(direction, edges, src):
if path[-1] in self.leaves(src, edges, direction):
yield list(path)
@check_arguments
def walk(
self,
src: str,
dst: str,
edges: str = "*",
traversal: str = "dfs",
direction: str = "forward",
limit: Optional[int] = None,
) -> Iterator[str]:
# TODO: implement algo="bfs"
# TODO: limit
match_path: Callable[[str], bool]
if ":" in dst:
match_path = dst.__eq__
self._check_swhid(dst)
else:
match_path = lambda node: node.startswith(f"swh:1:{dst}:") # noqa
for path in self.graph.iter_paths_dfs(direction, edges, src):
if match_path(path[-1]):
if not limit:
# 0 or None
yield from path
elif limit > 0:
yield from path[0:limit]
else:
yield from path[limit:]
@check_arguments
def random_walk(
self,
src: str,
dst: str,
edges: str = "*",
direction: str = "forward",
limit: Optional[int] = None,
):
# TODO: limit
yield from self.walk(src, dst, edges, "dfs", direction, limit)
@check_arguments
def count_leaves(
self, src: str, edges: str = "*", direction: str = "forward"
) -> int:
return len(list(self.leaves(src, edges, direction)))
@check_arguments
def count_neighbors(
self, src: str, edges: str = "*", direction: str = "forward"
) -> int:
return len(self.graph.get_filtered_neighbors(src, edges, direction))
@check_arguments
def count_visit_nodes(
self, src: str, edges: str = "*", direction: str = "forward"
) -> int:
return len(self.graph.get_subgraph(src, edges, direction))
class Graph:
def __init__(self, nodes: List[str], edges: List[Tuple[str, str]]):
self.nodes = nodes
self.forward_edges: Dict[str, List[str]] = {}
self.backward_edges: Dict[str, List[str]] = {}
for node in nodes:
self.forward_edges[node] = []
self.backward_edges[node] = []
for (src, dst) in edges:
self.forward_edges[src].append(dst)
self.backward_edges[dst].append(src)
def get_filtered_neighbors(
self, src: str, edges_fmt: str, direction: str,
) -> Set[str]:
if direction == "forward":
edges = self.forward_edges
elif direction == "backward":
edges = self.backward_edges
else:
raise GraphArgumentException(f"invalid direction: {direction}")
neighbors = edges.get(src, [])
if edges_fmt == "*":
return set(neighbors)
else:
filtered_neighbors: Set[str] = set()
for edges_fmt_item in edges_fmt.split(","):
(src_fmt, dst_fmt) = edges_fmt_item.split(":")
if src_fmt != "*" and not src.startswith(f"swh:1:{src_fmt}:"):
continue
if dst_fmt == "*":
filtered_neighbors.update(neighbors)
else:
prefix = f"swh:1:{dst_fmt}:"
filtered_neighbors.update(
n for n in neighbors if n.startswith(prefix)
)
return filtered_neighbors
def get_subgraph(self, src: str, edges_fmt: str, direction: str) -> Set[str]:
seen = set()
to_visit = {src}
while to_visit:
node = to_visit.pop()
seen.add(node)
neighbors = set(self.get_filtered_neighbors(node, edges_fmt, direction))
new_nodes = neighbors - seen
to_visit.update(new_nodes)
return seen
def iter_paths_dfs(
self, direction: str, edges_fmt: str, src: str
) -> Iterator[Tuple[str, ...]]:
for (path, node) in DfsSubgraphIterator(self, direction, edges_fmt, src):
yield path + (node,)
def iter_edges_dfs(
self, direction: str, edges_fmt: str, src: str
) -> Iterator[Tuple[str, str]]:
for (path, node) in DfsSubgraphIterator(self, direction, edges_fmt, src):
if len(path) > 0:
yield (path[-1], node)
class SubgraphIterator(Iterator[Tuple[Tuple[str, ...], str]]):
def __init__(self, graph: Graph, direction: str, edges_fmt: str, src: str):
self.graph = graph
self.direction = direction
self.edges_fmt = edges_fmt
self.seen: Set[str] = set()
self.src = src
def more_work(self) -> bool:
raise NotImplementedError()
def pop(self) -> Tuple[Tuple[str, ...], str]:
raise NotImplementedError()
def push(self, new_path: Tuple[str, ...], neighbor: str) -> None:
raise NotImplementedError()
def __next__(self) -> Tuple[Tuple[str, ...], str]:
# Stores (path, next_node)
if not self.more_work():
raise StopIteration()
(path, node) = self.pop()
new_path = path + (node,)
if node not in self.seen:
neighbors = self.graph.get_filtered_neighbors(
node, self.edges_fmt, self.direction
)
# We want to visit the first neighbor first, and to_visit is a stack;
# so we need to reversed() the list of neighbors to get it on top
# of the stack.
for neighbor in reversed(list(neighbors)):
self.push(new_path, neighbor)
self.seen.add(node)
return (path, node)
class DfsSubgraphIterator(SubgraphIterator):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self.to_visit: List[Tuple[Tuple[str, ...], str]] = [((), self.src)]
def more_work(self) -> bool:
return bool(self.to_visit)
def pop(self) -> Tuple[Tuple[str, ...], str]:
return self.to_visit.pop()
def push(self, new_path: Tuple[str, ...], neighbor: str) -> None:
self.to_visit.append((new_path, neighbor))
diff --git a/swh/graph/server/app.py b/swh/graph/server/app.py
index 0816a4b..bd21952 100644
--- a/swh/graph/server/app.py
+++ b/swh/graph/server/app.py
@@ -1,402 +1,402 @@
# Copyright (C) 2019-2020 The Software Heritage developers
# See the AUTHORS file at the top-level directory of this distribution
# License: GNU General Public License version 3, or any later version
# See top-level LICENSE file for more information
"""
A proxy HTTP server for swh-graph, talking to the Java code via py4j, and using
FIFO as a transport to stream integers between the two languages.
"""
import asyncio
from collections import deque
import json
import os
from typing import Optional
import aiohttp.web
from swh.core.api.asynchronous import RPCServerApp
from swh.core.config import read as config_read
from swh.graph.backend import Backend
from swh.model.exceptions import ValidationError
-from swh.model.identifiers import EXTENDED_SWHID_TYPES
+from swh.model.swhids import EXTENDED_SWHID_TYPES
try:
from contextlib import asynccontextmanager
except ImportError:
# Compatibility with 3.6 backport
from async_generator import asynccontextmanager # type: ignore
# maximum number of retries for random walks
RANDOM_RETRIES = 5 # TODO make this configurable via rpc-serve configuration
class GraphServerApp(RPCServerApp):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self.on_startup.append(self._start_gateway)
self.on_shutdown.append(self._stop_gateway)
@staticmethod
async def _start_gateway(app):
# Equivalent to entering `with app["backend"]:`
app["backend"].start_gateway()
@staticmethod
async def _stop_gateway(app):
# Equivalent to exiting `with app["backend"]:` with no error
app["backend"].stop_gateway()
async def index(request):
return aiohttp.web.Response(
content_type="text/html",
body="""<html>
<head><title>Software Heritage graph server</title></head>
<body>
<p>You have reached the <a href="https://www.softwareheritage.org/">
Software Heritage</a> graph API server.</p>
<p>See its
<a href="https://docs.softwareheritage.org/devel/swh-graph/api.html">API
documentation</a> for more information.</p>
</body>
</html>""",
)
class GraphView(aiohttp.web.View):
"""Base class for views working on the graph, with utility functions"""
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self.backend = self.request.app["backend"]
def node_of_swhid(self, swhid):
"""Lookup a SWHID in a swhid2node map, failing in an HTTP-nice way if
needed."""
try:
return self.backend.swhid2node[swhid]
except KeyError:
raise aiohttp.web.HTTPNotFound(text=f"SWHID not found: {swhid}")
except ValidationError:
raise aiohttp.web.HTTPBadRequest(text=f"malformed SWHID: {swhid}")
def swhid_of_node(self, node):
"""Lookup a node in a node2swhid map, failing in an HTTP-nice way if
needed."""
try:
return self.backend.node2swhid[node]
except KeyError:
raise aiohttp.web.HTTPInternalServerError(
text=f"reverse lookup failed for node id: {node}"
)
def get_direction(self):
"""Validate HTTP query parameter `direction`"""
s = self.request.query.get("direction", "forward")
if s not in ("forward", "backward"):
raise aiohttp.web.HTTPBadRequest(text=f"invalid direction: {s}")
return s
def get_edges(self):
"""Validate HTTP query parameter `edges`, i.e., edge restrictions"""
s = self.request.query.get("edges", "*")
if any(
[
node_type != "*" and node_type not in EXTENDED_SWHID_TYPES
for edge in s.split(":")
for node_type in edge.split(",", maxsplit=1)
]
):
raise aiohttp.web.HTTPBadRequest(text=f"invalid edge restriction: {s}")
return s
def get_return_types(self):
"""Validate HTTP query parameter 'return types', i.e,
a set of types which we will filter the query results with"""
s = self.request.query.get("return_types", "*")
if any(
node_type != "*" and node_type not in EXTENDED_SWHID_TYPES
for node_type in s.split(",")
):
raise aiohttp.web.HTTPBadRequest(
text=f"invalid type for filtering res: {s}"
)
# if the user puts a star,
# then we filter nothing, we don't need the other information
if "*" in s:
return "*"
else:
return s
def get_traversal(self):
"""Validate HTTP query parameter `traversal`, i.e., visit order"""
s = self.request.query.get("traversal", "dfs")
if s not in ("bfs", "dfs"):
raise aiohttp.web.HTTPBadRequest(text=f"invalid traversal order: {s}")
return s
def get_limit(self):
"""Validate HTTP query parameter `limit`, i.e., number of results"""
s = self.request.query.get("limit", "0")
try:
return int(s)
except ValueError:
raise aiohttp.web.HTTPBadRequest(text=f"invalid limit value: {s}")
def get_max_edges(self):
"""Validate HTTP query parameter 'max_edges', i.e.,
the limit of the number of edges that can be visited"""
s = self.request.query.get("max_edges", "0")
try:
return int(s)
except ValueError:
raise aiohttp.web.HTTPBadRequest(text=f"invalid max_edges value: {s}")
class StreamingGraphView(GraphView):
"""Base class for views streaming their response line by line."""
content_type = "text/plain"
@asynccontextmanager
async def response_streamer(self, *args, **kwargs):
"""Context manager to prepare then close a StreamResponse"""
response = aiohttp.web.StreamResponse(*args, **kwargs)
response.content_type = self.content_type
await response.prepare(self.request)
yield response
await response.write_eof()
async def get(self):
await self.prepare_response()
async with self.response_streamer() as self.response_stream:
self._buf = []
try:
await self.stream_response()
finally:
await self._flush_buffer()
return self.response_stream
async def prepare_response(self):
"""This can be overridden with some setup to be run before the response
actually starts streaming.
"""
pass
async def stream_response(self):
"""Override this to perform the response streaming. Implementations of
this should await self.stream_line(line) to write each line.
"""
raise NotImplementedError
async def stream_line(self, line):
"""Write a line in the response stream."""
self._buf.append(line)
if len(self._buf) > 100:
await self._flush_buffer()
async def _flush_buffer(self):
await self.response_stream.write("\n".join(self._buf).encode() + b"\n")
self._buf = []
class StatsView(GraphView):
"""View showing some statistics on the graph"""
async def get(self):
stats = self.backend.stats()
return aiohttp.web.Response(body=stats, content_type="application/json")
class SimpleTraversalView(StreamingGraphView):
"""Base class for views of simple traversals"""
simple_traversal_type: Optional[str] = None
async def prepare_response(self):
src = self.request.match_info["src"]
self.src_node = self.node_of_swhid(src)
self.edges = self.get_edges()
self.direction = self.get_direction()
self.max_edges = self.get_max_edges()
self.return_types = self.get_return_types()
async def stream_response(self):
async for res_node in self.backend.simple_traversal(
self.simple_traversal_type,
self.direction,
self.edges,
self.src_node,
self.max_edges,
self.return_types,
):
res_swhid = self.swhid_of_node(res_node)
await self.stream_line(res_swhid)
class LeavesView(SimpleTraversalView):
simple_traversal_type = "leaves"
class NeighborsView(SimpleTraversalView):
simple_traversal_type = "neighbors"
class VisitNodesView(SimpleTraversalView):
simple_traversal_type = "visit_nodes"
class WalkView(StreamingGraphView):
async def prepare_response(self):
src = self.request.match_info["src"]
dst = self.request.match_info["dst"]
self.src_node = self.node_of_swhid(src)
if dst not in EXTENDED_SWHID_TYPES:
self.dst_thing = self.node_of_swhid(dst)
else:
self.dst_thing = dst
self.edges = self.get_edges()
self.direction = self.get_direction()
self.algo = self.get_traversal()
self.limit = self.get_limit()
self.return_types = self.get_return_types()
async def get_walk_iterator(self):
return self.backend.walk(
self.direction, self.edges, self.algo, self.src_node, self.dst_thing
)
async def stream_response(self):
it = self.get_walk_iterator()
if self.limit < 0:
queue = deque(maxlen=-self.limit)
async for res_node in it:
res_swhid = self.swhid_of_node(res_node)
queue.append(res_swhid)
while queue:
await self.stream_line(queue.popleft())
else:
count = 0
async for res_node in it:
if self.limit == 0 or count < self.limit:
res_swhid = self.swhid_of_node(res_node)
await self.stream_line(res_swhid)
count += 1
else:
break
class RandomWalkView(WalkView):
def get_walk_iterator(self):
return self.backend.random_walk(
self.direction,
self.edges,
RANDOM_RETRIES,
self.src_node,
self.dst_thing,
self.return_types,
)
class VisitEdgesView(SimpleTraversalView):
async def stream_response(self):
it = self.backend.visit_edges(
self.direction, self.edges, self.src_node, self.max_edges
)
async for (res_src, res_dst) in it:
res_src_swhid = self.swhid_of_node(res_src)
res_dst_swhid = self.swhid_of_node(res_dst)
await self.stream_line("{} {}".format(res_src_swhid, res_dst_swhid))
class VisitPathsView(SimpleTraversalView):
content_type = "application/x-ndjson"
async def stream_response(self):
it = self.backend.visit_paths(
self.direction, self.edges, self.src_node, self.max_edges
)
async for res_path in it:
res_path_swhid = [self.swhid_of_node(n) for n in res_path]
line = json.dumps(res_path_swhid)
await self.stream_line(line)
class CountView(GraphView):
"""Base class for counting views."""
count_type: Optional[str] = None
async def get(self):
src = self.request.match_info["src"]
self.src_node = self.node_of_swhid(src)
self.edges = self.get_edges()
self.direction = self.get_direction()
loop = asyncio.get_event_loop()
cnt = await loop.run_in_executor(
None,
self.backend.count,
self.count_type,
self.direction,
self.edges,
self.src_node,
)
return aiohttp.web.Response(body=str(cnt), content_type="application/json")
class CountNeighborsView(CountView):
count_type = "neighbors"
class CountLeavesView(CountView):
count_type = "leaves"
class CountVisitNodesView(CountView):
count_type = "visit_nodes"
def make_app(config=None, backend=None, **kwargs):
if (config is None) == (backend is None):
raise ValueError("make_app() expects exactly one of 'config' or 'backend'")
if backend is None:
backend = Backend(graph_path=config["graph"]["path"], config=config["graph"])
app = GraphServerApp(**kwargs)
app.add_routes(
[
aiohttp.web.get("/", index),
aiohttp.web.get("/graph", index),
aiohttp.web.view("/graph/stats", StatsView),
aiohttp.web.view("/graph/leaves/{src}", LeavesView),
aiohttp.web.view("/graph/neighbors/{src}", NeighborsView),
aiohttp.web.view("/graph/visit/nodes/{src}", VisitNodesView),
aiohttp.web.view("/graph/visit/edges/{src}", VisitEdgesView),
aiohttp.web.view("/graph/visit/paths/{src}", VisitPathsView),
# temporarily disabled in wait of a proper fix for T1969
# aiohttp.web.view("/graph/walk/{src}/{dst}", WalkView)
aiohttp.web.view("/graph/randomwalk/{src}/{dst}", RandomWalkView),
aiohttp.web.view("/graph/neighbors/count/{src}", CountNeighborsView),
aiohttp.web.view("/graph/leaves/count/{src}", CountLeavesView),
aiohttp.web.view("/graph/visit/nodes/count/{src}", CountVisitNodesView),
]
)
app["backend"] = backend
return app
def make_app_from_configfile():
"""Load configuration and then build application to run
"""
config_file = os.environ.get("SWH_CONFIG_FILENAME")
config = config_read(config_file)
return make_app(config=config)
diff --git a/swh/graph/swhid.py b/swh/graph/swhid.py
index aadd0d0..90db73f 100644
--- a/swh/graph/swhid.py
+++ b/swh/graph/swhid.py
@@ -1,419 +1,419 @@
# Copyright (C) 2019-2021 The Software Heritage developers
# See the AUTHORS file at the top-level directory of this distribution
# License: GNU General Public License version 3, or any later version
# See top-level LICENSE file for more information
from __future__ import annotations
from collections.abc import MutableMapping
from enum import Enum
import mmap
from mmap import MAP_SHARED, PROT_READ, PROT_WRITE
import os
import struct
from typing import BinaryIO, Iterator, Tuple
from swh.model.hashutil import hash_to_hex
-from swh.model.identifiers import ExtendedObjectType, ExtendedSWHID
+from swh.model.swhids import ExtendedObjectType, ExtendedSWHID
SWHID_BIN_FMT = "BB20s" # 2 unsigned chars + 20 bytes
INT_BIN_FMT = ">q" # big endian, 8-byte integer
SWHID_BIN_SIZE = 22 # in bytes
INT_BIN_SIZE = 8 # in bytes
class SwhidType(Enum):
"""types of existing SWHIDs, used to serialize ExtendedSWHID type as a (char)
integer
Note that the order does matter also for driving the binary search in
SWHID-indexed maps. Integer values also matter, for compatibility with the
Java layer.
"""
content = 0
directory = 1
origin = 2
release = 3
revision = 4
snapshot = 5
@classmethod
def from_extended_object_type(cls, object_type: ExtendedObjectType) -> SwhidType:
return cls[object_type.name.lower()]
def to_extended_object_type(self) -> ExtendedObjectType:
return ExtendedObjectType[SwhidType(self).name.upper()]
def str_to_bytes(swhid_str: str) -> bytes:
"""Convert a SWHID to a byte sequence
The binary format used to represent SWHIDs as 22-byte long byte sequences as
follows:
- 1 byte for the namespace version represented as a C `unsigned char`
- 1 byte for the object type, as the int value of :class:`SwhidType` enums,
represented as a C `unsigned char`
- 20 bytes for the SHA1 digest as a byte sequence
Args:
swhid: persistent identifier
Returns:
bytes: byte sequence representation of swhid
"""
swhid = ExtendedSWHID.from_string(swhid_str)
return struct.pack(
SWHID_BIN_FMT,
swhid.scheme_version,
SwhidType.from_extended_object_type(swhid.object_type).value,
swhid.object_id,
)
def bytes_to_str(bytes: bytes) -> str:
"""Inverse function of :func:`str_to_bytes`
See :func:`str_to_bytes` for a description of the binary SWHID format.
Args:
bytes: byte sequence representation of swhid
Returns:
swhid: persistent identifier
"""
(version, type, bin_digest) = struct.unpack(SWHID_BIN_FMT, bytes)
# The following is equivalent to:
# return str(ExtendedSWHID(
# object_type=SwhidType(type).to_extended_object_type(), object_id=bin_digest
# )
# but more efficient, because ExtendedSWHID.__init__ is extremely slow.
object_type = ExtendedObjectType[SwhidType(type).name.upper()]
return f"swh:1:{object_type.value}:{hash_to_hex(bin_digest)}"
class _OnDiskMap:
"""mmap-ed on-disk sequence of fixed size records"""
def __init__(
self, record_size: int, fname: str, mode: str = "rb", length: int = None
):
"""open an existing on-disk map
Args:
record_size: size of each record in bytes
fname: path to the on-disk map
mode: file open mode, usually either 'rb' for read-only maps, 'wb'
for creating new maps, or 'rb+' for updating existing ones
(default: 'rb')
length: map size in number of logical records; used to initialize
writable maps at creation time. Must be given when mode is 'wb'
and the map doesn't exist on disk; ignored otherwise
"""
os_modes = {"rb": os.O_RDONLY, "wb": os.O_RDWR | os.O_CREAT, "rb+": os.O_RDWR}
if mode not in os_modes:
raise ValueError("invalid file open mode: " + mode)
new_map = mode == "wb"
writable_map = mode in ["wb", "rb+"]
self.record_size = record_size
self.fd = os.open(fname, os_modes[mode])
if new_map:
if length is None:
raise ValueError("missing length when creating new map")
os.truncate(self.fd, length * self.record_size)
self.size = os.path.getsize(fname)
(self.length, remainder) = divmod(self.size, record_size)
if remainder:
raise ValueError(
"map size {} is not a multiple of the record size {}".format(
self.size, record_size
)
)
self.mm = mmap.mmap(
self.fd,
self.size,
prot=(PROT_READ | PROT_WRITE if writable_map else PROT_READ),
flags=MAP_SHARED,
)
def close(self) -> None:
"""close the map
shuts down both the mmap and the underlying file descriptor
"""
if not self.mm.closed:
self.mm.close()
os.close(self.fd)
def __len__(self) -> int:
return self.length
def __delitem__(self, pos: int) -> None:
raise NotImplementedError("cannot delete records from fixed-size map")
class SwhidToNodeMap(_OnDiskMap, MutableMapping):
"""memory mapped map from :ref:`SWHIDs <persistent-identifiers>` to a
continuous range 0..N of (8-byte long) integers
This is the converse mapping of :class:`NodeToSwhidMap`.
The on-disk serialization format is a sequence of fixed length (30 bytes)
records with the following fields:
- SWHID (22 bytes): binary SWHID representation as per :func:`str_to_bytes`
- long (8 bytes): big endian long integer
The records are sorted lexicographically by SWHID type and checksum, where
type is the integer value of :class:`SwhidType`. SWHID lookup in the map is
performed via binary search. Hence a huge map with, say, 11 B entries,
will require ~30 disk seeks.
Note that, due to fixed size + ordering, it is not possible to create these
maps by random writing. Hence, __setitem__ can be used only to *update* the
value associated to an existing key, rather than to add a missing item. To
create an entire map from scratch, you should do so *sequentially*, using
static method :meth:`write_record` (or, at your own risk, by hand via the
mmap :attr:`mm`).
"""
# record binary format: SWHID + a big endian 8-byte big endian integer
RECORD_BIN_FMT = ">" + SWHID_BIN_FMT + "q"
RECORD_SIZE = SWHID_BIN_SIZE + INT_BIN_SIZE
def __init__(self, fname: str, mode: str = "rb", length: int = None):
"""open an existing on-disk map
Args:
fname: path to the on-disk map
mode: file open mode, usually either 'rb' for read-only maps, 'wb'
for creating new maps, or 'rb+' for updating existing ones
(default: 'rb')
length: map size in number of logical records; used to initialize
read-write maps at creation time. Must be given when mode is
'wb'; ignored otherwise
"""
super().__init__(self.RECORD_SIZE, fname, mode=mode, length=length)
def _get_bin_record(self, pos: int) -> Tuple[bytes, bytes]:
"""seek and return the (binary) record at a given (logical) position
see :func:`_get_record` for an equivalent function with additional
deserialization
Args:
pos: 0-based record number
Returns:
a pair `(swhid, int)`, where swhid and int are bytes
"""
rec_pos = pos * self.RECORD_SIZE
int_pos = rec_pos + SWHID_BIN_SIZE
return (self.mm[rec_pos:int_pos], self.mm[int_pos : int_pos + INT_BIN_SIZE])
def _get_record(self, pos: int) -> Tuple[str, int]:
"""seek and return the record at a given (logical) position
moral equivalent of :func:`_get_bin_record`, with additional
deserialization to non-bytes types
Args:
pos: 0-based record number
Returns:
a pair `(swhid, int)`, where swhid is a string-based SWHID and int the
corresponding integer identifier
"""
(swhid_bytes, int_bytes) = self._get_bin_record(pos)
return (bytes_to_str(swhid_bytes), struct.unpack(INT_BIN_FMT, int_bytes)[0])
@classmethod
def write_record(cls, f: BinaryIO, swhid: str, int: int) -> None:
"""write a logical record to a file-like object
Args:
f: file-like object to write the record to
swhid: textual SWHID
int: SWHID integer identifier
"""
f.write(str_to_bytes(swhid))
f.write(struct.pack(INT_BIN_FMT, int))
def _bisect_pos(self, swhid_str: str) -> int:
"""bisect the position of the given identifier. If the identifier is
not found, the position of the swhid immediately after is returned.
Args:
swhid_str: the swhid as a string
Returns:
the logical record of the bisected position in the map
"""
if not isinstance(swhid_str, str):
raise TypeError("SWHID must be a str, not {}".format(type(swhid_str)))
try:
target = str_to_bytes(swhid_str) # desired SWHID as bytes
except ValueError:
raise ValueError('invalid SWHID: "{}"'.format(swhid_str))
lo = 0
hi = self.length - 1
while lo < hi:
mid = (lo + hi) // 2
(swhid, _value) = self._get_bin_record(mid)
if swhid < target:
lo = mid + 1
else:
hi = mid
return lo
def _find(self, swhid_str: str) -> Tuple[int, int]:
"""lookup the integer identifier of a swhid and its position
Args:
swhid_str: the swhid as a string
Returns:
a pair `(swhid, pos)` with swhid integer identifier and its logical
record position in the map
"""
pos = self._bisect_pos(swhid_str)
swhid_found, value = self._get_record(pos)
if swhid_found == swhid_str:
return (value, pos)
raise KeyError(swhid_str)
def __getitem__(self, swhid_str: str) -> int:
"""lookup the integer identifier of a SWHID
Args:
swhid: the SWHID as a string
Returns:
the integer identifier of swhid
"""
return self._find(swhid_str)[0] # return element, ignore position
def __setitem__(self, swhid_str: str, int: str) -> None:
(_swhid, pos) = self._find(swhid_str) # might raise KeyError and that's OK
rec_pos = pos * self.RECORD_SIZE
int_pos = rec_pos + SWHID_BIN_SIZE
self.mm[rec_pos:int_pos] = str_to_bytes(swhid_str)
self.mm[int_pos : int_pos + INT_BIN_SIZE] = struct.pack(INT_BIN_FMT, int)
def __iter__(self) -> Iterator[Tuple[str, int]]:
for pos in range(self.length):
yield self._get_record(pos)
def iter_prefix(self, prefix: str):
swh, n, t, sha = prefix.split(":")
sha = sha.ljust(40, "0")
start_swhid = ":".join([swh, n, t, sha])
start = self._bisect_pos(start_swhid)
for pos in range(start, self.length):
swhid, value = self._get_record(pos)
if not swhid.startswith(prefix):
break
yield swhid, value
def iter_type(self, swhid_type: str) -> Iterator[Tuple[str, int]]:
prefix = "swh:1:{}:".format(swhid_type)
yield from self.iter_prefix(prefix)
class NodeToSwhidMap(_OnDiskMap, MutableMapping):
"""memory mapped map from a continuous range of 0..N (8-byte long) integers to
:ref:`SWHIDs <persistent-identifiers>`
This is the converse mapping of :class:`SwhidToNodeMap`.
The on-disk serialization format is a sequence of fixed length records (22
bytes), each being the binary representation of a SWHID as per
:func:`str_to_bytes`.
The records are sorted by long integer, so that integer lookup is possible
via fixed-offset seek.
"""
RECORD_BIN_FMT = SWHID_BIN_FMT
RECORD_SIZE = SWHID_BIN_SIZE
def __init__(self, fname: str, mode: str = "rb", length: int = None):
"""open an existing on-disk map
Args:
fname: path to the on-disk map
mode: file open mode, usually either 'rb' for read-only maps, 'wb'
for creating new maps, or 'rb+' for updating existing ones
(default: 'rb')
size: map size in number of logical records; used to initialize
read-write maps at creation time. Must be given when mode is
'wb'; ignored otherwise
length: passed to :class:`_OnDiskMap`
"""
super().__init__(self.RECORD_SIZE, fname, mode=mode, length=length)
def _get_bin_record(self, pos: int) -> bytes:
"""seek and return the (binary) SWHID at a given (logical) position
Args:
pos: 0-based record number
Returns:
SWHID as a byte sequence
"""
rec_pos = pos * self.RECORD_SIZE
return self.mm[rec_pos : rec_pos + self.RECORD_SIZE]
@classmethod
def write_record(cls, f: BinaryIO, swhid: str) -> None:
"""write a SWHID to a file-like object
Args:
f: file-like object to write the record to
swhid: textual SWHID
"""
f.write(str_to_bytes(swhid))
def __getitem__(self, pos: int) -> str:
orig_pos = pos
if pos < 0:
pos = len(self) + pos
if not (0 <= pos < len(self)):
raise IndexError(orig_pos)
return bytes_to_str(self._get_bin_record(pos))
def __setitem__(self, pos: int, swhid: str) -> None:
rec_pos = pos * self.RECORD_SIZE
self.mm[rec_pos : rec_pos + self.RECORD_SIZE] = str_to_bytes(swhid)
def __iter__(self) -> Iterator[Tuple[int, str]]:
for pos in range(self.length):
yield (pos, self[pos])
diff --git a/swh/graph/tests/test_swhid.py b/swh/graph/tests/test_swhid.py
index 722e6b1..6053215 100644
--- a/swh/graph/tests/test_swhid.py
+++ b/swh/graph/tests/test_swhid.py
@@ -1,196 +1,196 @@
# Copyright (C) 2019-2020 The Software Heritage developers
# See the AUTHORS file at the top-level directory of this distribution
# License: GNU General Public License version 3, or any later version
# See top-level LICENSE file for more information
from itertools import islice
import os
import shutil
import tempfile
import unittest
from swh.graph.swhid import NodeToSwhidMap, SwhidToNodeMap, bytes_to_str, str_to_bytes
-from swh.model.identifiers import SWHID_TYPES
+from swh.model.swhids import SWHID_TYPES
class TestSwhidSerialization(unittest.TestCase):
pairs = [
(
"swh:1:cnt:94a9ed024d3859793618152ea559a168bbcbb5e2",
bytes.fromhex("01" + "00" + "94a9ed024d3859793618152ea559a168bbcbb5e2"),
),
(
"swh:1:dir:d198bc9d7a6bcf6db04f476d29314f157507d505",
bytes.fromhex("01" + "01" + "d198bc9d7a6bcf6db04f476d29314f157507d505"),
),
(
"swh:1:ori:b63a575fe3faab7692c9f38fb09d4bb45651bb0f",
bytes.fromhex("01" + "02" + "b63a575fe3faab7692c9f38fb09d4bb45651bb0f"),
),
(
"swh:1:rel:22ece559cc7cc2364edc5e5593d63ae8bd229f9f",
bytes.fromhex("01" + "03" + "22ece559cc7cc2364edc5e5593d63ae8bd229f9f"),
),
(
"swh:1:rev:309cf2674ee7a0749978cf8265ab91a60aea0f7d",
bytes.fromhex("01" + "04" + "309cf2674ee7a0749978cf8265ab91a60aea0f7d"),
),
(
"swh:1:snp:c7c108084bc0bf3d81436bf980b46e98bd338453",
bytes.fromhex("01" + "05" + "c7c108084bc0bf3d81436bf980b46e98bd338453"),
),
]
def test_str_to_bytes(self):
for (swhid_str, swhid_bytes) in self.pairs:
self.assertEqual(str_to_bytes(swhid_str), swhid_bytes)
def test_bytes_to_str(self):
for (swhid_str, swhid_bytes) in self.pairs:
self.assertEqual(bytes_to_str(swhid_bytes), swhid_str)
def test_round_trip(self):
for (swhid_str, swhid_bytes) in self.pairs:
self.assertEqual(swhid_str, bytes_to_str(str_to_bytes(swhid_str)))
self.assertEqual(swhid_bytes, str_to_bytes(bytes_to_str(swhid_bytes)))
def gen_records(types=["cnt", "dir", "ori", "rel", "rev", "snp"], length=10000):
"""generate sequential SWHID/int records, suitable for filling int<->swhid maps for
testing swh-graph on-disk binary databases
Args:
types (list): list of SWHID types to be generated, specified as the
corresponding 3-letter component in SWHIDs
length (int): number of SWHIDs to generate *per type*
Yields:
pairs (swhid, int) where swhid is a textual SWHID and int its sequential
integer identifier
"""
pos = 0
for t in sorted(types):
for i in range(0, length):
seq = format(pos, "x") # current position as hex string
swhid = "swh:1:{}:{}{}".format(t, "0" * (40 - len(seq)), seq)
yield (swhid, pos)
pos += 1
# pairs SWHID/position in the sequence generated by :func:`gen_records` above
MAP_PAIRS = [
("swh:1:cnt:0000000000000000000000000000000000000000", 0),
("swh:1:cnt:000000000000000000000000000000000000002a", 42),
("swh:1:dir:0000000000000000000000000000000000002afc", 11004),
("swh:1:ori:00000000000000000000000000000000000056ce", 22222),
("swh:1:rel:0000000000000000000000000000000000008235", 33333),
("swh:1:rev:000000000000000000000000000000000000ad9c", 44444),
("swh:1:snp:000000000000000000000000000000000000ea5f", 59999),
]
class TestSwhidToNodeMap(unittest.TestCase):
@classmethod
def setUpClass(cls):
"""create reasonably sized (~2 MB) SWHID->int map to test on-disk DB"""
cls.tmpdir = tempfile.mkdtemp(prefix="swh.graph.test.")
cls.fname = os.path.join(cls.tmpdir, "swhid2int.bin")
with open(cls.fname, "wb") as f:
for (swhid, i) in gen_records(length=10000):
SwhidToNodeMap.write_record(f, swhid, i)
@classmethod
def tearDownClass(cls):
shutil.rmtree(cls.tmpdir)
def setUp(self):
self.map = SwhidToNodeMap(self.fname)
def tearDown(self):
self.map.close()
def test_lookup(self):
for (swhid, pos) in MAP_PAIRS:
self.assertEqual(self.map[swhid], pos)
def test_missing(self):
with self.assertRaises(KeyError):
self.map["swh:1:ori:0101010100000000000000000000000000000000"],
with self.assertRaises(KeyError):
self.map["swh:1:cnt:0101010100000000000000000000000000000000"],
def test_type_error(self):
with self.assertRaises(TypeError):
self.map[42]
with self.assertRaises(TypeError):
self.map[1.2]
def test_update(self):
fname2 = self.fname + ".update"
shutil.copy(self.fname, fname2) # fresh map copy
map2 = SwhidToNodeMap(fname2, mode="rb+")
for (swhid, int) in islice(map2, 11): # update the first N items
new_int = int + 42
map2[swhid] = new_int
self.assertEqual(map2[swhid], new_int) # check updated value
os.unlink(fname2) # tmpdir will be cleaned even if we don't reach this
def test_iter_type(self):
for t in SWHID_TYPES + ["ori"]:
first_20 = list(islice(self.map.iter_type(t), 20))
k = first_20[0][1]
expected = [("swh:1:{}:{:040x}".format(t, i), i) for i in range(k, k + 20)]
assert first_20 == expected
def test_iter_prefix(self):
for t in SWHID_TYPES + ["ori"]:
prefix = self.map.iter_prefix("swh:1:{}:00".format(t))
first_20 = list(islice(prefix, 20))
k = first_20[0][1]
expected = [("swh:1:{}:{:040x}".format(t, i), i) for i in range(k, k + 20)]
assert first_20 == expected
class TestNodeToSwhidMap(unittest.TestCase):
@classmethod
def setUpClass(cls):
"""create reasonably sized (~1 MB) int->SWHID map to test on-disk DB"""
cls.tmpdir = tempfile.mkdtemp(prefix="swh.graph.test.")
cls.fname = os.path.join(cls.tmpdir, "int2swhid.bin")
with open(cls.fname, "wb") as f:
for (swhid, _i) in gen_records(length=10000):
NodeToSwhidMap.write_record(f, swhid)
@classmethod
def tearDownClass(cls):
shutil.rmtree(cls.tmpdir)
def setUp(self):
self.map = NodeToSwhidMap(self.fname)
def tearDown(self):
self.map.close()
def test_lookup(self):
for (swhid, pos) in MAP_PAIRS:
self.assertEqual(self.map[pos], swhid)
def test_out_of_bounds(self):
with self.assertRaises(IndexError):
self.map[1000000]
with self.assertRaises(IndexError):
self.map[-1000000]
def test_update(self):
fname2 = self.fname + ".update"
shutil.copy(self.fname, fname2) # fresh map copy
map2 = NodeToSwhidMap(fname2, mode="rb+")
for (int, swhid) in islice(map2, 11): # update the first N items
new_swhid = swhid.replace(":0", ":f") # mangle first hex digit
map2[int] = new_swhid
self.assertEqual(map2[int], new_swhid) # check updated value
os.unlink(fname2) # tmpdir will be cleaned even if we don't reach this