diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 380c658..5aaa4c1 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -1,46 +1,45 @@ repos: - repo: https://github.com/pre-commit/pre-commit-hooks rev: v2.4.0 hooks: - id: trailing-whitespace - id: flake8 - id: check-json - id: check-yaml - repo: https://github.com/codespell-project/codespell rev: v1.16.0 hooks: - id: codespell - repo: local hooks: - id: mypy name: mypy entry: mypy args: [swh] pass_filenames: false language: system types: [python] +- repo: https://github.com/python/black + rev: 19.10b0 + hooks: + - id: black + # unfortunately, we are far from being able to enable this... # - repo: https://github.com/PyCQA/pydocstyle.git # rev: 4.0.0 # hooks: # - id: pydocstyle # name: pydocstyle # description: pydocstyle is a static analysis tool for checking compliance with Python docstring conventions. # entry: pydocstyle --convention=google # language: python # types: [python] -# black requires py3.6+ -#- repo: https://github.com/python/black -# rev: 19.3b0 -# hooks: -# - id: black -# language_version: python3 #- repo: https://github.com/asottile/blacken-docs # rev: v1.0.0-1 # hooks: # - id: blacken-docs # additional_dependencies: [black==19.3b0] diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..b5413f6 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,2 @@ +[tool.black] +target-version = ['py37'] diff --git a/setup.cfg b/setup.cfg new file mode 100644 index 0000000..8d79b7e --- /dev/null +++ b/setup.cfg @@ -0,0 +1,6 @@ +[flake8] +# E203: whitespaces before ':' +# E231: missing whitespace after ',' +# W503: line break before binary operator +ignore = E203,E231,W503 +max-line-length = 88 diff --git a/setup.py b/setup.py index 66bcd5a..e6c9616 100755 --- a/setup.py +++ b/setup.py @@ -1,70 +1,70 @@ #!/usr/bin/env python3 # Copyright (C) 2015-2018 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 setuptools import setup, find_packages from os import path from io import open here = path.abspath(path.dirname(__file__)) # Get the long description from the README file with open(path.join(here, "README.md"), encoding="utf-8") as f: long_description = f.read() def parse_requirements(name=None): if name: reqf = "requirements-%s.txt" % name else: reqf = "requirements.txt" requirements = [] if not path.exists(reqf): return requirements with open(reqf) as f: for line in f.readlines(): line = line.strip() if not line or line.startswith("#"): continue requirements.append(line) return requirements setup( name="swh.dataset", description="Software Heritage dataset tools", long_description=long_description, long_description_content_type="text/markdown", author="Software Heritage developers", author_email="swh-devel@inria.fr", url="https://forge.softwareheritage.org/source/swh-dataset/", packages=find_packages(), # packages's modules install_requires=parse_requirements() + parse_requirements("swh"), tests_require=parse_requirements("test"), - entry_points=''' + entry_points=""" [swh.cli.subcommands] dataset=swh.dataset.cli:cli - ''', + """, setup_requires=["vcversioner"], extras_require={"testing": parse_requirements("test")}, vcversioner={}, include_package_data=True, classifiers=[ "Programming Language :: Python :: 3", "Intended Audience :: Developers", "License :: OSI Approved :: GNU General Public License v3 (GPLv3)", "Operating System :: OS Independent", "Development Status :: 3 - Alpha", ], project_urls={ "Bug Reports": "https://forge.softwareheritage.org/maniphest", "Funding": "https://www.softwareheritage.org/donate", "Source": "https://forge.softwareheritage.org/source/swh-dataset", "Documentation": "https://docs.softwareheritage.org/devel/swh-dataset/", }, ) diff --git a/swh/dataset/cli.py b/swh/dataset/cli.py index 33291ae..55ac9ab 100644 --- a/swh/dataset/cli.py +++ b/swh/dataset/cli.py @@ -1,54 +1,57 @@ # Copyright (C) 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 click import uuid from swh.core import config from swh.core.cli import CONTEXT_SETTINGS from swh.dataset.graph import export_edges, sort_graph_nodes -@click.group(name='dataset', context_settings=CONTEXT_SETTINGS) -@click.option('--config-file', '-C', default=None, - type=click.Path(exists=True, dir_okay=False), - help="Configuration file.") +@click.group(name="dataset", context_settings=CONTEXT_SETTINGS) +@click.option( + "--config-file", + "-C", + default=None, + type=click.Path(exists=True, dir_okay=False), + help="Configuration file.", +) @click.pass_context def cli(ctx, config_file): - '''Software Heritage Dataset Tools''' + """Software Heritage Dataset Tools""" ctx.ensure_object(dict) conf = config.read(config_file) - ctx.obj['config'] = conf + ctx.obj["config"] = conf -@cli.group('graph') +@cli.group("graph") @click.pass_context def graph(ctx): """Manage graph edges export""" pass -@graph.command('export') -@click.argument('export-path', type=click.Path()) -@click.option('--export-id', '-e', help="Unique ID of the export run.") -@click.option('--processes', '-p', default=1, - help="Number of parallel processes") +@graph.command("export") +@click.argument("export-path", type=click.Path()) +@click.option("--export-id", "-e", help="Unique ID of the export run.") +@click.option("--processes", "-p", default=1, help="Number of parallel processes") @click.pass_context def export_graph(ctx, export_path, export_id, processes): """Export the Software Heritage graph as an edge dataset.""" - config = ctx.obj['config'] + config = ctx.obj["config"] if not export_id: export_id = str(uuid.uuid4()) export_edges(config, export_path, export_id, processes) -@graph.command('sort') -@click.argument('export-path', type=click.Path()) +@graph.command("sort") +@click.argument("export-path", type=click.Path()) @click.pass_context def sort_graph(ctx, export_path): - config = ctx.obj['config'] + config = ctx.obj["config"] sort_graph_nodes(export_path, config) diff --git a/swh/dataset/exporter.py b/swh/dataset/exporter.py index 7e3e42e..1b9d870 100644 --- a/swh/dataset/exporter.py +++ b/swh/dataset/exporter.py @@ -1,232 +1,237 @@ # Copyright (C) 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 concurrent.futures import multiprocessing import time import tqdm from typing import Mapping, Sequence, Tuple from concurrent.futures import FIRST_EXCEPTION, ProcessPoolExecutor from confluent_kafka import TopicPartition from swh.journal.client import JournalClient class JournalClientOffsetRanges(JournalClient): """ A subclass of JournalClient reading only inside some specific offset range. Partition assignments have to be manually given to the class. This client can only read a single topic at a time. """ + def __init__( - self, - *args, - offset_ranges: Mapping[int, Tuple[int, int]] = None, - assignment: Sequence[int] = None, - progress_queue: multiprocessing.Queue = None, - refresh_every: int = 200, - **kwargs, + self, + *args, + offset_ranges: Mapping[int, Tuple[int, int]] = None, + assignment: Sequence[int] = None, + progress_queue: multiprocessing.Queue = None, + refresh_every: int = 200, + **kwargs, ): """ Args: offset_ranges: A mapping of partition_id -> (low, high) offsets that define the boundaries of the messages to consume. assignment: The list of partitions to assign to this client. progress_queue: a multiprocessing.Queue where the current progress will be reported. refresh_every: the refreshing rate of the progress reporting. """ self.offset_ranges = offset_ranges self.progress_queue = progress_queue self.refresh_every = refresh_every self.assignment = assignment self.count = None self.topic_name = None super().__init__(*args, **kwargs) def subscribe(self): self.topic_name = self.subscription[0] time.sleep(0.1) # https://github.com/edenhill/librdkafka/issues/1983 - self.consumer.assign([ - TopicPartition(self.topic_name, pid) for pid in self.assignment - ]) + self.consumer.assign( + [TopicPartition(self.topic_name, pid) for pid in self.assignment] + ) def process(self, *args, **kwargs): self.count = 0 try: self.handle_committed_offsets() super().process(*args, **kwargs) except EOFError: self.progress_queue.put(None) - def handle_committed_offsets(self, ): + def handle_committed_offsets(self,): """ Handle already committed partition offsets before starting processing. """ - committed = self.consumer.committed([ - TopicPartition(self.topic_name, pid) for pid in self.assignment - ]) + committed = self.consumer.committed( + [TopicPartition(self.topic_name, pid) for pid in self.assignment] + ) for tp in committed: self.handle_offset(tp.partition, tp.offset) def handle_offset(self, partition_id, offset): """ Check whether the client has reached the end of the current partition, and trigger a reassignment if that is the case. Raise EOFError if all the partitions have reached the end. """ if offset < 0: # Uninitialized partition offset return if self.count % self.refresh_every == 0: self.progress_queue.put({partition_id: offset}) if offset >= self.offset_ranges[partition_id][1] - 1: - self.assignment = [pid for pid in self.assignment - if pid != partition_id] + self.assignment = [pid for pid in self.assignment if pid != partition_id] self.subscribe() if not self.assignment: raise EOFError def deserialize_message(self, message): """ Override of the message deserialization to hook the handling of the message offset. """ self.handle_offset(message.partition(), message.offset()) self.count += 1 return super().deserialize_message(message) class ParallelExporter: """ Base class for all the Journal exporters. Each exporter should override the `export_worker` function with an implementation of how to run the message processing. """ + def __init__(self, config, export_id: str, obj_type, processes=1): """ Args: config: the exporter config, which should also include the JournalClient configuration. export_id: a unique identifier for the export that will be used as part of a Kafka consumer group ID. obj_type: The type of SWH object to export. processes: The number of processes to run. """ self.config = config - self.export_id = 'swh-dataset-export-{}'.format(export_id) + self.export_id = "swh-dataset-export-{}".format(export_id) self.obj_type = obj_type self.processes = processes self.offsets = None def get_offsets(self): """ First pass to fetch all the current low and high offsets of each partition to define the consumption boundaries. """ if self.offsets is None: client = JournalClient( - **self.config['journal'], + **self.config["journal"], object_types=[self.obj_type], group_id=self.export_id, ) topic_name = client.subscription[0] topics = client.consumer.list_topics(topic_name).topics partitions = topics[topic_name].partitions self.offsets = {} - for partition_id in tqdm.tqdm(partitions.keys(), - desc=" - Partition offsets"): + for partition_id in tqdm.tqdm( + partitions.keys(), desc=" - Partition offsets" + ): tp = TopicPartition(topic_name, partition_id) (lo, hi) = client.consumer.get_watermark_offsets(tp) self.offsets[partition_id] = (lo, hi) return self.offsets def run(self, *args): """ Run the parallel export. """ offsets = self.get_offsets() to_assign = list(offsets.keys()) manager = multiprocessing.Manager() q = manager.Queue() with ProcessPoolExecutor(self.processes + 1) as pool: futures = [] for i in range(self.processes): - futures.append(pool.submit( - self.export_worker, - *args, - assignment=to_assign[i::self.processes], - queue=q - )) + futures.append( + pool.submit( + self.export_worker, + *args, + assignment=to_assign[i :: self.processes], + queue=q, + ) + ) futures.append(pool.submit(self.progress_worker, queue=q)) concurrent.futures.wait(futures, return_when=FIRST_EXCEPTION) for f in futures: if f.running(): continue exc = f.exception() if exc: pool.shutdown(wait=False) f.result() raise exc def progress_worker(self, *args, queue=None): """ An additional worker process that reports the current progress of the export between all the different parallel consumers and across all the partitions, by consuming the shared progress reporting Queue. """ d = {} active_workers = self.processes offset_diff = sum((hi - lo) for lo, hi in self.offsets.values()) with tqdm.tqdm(total=offset_diff, desc=" - Journal export") as pbar: while active_workers: item = queue.get() if item is None: active_workers -= 1 continue d.update(item) progress = sum(n - self.offsets[p][0] for p, n in d.items()) - pbar.set_postfix(active_workers=active_workers, - total_workers=self.processes) + pbar.set_postfix( + active_workers=active_workers, total_workers=self.processes + ) pbar.update(progress - pbar.n) def process(self, callback, assignment=None, queue=None): client = JournalClientOffsetRanges( - **self.config['journal'], + **self.config["journal"], object_types=[self.obj_type], group_id=self.export_id, - debug='cgrp,broker', + debug="cgrp,broker", offset_ranges=self.offsets, assignment=assignment, progress_queue=queue, ) client.process(callback) def export_worker(self, *args, **kwargs): """ Override this with a custom implementation of a worker function. A worker function should call `self.process(fn, **kwargs)` with `fn` being a callback that will be called in the same fashion as with `JournalClient.process()`. A simple exporter to print all the objects in the log would look like this: ``` class PrintExporter(ParallelExporter): def export_worker(self, **kwargs): self.process(print, **kwargs) ``` """ raise NotImplementedError diff --git a/swh/dataset/graph.py b/swh/dataset/graph.py index 7ac0551..395575b 100644 --- a/swh/dataset/graph.py +++ b/swh/dataset/graph.py @@ -1,197 +1,201 @@ # Copyright (C) 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 functools import os import os.path import pathlib import shlex import subprocess import tempfile import uuid from swh.dataset.exporter import ParallelExporter from swh.dataset.utils import ZSTFile from swh.model.identifiers import origin_identifier, persistent_identifier from swh.storage.fixer import fix_objects def process_messages(messages, config, node_writer, edge_writer): """ Args: messages: A sequence of messages to process config: The exporter configuration node_writer: A file-like object where to write nodes edge_writer: A file-like object where to write edges """ + def write_node(node): node_type, node_id = node if node_id is None: return - node_pid = persistent_identifier(object_type=node_type, - object_id=node_id) - node_writer.write('{}\n'.format(node_pid)) + node_pid = persistent_identifier(object_type=node_type, object_id=node_id) + node_writer.write("{}\n".format(node_pid)) def write_edge(src, dst): src_type, src_id = src dst_type, dst_id = dst if src_id is None or dst_id is None: return src_pid = persistent_identifier(object_type=src_type, object_id=src_id) dst_pid = persistent_identifier(object_type=dst_type, object_id=dst_id) - edge_writer.write('{} {}\n'.format(src_pid, dst_pid)) + edge_writer.write("{} {}\n".format(src_pid, dst_pid)) messages = {k: fix_objects(k, v) for k, v in messages.items()} - for visit in messages.get('origin_visit', []): - origin_id = origin_identifier({'url': visit['origin']}) - write_node(('origin', origin_id)) - write_edge(('origin', origin_id), ('snapshot', visit['snapshot'])) - - for snapshot in messages.get('snapshot', []): - write_node(('snapshot', snapshot['id'])) - for branch_name, branch in snapshot['branches'].items(): - while branch and branch.get('target_type') == 'alias': - branch_name = branch['target'] - branch = snapshot['branches'][branch_name] + for visit in messages.get("origin_visit", []): + origin_id = origin_identifier({"url": visit["origin"]}) + write_node(("origin", origin_id)) + write_edge(("origin", origin_id), ("snapshot", visit["snapshot"])) + + for snapshot in messages.get("snapshot", []): + write_node(("snapshot", snapshot["id"])) + for branch_name, branch in snapshot["branches"].items(): + while branch and branch.get("target_type") == "alias": + branch_name = branch["target"] + branch = snapshot["branches"][branch_name] if branch is None or not branch_name: continue - if (config.get('remove_pull_requests') - and (branch_name.startswith(b'refs/pull') - or branch_name.startswith(b'refs/merge-requests'))): + if config.get("remove_pull_requests") and ( + branch_name.startswith(b"refs/pull") + or branch_name.startswith(b"refs/merge-requests") + ): continue - write_edge(('snapshot', snapshot['id']), - (branch['target_type'], branch['target'])) - - for release in messages.get('release', []): - write_node(('release', release['id'])) - write_edge(('release', release['id']), - (release['target_type'], release['target'])) - - for revision in messages.get('revision', []): - write_node(('revision', revision['id'])) - write_edge(('revision', revision['id']), - ('directory', revision['directory'])) - for parent in revision['parents']: - write_edge(('revision', revision['id']), - ('revision', parent)) - - for directory in messages.get('directory', []): - write_node(('directory', directory['id'])) - for entry in directory['entries']: + write_edge( + ("snapshot", snapshot["id"]), (branch["target_type"], branch["target"]) + ) + + for release in messages.get("release", []): + write_node(("release", release["id"])) + write_edge( + ("release", release["id"]), (release["target_type"], release["target"]) + ) + + for revision in messages.get("revision", []): + write_node(("revision", revision["id"])) + write_edge(("revision", revision["id"]), ("directory", revision["directory"])) + for parent in revision["parents"]: + write_edge(("revision", revision["id"]), ("revision", parent)) + + for directory in messages.get("directory", []): + write_node(("directory", directory["id"])) + for entry in directory["entries"]: entry_type_mapping = { - 'file': 'content', - 'dir': 'directory', - 'rev': 'revision' + "file": "content", + "dir": "directory", + "rev": "revision", } - write_edge(('directory', directory['id']), - (entry_type_mapping[entry['type']], entry['target'])) + write_edge( + ("directory", directory["id"]), + (entry_type_mapping[entry["type"]], entry["target"]), + ) - for content in messages.get('content', []): - write_node(('content', content['sha1_git'])) + for content in messages.get("content", []): + write_node(("content", content["sha1_git"])) class GraphEdgeExporter(ParallelExporter): """ Implementation of ParallelExporter which writes all the graph edges of a specific type in a Zstandard-compressed CSV file. Each row of the CSV is in the format: ` """ + def export_worker(self, export_path, **kwargs): dataset_path = pathlib.Path(export_path) dataset_path.mkdir(exist_ok=True, parents=True) - nodes_file = dataset_path / ('graph-{}.nodes.csv.zst' - .format(str(uuid.uuid4()))) - edges_file = dataset_path / ('graph-{}.edges.csv.zst' - .format(str(uuid.uuid4()))) - - with \ - ZSTFile(nodes_file, 'w') as nodes_writer, \ - ZSTFile(edges_file, 'w') as edges_writer: + nodes_file = dataset_path / ("graph-{}.nodes.csv.zst".format(str(uuid.uuid4()))) + edges_file = dataset_path / ("graph-{}.edges.csv.zst".format(str(uuid.uuid4()))) + + with ZSTFile(nodes_file, "w") as nodes_writer, ZSTFile( + edges_file, "w" + ) as edges_writer: process_fn = functools.partial( process_messages, config=self.config, nodes_writer=nodes_writer, - edges_writer=edges_writer + edges_writer=edges_writer, ) self.process(process_fn, **kwargs) def export_edges(config, export_path, export_id, processes): """Run the edge exporter for each edge type.""" object_types = [ - 'origin_visit', - 'snapshot', - 'release', - 'revision', - 'directory', + "origin_visit", + "snapshot", + "release", + "revision", + "directory", ] for obj_type in object_types: - print('{} edges:'.format(obj_type)) + print("{} edges:".format(obj_type)) exporter = GraphEdgeExporter(config, export_id, obj_type, processes) exporter.run(os.path.join(export_path, obj_type)) def sort_graph_nodes(export_path, config): """ Generate the node list from the edges files. We cannot solely rely on the object IDs that are read in the journal, as some nodes that are referred to as destinations in the edge file might not be present in the archive (e.g a rev_entry referring to a revision that we do not have crawled yet). The most efficient way of getting all the nodes that are mentioned in the edges file is therefore to use sort(1) on the gigantic edge files to get all the unique node IDs, while using the disk as a temporary buffer. This pipeline does, in order: - concatenate and write all the compressed edges files in graph.edges.csv.zst (using the fact that ZST compression is an additive function) ; - deflate the edges ; - count the number of edges and write it in graph.edges.count.txt ; - concatenate all the (deflated) nodes from the export with the destination edges, and sort the output to get the list of unique graph nodes ; - count the number of unique graph nodes and write it in graph.nodes.count.txt ; - compress and write the resulting nodes in graph.nodes.csv.zst. """ # Use bytes for the sorting algorithm (faster than being locale-specific) env = { **os.environ.copy(), - 'LC_ALL': 'C', - 'LC_COLLATE': 'C', - 'LANG': 'C', + "LC_ALL": "C", + "LC_COLLATE": "C", + "LANG": "C", } - sort_buffer_size = config.get('sort_buffer_size', '4G') - disk_buffer_dir = config.get('disk_buffer_dir', export_path) - with tempfile.TemporaryDirectory(prefix='.graph_node_sort_', - dir=disk_buffer_dir) as buffer_path: + sort_buffer_size = config.get("sort_buffer_size", "4G") + disk_buffer_dir = config.get("disk_buffer_dir", export_path) + with tempfile.TemporaryDirectory( + prefix=".graph_node_sort_", dir=disk_buffer_dir + ) as buffer_path: subprocess.run( [ "bash", "-c", - ("pv {export_path}/*/*.edges.csv.zst | " - "tee {export_path}/graph.edges.csv.zst |" - "zstdcat |" - "tee >( wc -l > {export_path}/graph.edges.count.txt ) |" - "cut -d' ' -f2 | " - "cat - <( zstdcat {export_path}/*/*.nodes.csv.zst ) | " - "sort -u -S{sort_buffer_size} -T{buffer_path} | " - "tee >( wc -l > {export_path}/graph.nodes.count.txt ) |" - "zstdmt > {export_path}/graph.nodes.csv.zst") - .format( + ( + "pv {export_path}/*/*.edges.csv.zst | " + "tee {export_path}/graph.edges.csv.zst |" + "zstdcat |" + "tee >( wc -l > {export_path}/graph.edges.count.txt ) |" + "cut -d' ' -f2 | " + "cat - <( zstdcat {export_path}/*/*.nodes.csv.zst ) | " + "sort -u -S{sort_buffer_size} -T{buffer_path} | " + "tee >( wc -l > {export_path}/graph.nodes.count.txt ) |" + "zstdmt > {export_path}/graph.nodes.csv.zst" + ).format( export_path=shlex.quote(str(export_path)), buffer_path=shlex.quote(str(buffer_path)), sort_buffer_size=shlex.quote(sort_buffer_size), ), ], env=env, ) diff --git a/swh/dataset/test/test_graph.py b/swh/dataset/test/test_graph.py index c6fb54a..b0c53e7 100644 --- a/swh/dataset/test/test_graph.py +++ b/swh/dataset/test/test_graph.py @@ -1,437 +1,420 @@ # Copyright (C) 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 hashlib from typing import Tuple import pytest from unittest.mock import Mock, call from swh.dataset.graph import process_messages, sort_graph_nodes from swh.dataset.utils import ZSTFile from swh.model.hashutil import MultiHash, hash_to_bytes DATE = { "timestamp": {"seconds": 1234567891, "microseconds": 0,}, "offset": 120, "negative_utc": False, } TEST_CONTENT = { **MultiHash.from_data(b"foo").digest(), "length": 3, "status": "visible", } TEST_REVISION = { "id": hash_to_bytes("7026b7c1a2af56521e951c01ed20f255fa054238"), "message": b"hello", "date": DATE, "committer": {"fullname": b"foo", "name": b"foo", "email": b""}, "author": {"fullname": b"foo", "name": b"foo", "email": b""}, "committer_date": DATE, "type": "git", "directory": b"\x01" * 20, "synthetic": False, "metadata": None, "parents": [], } TEST_RELEASE = { "id": hash_to_bytes("d81cc0710eb6cf9efd5b920a8453e1e07157b6cd"), "name": b"v0.0.1", "date": { "timestamp": {"seconds": 1234567890, "microseconds": 0,}, "offset": 120, "negative_utc": False, }, "author": {"author": {"fullname": b"foo", "name": b"foo", "email": b""}}, "target_type": "revision", "target": b"\x04" * 20, "message": b"foo", "synthetic": False, } TEST_ORIGIN = {"url": "https://somewhere.org/den/fox"} TEST_ORIGIN_VISIT = { "origin": TEST_ORIGIN["url"], "date": "2013-05-07 04:20:39.369271+00:00", "snapshot": None, # TODO "status": "ongoing", # TODO "metadata": {"foo": "bar"}, "type": "git", } @pytest.fixture def exporter(): def wrapped(messages, config=None) -> Tuple[Mock, Mock]: if config is None: config = {} node_writer = Mock() edge_writer = Mock() process_messages( messages, config=config, node_writer=node_writer, edge_writer=edge_writer, ) return node_writer.write, edge_writer.write return wrapped def binhash(s): return hashlib.sha1(s.encode()).digest() def hexhash(s): return hashlib.sha1(s.encode()).hexdigest() def test_export_origin_visits(exporter): node_writer, edge_writer = exporter( { "origin_visit": [ { **TEST_ORIGIN_VISIT, "origin": {"url": "ori1"}, "snapshot": binhash("snp1"), }, { **TEST_ORIGIN_VISIT, "origin": {"url": "ori2"}, "snapshot": binhash("snp2"), }, ] } ) assert node_writer.mock_calls == [ call(f"swh:1:ori:{hexhash('ori1')}\n"), call(f"swh:1:ori:{hexhash('ori2')}\n"), ] assert edge_writer.mock_calls == [ call(f"swh:1:ori:{hexhash('ori1')} swh:1:snp:{hexhash('snp1')}\n"), call(f"swh:1:ori:{hexhash('ori2')} swh:1:snp:{hexhash('snp2')}\n"), ] def test_export_snapshot_simple(exporter): node_writer, edge_writer = exporter( { "snapshot": [ { "id": binhash("snp1"), "branches": { b"refs/heads/master": { "target": binhash("rev1"), "target_type": "revision", }, b"HEAD": {"target": binhash("rev1"), "target_type": "revision"}, }, }, { "id": binhash("snp2"), "branches": { b"refs/heads/master": { "target": binhash("rev1"), "target_type": "revision", }, b"HEAD": {"target": binhash("rev2"), "target_type": "revision"}, b"bcnt": {"target": binhash("cnt1"), "target_type": "content"}, b"bdir": { "target": binhash("dir1"), "target_type": "directory", }, b"brel": {"target": binhash("rel1"), "target_type": "release"}, b"bsnp": {"target": binhash("snp1"), "target_type": "snapshot"}, }, }, {"id": binhash("snp3"), "branches": {}}, ] } ) assert node_writer.mock_calls == [ call(f"swh:1:snp:{hexhash('snp1')}\n"), call(f"swh:1:snp:{hexhash('snp2')}\n"), call(f"swh:1:snp:{hexhash('snp3')}\n"), ] assert edge_writer.mock_calls == [ call(f"swh:1:snp:{hexhash('snp1')} swh:1:rev:{hexhash('rev1')}\n"), call(f"swh:1:snp:{hexhash('snp1')} swh:1:rev:{hexhash('rev1')}\n"), call(f"swh:1:snp:{hexhash('snp2')} swh:1:rev:{hexhash('rev1')}\n"), call(f"swh:1:snp:{hexhash('snp2')} swh:1:rev:{hexhash('rev2')}\n"), call(f"swh:1:snp:{hexhash('snp2')} swh:1:cnt:{hexhash('cnt1')}\n"), call(f"swh:1:snp:{hexhash('snp2')} swh:1:dir:{hexhash('dir1')}\n"), call(f"swh:1:snp:{hexhash('snp2')} swh:1:rel:{hexhash('rel1')}\n"), call(f"swh:1:snp:{hexhash('snp2')} swh:1:snp:{hexhash('snp1')}\n"), ] def test_export_snapshot_aliases(exporter): node_writer, edge_writer = exporter( { "snapshot": [ { "id": binhash("snp1"), "branches": { b"origin_branch": { "target": binhash("rev1"), "target_type": "revision", }, b"alias1": {"target": b"origin_branch", "target_type": "alias"}, b"alias2": {"target": b"alias1", "target_type": "alias"}, b"alias3": {"target": b"alias2", "target_type": "alias"}, }, }, ] } ) assert node_writer.mock_calls == [call(f"swh:1:snp:{hexhash('snp1')}\n")] assert edge_writer.mock_calls == ( [call(f"swh:1:snp:{hexhash('snp1')} swh:1:rev:{hexhash('rev1')}\n")] * 4 ) def test_export_snapshot_no_pull_requests(exporter): snp = { "id": binhash("snp1"), "branches": { b"refs/heads/master": { "target": binhash("rev1"), "target_type": "revision", }, b"refs/pull/42": {"target": binhash("rev2"), "target_type": "revision"}, b"refs/merge-requests/lol": { "target": binhash("rev3"), "target_type": "revision", }, }, } node_writer, edge_writer = exporter({"snapshot": [snp]}) assert edge_writer.mock_calls == [ call(f"swh:1:snp:{hexhash('snp1')} swh:1:rev:{hexhash('rev1')}\n"), call(f"swh:1:snp:{hexhash('snp1')} swh:1:rev:{hexhash('rev2')}\n"), call(f"swh:1:snp:{hexhash('snp1')} swh:1:rev:{hexhash('rev3')}\n"), ] node_writer, edge_writer = exporter( {"snapshot": [snp]}, config={"remove_pull_requests": True} ) assert edge_writer.mock_calls == [ call(f"swh:1:snp:{hexhash('snp1')} swh:1:rev:{hexhash('rev1')}\n"), ] def test_export_releases(exporter): node_writer, edge_writer = exporter( { "release": [ { **TEST_RELEASE, "id": binhash("rel1"), "target": binhash("rev1"), "target_type": "revision", }, { **TEST_RELEASE, "id": binhash("rel2"), "target": binhash("rel1"), "target_type": "release", }, { **TEST_RELEASE, "id": binhash("rel3"), "target": binhash("dir1"), "target_type": "directory", }, { **TEST_RELEASE, "id": binhash("rel4"), "target": binhash("cnt1"), "target_type": "content", }, ] } ) assert node_writer.mock_calls == [ call(f"swh:1:rel:{hexhash('rel1')}\n"), call(f"swh:1:rel:{hexhash('rel2')}\n"), call(f"swh:1:rel:{hexhash('rel3')}\n"), call(f"swh:1:rel:{hexhash('rel4')}\n"), ] assert edge_writer.mock_calls == [ call(f"swh:1:rel:{hexhash('rel1')} swh:1:rev:{hexhash('rev1')}\n"), call(f"swh:1:rel:{hexhash('rel2')} swh:1:rel:{hexhash('rel1')}\n"), call(f"swh:1:rel:{hexhash('rel3')} swh:1:dir:{hexhash('dir1')}\n"), call(f"swh:1:rel:{hexhash('rel4')} swh:1:cnt:{hexhash('cnt1')}\n"), ] def test_export_revision(exporter): node_writer, edge_writer = exporter( { "revision": [ { **TEST_REVISION, "id": binhash("rev1"), "directory": binhash("dir1"), - "parents": [ - binhash("rev2"), - binhash("rev3"), - ], + "parents": [binhash("rev2"), binhash("rev3"),], }, { **TEST_REVISION, "id": binhash("rev2"), "directory": binhash("dir2"), "parents": [], }, ] } ) assert node_writer.mock_calls == [ call(f"swh:1:rev:{hexhash('rev1')}\n"), call(f"swh:1:rev:{hexhash('rev2')}\n"), ] assert edge_writer.mock_calls == [ call(f"swh:1:rev:{hexhash('rev1')} swh:1:dir:{hexhash('dir1')}\n"), call(f"swh:1:rev:{hexhash('rev1')} swh:1:rev:{hexhash('rev2')}\n"), call(f"swh:1:rev:{hexhash('rev1')} swh:1:rev:{hexhash('rev3')}\n"), - call(f"swh:1:rev:{hexhash('rev2')} swh:1:dir:{hexhash('dir2')}\n"), ] def test_export_directory(exporter): node_writer, edge_writer = exporter( { "directory": [ { "id": binhash("dir1"), "entries": [ {"type": "file", "target": binhash("cnt1")}, {"type": "dir", "target": binhash("dir2")}, {"type": "rev", "target": binhash("rev1")}, ], }, - { - "id": binhash("dir2"), - "entries": [], - }, + {"id": binhash("dir2"), "entries": [],}, ] } ) assert node_writer.mock_calls == [ call(f"swh:1:dir:{hexhash('dir1')}\n"), call(f"swh:1:dir:{hexhash('dir2')}\n"), ] assert edge_writer.mock_calls == [ call(f"swh:1:dir:{hexhash('dir1')} swh:1:cnt:{hexhash('cnt1')}\n"), call(f"swh:1:dir:{hexhash('dir1')} swh:1:dir:{hexhash('dir2')}\n"), call(f"swh:1:dir:{hexhash('dir1')} swh:1:rev:{hexhash('rev1')}\n"), ] def test_export_content(exporter): node_writer, edge_writer = exporter( { "content": [ - { - **TEST_CONTENT, - "sha1_git": binhash("cnt1"), - }, - { - **TEST_CONTENT, - "sha1_git": binhash("cnt2"), - }, + {**TEST_CONTENT, "sha1_git": binhash("cnt1"),}, + {**TEST_CONTENT, "sha1_git": binhash("cnt2"),}, ] } ) assert node_writer.mock_calls == [ call(f"swh:1:cnt:{hexhash('cnt1')}\n"), call(f"swh:1:cnt:{hexhash('cnt2')}\n"), ] assert edge_writer.mock_calls == [] def zstwrite(fp, lines): - with ZSTFile(fp, 'w') as writer: + with ZSTFile(fp, "w") as writer: for l in lines: writer.write(l + "\n") def zstread(fp): - with ZSTFile(fp, 'r') as reader: + with ZSTFile(fp, "r") as reader: return reader.read() def test_sort_pipeline(tmp_path): short_type_mapping = { - 'origin_visit': 'ori', - 'snapshot': 'snp', - 'release': 'rel', - 'revision': 'rev', - 'directory': 'dir', - 'content': 'cnt', + "origin_visit": "ori", + "snapshot": "snp", + "release": "rel", + "revision": "rev", + "directory": "dir", + "content": "cnt", } input_nodes = [ f"swh:1:{short}:{hexhash(short + str(x))}" for short in short_type_mapping.values() for x in range(4) ] input_edges = [ f"swh:1:ori:{hexhash('ori1')} swh:1:snp:{hexhash('snp1')}", f"swh:1:ori:{hexhash('ori2')} swh:1:snp:{hexhash('snp2')}", f"swh:1:ori:{hexhash('ori3')} swh:1:snp:{hexhash('snp3')}", f"swh:1:ori:{hexhash('ori4')} swh:1:snp:{hexhash('snpX')}", # missing dest - f"swh:1:snp:{hexhash('snp1')} swh:1:rev:{hexhash('rev1')}", # dup f"swh:1:snp:{hexhash('snp1')} swh:1:rev:{hexhash('rev1')}", # dup f"swh:1:snp:{hexhash('snp3')} swh:1:cnt:{hexhash('cnt1')}", f"swh:1:snp:{hexhash('snp4')} swh:1:rel:{hexhash('rel1')}", - f"swh:1:rel:{hexhash('rel1')} swh:1:rel:{hexhash('rel2')}", f"swh:1:rel:{hexhash('rel2')} swh:1:rev:{hexhash('rev1')}", f"swh:1:rel:{hexhash('rel3')} swh:1:rev:{hexhash('rev2')}", f"swh:1:rel:{hexhash('rel4')} swh:1:dir:{hexhash('dir1')}", - f"swh:1:rev:{hexhash('rev1')} swh:1:rev:{hexhash('rev1')}", # dup f"swh:1:rev:{hexhash('rev1')} swh:1:rev:{hexhash('rev1')}", # dup f"swh:1:rev:{hexhash('rev1')} swh:1:rev:{hexhash('rev2')}", f"swh:1:rev:{hexhash('rev2')} swh:1:rev:{hexhash('revX')}", # missing dest f"swh:1:rev:{hexhash('rev3')} swh:1:rev:{hexhash('rev2')}", f"swh:1:rev:{hexhash('rev4')} swh:1:dir:{hexhash('dir1')}", - f"swh:1:dir:{hexhash('dir1')} swh:1:cnt:{hexhash('cnt1')}", f"swh:1:dir:{hexhash('dir1')} swh:1:dir:{hexhash('dir1')}", f"swh:1:dir:{hexhash('dir1')} swh:1:rev:{hexhash('rev1')}", ] for obj_type, short_obj_type in short_type_mapping.items(): - p = (tmp_path / obj_type) + p = tmp_path / obj_type p.mkdir() edges = [e for e in input_edges if e.startswith(f"swh:1:{short_obj_type}")] - zstwrite(p / '00.edges.csv.zst', edges[0::2]) - zstwrite(p / '01.edges.csv.zst', edges[1::2]) + zstwrite(p / "00.edges.csv.zst", edges[0::2]) + zstwrite(p / "01.edges.csv.zst", edges[1::2]) nodes = [n for n in input_nodes if n.startswith(f"swh:1:{short_obj_type}")] - zstwrite(p / '00.nodes.csv.zst', nodes[0::2]) - zstwrite(p / '01.nodes.csv.zst', nodes[1::2]) + zstwrite(p / "00.nodes.csv.zst", nodes[0::2]) + zstwrite(p / "01.nodes.csv.zst", nodes[1::2]) - sort_graph_nodes(tmp_path, config={'sort_buffer_size': '1M'}) + sort_graph_nodes(tmp_path, config={"sort_buffer_size": "1M"}) output_nodes = zstread(tmp_path / "graph.nodes.csv.zst").split("\n") output_edges = zstread(tmp_path / "graph.edges.csv.zst").split("\n") output_nodes = list(filter(bool, output_nodes)) output_edges = list(filter(bool, output_edges)) expected_nodes = set(input_nodes) | set(l.split()[1] for l in input_edges) assert output_nodes == sorted(expected_nodes) - assert int((tmp_path / 'graph.nodes.count.txt').read_text()) == len(expected_nodes) + assert int((tmp_path / "graph.nodes.count.txt").read_text()) == len(expected_nodes) assert sorted(output_edges) == sorted(input_edges) - assert int((tmp_path / 'graph.edges.count.txt').read_text()) == len(input_edges) + assert int((tmp_path / "graph.edges.count.txt").read_text()) == len(input_edges) diff --git a/swh/dataset/utils.py b/swh/dataset/utils.py index 2ea1eb2..646c614 100644 --- a/swh/dataset/utils.py +++ b/swh/dataset/utils.py @@ -1,40 +1,37 @@ # Copyright (C) 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 subprocess class ZSTFile: - def __init__(self, path, mode='r'): - if mode not in ('r', 'rb', 'w', 'wb'): + def __init__(self, path, mode="r"): + if mode not in ("r", "rb", "w", "wb"): raise ValueError(f"ZSTFile mode {mode} is invalid.") self.path = path self.mode = mode def __enter__(self): - is_text = not (self.mode in ('rb', 'wb')) - writing = self.mode in ('w', 'wb') + is_text = not (self.mode in ("rb", "wb")) + writing = self.mode in ("w", "wb") if writing: - cmd = ['zstd', '-q', '-o', self.path] + cmd = ["zstd", "-q", "-o", self.path] else: - cmd = ['zstdcat', self.path] + cmd = ["zstdcat", self.path] self.process = subprocess.Popen( - cmd, - text=is_text, - stdin=subprocess.PIPE, - stdout=subprocess.PIPE, + cmd, text=is_text, stdin=subprocess.PIPE, stdout=subprocess.PIPE, ) return self def __exit__(self, exc_type, exc_value, tb): self.process.stdin.close() self.process.stdout.close() self.process.wait() def read(self, *args): return self.process.stdout.read(*args) def write(self, buf): self.process.stdin.write(buf)