swh/model/fields/hashes.py

@@ -4,6 +4,7 @@
 # See top-level LICENSE file for more information
 
 import string
+
 from ..exceptions import ValidationError
 
 
@@ -22,22 +23,22 @@ def validate_hash(value, hash_type):
     """
 
     hash_lengths = {
-        'sha1': 20,
-        'sha1_git': 20,
-        'sha256': 32,
+        "sha1": 20,
+        "sha1_git": 20,
+        "sha256": 32,
     }
 
     hex_digits = set(string.hexdigits)
 
     if hash_type not in hash_lengths:
         raise ValidationError(
-            'Unexpected hash type %(hash_type)s, expected one of'
-            ' %(hash_types)s',
+            "Unexpected hash type %(hash_type)s, expected one of" " %(hash_types)s",
             params={
-                'hash_type': hash_type,
-                'hash_types': ', '.join(sorted(hash_lengths)),
+                "hash_type": hash_type,
+                "hash_types": ", ".join(sorted(hash_lengths)),
             },
-            code='unexpected-hash-type')
+            code="unexpected-hash-type",
+        )
 
     if isinstance(value, str):
         errors = []
@@ -48,10 +49,10 @@ def validate_hash(value, hash_type):
                     "Unexpected characters `%(unexpected_chars)s' for hash "
                     "type %(hash_type)s",
                     params={
-                        'unexpected_chars': ', '.join(sorted(extra_chars)),
-                        'hash_type': hash_type,
+                        "unexpected_chars": ", ".join(sorted(extra_chars)),
+                        "hash_type": hash_type,
                     },
-                    code='unexpected-hash-contents',
+                    code="unexpected-hash-contents",
                 )
             )
 
@@ -60,14 +61,14 @@ def validate_hash(value, hash_type):
         if length != expected_length:
             errors.append(
                 ValidationError(
-                    'Unexpected length %(length)d for hash type '
-                    '%(hash_type)s, expected %(expected_length)d',
+                    "Unexpected length %(length)d for hash type "
+                    "%(hash_type)s, expected %(expected_length)d",
                     params={
-                        'length': length,
-                        'expected_length': expected_length,
-                        'hash_type': hash_type,
+                        "length": length,
+                        "expected_length": expected_length,
+                        "hash_type": hash_type,
                     },
-                    code='unexpected-hash-length',
+                    code="unexpected-hash-length",
                 )
             )
 
@@ -81,37 +82,37 @@ def validate_hash(value, hash_type):
         expected_length = hash_lengths[hash_type]
         if length != expected_length:
             raise ValidationError(
-                'Unexpected length %(length)d for hash type '
-                '%(hash_type)s, expected %(expected_length)d',
+                "Unexpected length %(length)d for hash type "
+                "%(hash_type)s, expected %(expected_length)d",
                 params={
-                    'length': length,
-                    'expected_length': expected_length,
-                    'hash_type': hash_type,
+                    "length": length,
+                    "expected_length": expected_length,
+                    "hash_type": hash_type,
                 },
-                code='unexpected-hash-length',
+                code="unexpected-hash-length",
             )
 
         return True
 
     raise ValidationError(
-        'Unexpected type %(type)s for hash, expected str or bytes',
+        "Unexpected type %(type)s for hash, expected str or bytes",
         params={
-            'type': value.__class__.__name__,
+            "type": value.__class__.__name__,
         },
-        code='unexpected-hash-value-type',
+        code="unexpected-hash-value-type",
     )
 
 
 def validate_sha1(sha1):
     """Validate that sha1 is a valid sha1 hash"""
-    return validate_hash(sha1, 'sha1')
+    return validate_hash(sha1, "sha1")
 
 
 def validate_sha1_git(sha1_git):
     """Validate that sha1_git is a valid sha1_git hash"""
-    return validate_hash(sha1_git, 'sha1_git')
+    return validate_hash(sha1_git, "sha1_git")
 
 
 def validate_sha256(sha256):
     """Validate that sha256 is a valid sha256 hash"""
-    return validate_hash(sha256, 'sha256')
+    return validate_hash(sha256, "sha256")

swh/model/fields/simple.py

@@ -13,16 +13,16 @@ def validate_type(value, type):
     """Validate that value is an integer"""
     if not isinstance(value, type):
         if isinstance(type, tuple):
-            typestr = 'one of %s' % ', '.join(typ.__name__ for typ in type)
+            typestr = "one of %s" % ", ".join(typ.__name__ for typ in type)
         else:
             typestr = type.__name__
         raise ValidationError(
-            'Unexpected type %(type)s, expected %(expected_type)s',
+            "Unexpected type %(type)s, expected %(expected_type)s",
             params={
-                'type': value.__class__.__name__,
-                'expected_type': typestr,
+                "type": value.__class__.__name__,
+                "expected_type": typestr,
             },
-            code='unexpected-type'
+            code="unexpected-type",
         )
 
     return True
@@ -54,10 +54,12 @@ def validate_datetime(value):
         errors.append(e)
 
     if isinstance(value, datetime.datetime) and value.tzinfo is None:
-        errors.append(ValidationError(
-            'Datetimes must be timezone-aware in swh',
-            code='datetime-without-tzinfo',
-        ))
+        errors.append(
+            ValidationError(
+                "Datetimes must be timezone-aware in swh",
+                code="datetime-without-tzinfo",
+            )
+        )
 
     if errors:
         raise ValidationError(errors)
@@ -69,12 +71,12 @@ def validate_enum(value, expected_values):
     """Validate that value is contained in expected_values"""
     if value not in expected_values:
         raise ValidationError(
-            'Unexpected value %(value)s, expected one of %(expected_values)s',
+            "Unexpected value %(value)s, expected one of %(expected_values)s",
             params={
-                'value': value,
-                'expected_values': ', '.join(sorted(expected_values)),
+                "value": value,
+                "expected_values": ", ".join(sorted(expected_values)),
             },
-            code='unexpected-value',
+            code="unexpected-value",
         )
 
     return True

swh/model/hashutil.py

-# Copyright (C) 2015-2017  The Software Heritage developers
+# 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
@@ -10,31 +10,60 @@ Only a subset of hashing algorithms is supported as defined in the
 ALGORITHMS set. Any provided algorithms not in that list will result
 in a ValueError explaining the error.
 
-This modules defines the following hashing functions:
+This module defines a MultiHash class to ease the softwareheritage
+hashing algorithms computation. This allows to compute hashes from
+file object, path, data using a similar interface as what the standard
+hashlib module provides.
 
-- hash_file: Hash the contents of the given file object with the given
-  algorithms (defaulting to DEFAULT_ALGORITHMS if none provided).
+Basic usage examples:
 
-- hash_data: Hash the given binary blob with the given algorithms
-  (defaulting to DEFAULT_ALGORITHMS if none provided).
+- file object: MultiHash.from_file(
+                 file_object, hash_names=DEFAULT_ALGORITHMS).digest()
+
+- path (filepath): MultiHash.from_path(b'foo').hexdigest()
+
+- data (bytes): MultiHash.from_data(b'foo').bytehexdigest()
+
+
+"Complex" usage, defining a swh hashlib instance first:
+
+- To compute length, integrate the length to the set of algorithms to
+  compute, for example:
+
+  .. code-block:: python
+
+     h = MultiHash(hash_names=set({'length'}).union(DEFAULT_ALGORITHMS))
+     with open(filepath, 'rb') as f:
+         h.update(f.read(HASH_BLOCK_SIZE))
+     hashes = h.digest()  # returns a dict of {hash_algo_name: hash_in_bytes}
+
+- Write alongside computing hashing algorithms (from a stream), example:
+
+  .. code-block:: python
+
+     h = MultiHash(length=length)
+     with open(filepath, 'wb') as f:
+         for chunk in r.iter_content():  # r a stream of sort
+             h.update(chunk)
+             f.write(chunk)
+     hashes = h.hexdigest()  # returns a dict of {hash_algo_name: hash_in_hex}
 
-- hash_path: Hash the contents of the file at the given path with the
-  given algorithms (defaulting to DEFAULT_ALGORITHMS if none
-  provided).
 
 """
 
 import binascii
 import functools
 import hashlib
-import os
-
 from io import BytesIO
+import os
+from typing import Callable, Dict, Optional, Union
 
-ALGORITHMS = set(['sha1', 'sha256', 'sha1_git', 'blake2s256', 'blake2b512'])
+ALGORITHMS = set(
+    ["sha1", "sha256", "sha1_git", "blake2s256", "blake2b512", "md5", "sha512"]
+)
 """Hashing algorithms supported by this module"""
 
-DEFAULT_ALGORITHMS = set(['sha1', 'sha256', 'sha1_git', 'blake2s256'])
+DEFAULT_ALGORITHMS = set(["sha1", "sha256", "sha1_git", "blake2s256"])
 """Algorithms computed by default when calling the functions from this module.
 
 Subset of :const:`ALGORITHMS`.
@@ -43,23 +72,138 @@ Subset of :const:`ALGORITHMS`.
 HASH_BLOCK_SIZE = 32768
 """Block size for streaming hash computations made in this module"""
 
-# Load blake2 hashes from pyblake2 if they are not available in the builtin
-# hashlib
-__pyblake2_hashes = {'blake2s256': 'blake2s',
-                     'blake2b512': 'blake2b'}
-__cache = hashlib.__builtin_constructor_cache
-for __hash, __pyblake2_fn in __pyblake2_hashes.items():
-    if __hash not in hashlib.algorithms_available:
-        import pyblake2
-        __cache[__hash] = getattr(pyblake2, __pyblake2_fn)
+_blake2_hash_cache: Dict[str, Callable] = {}
+
 
+class MultiHash:
+    """Hashutil class to support multiple hashes computation.
 
-def _new_git_hash(base_algo, git_type, length):
-    """Initialize a digest object (as returned by python's hashlib) for the
-    requested algorithm, and feed it with the header for a git object of the
-    given type and length.
+    Args:
 
-    The header for hashing a git object consists of:
+        hash_names (set): Set of hash algorithms (+ optionally length)
+                          to compute hashes (cf. DEFAULT_ALGORITHMS)
+        length (int): Length of the total sum of chunks to read
+
+    If the length is provided as algorithm, the length is also
+    computed and returned.
+
+    """
+
+    def __init__(self, hash_names=DEFAULT_ALGORITHMS, length=None):
+        self.state = {}
+        self.track_length = False
+        for name in hash_names:
+            if name == "length":
+                self.state["length"] = 0
+                self.track_length = True
+            else:
+                self.state[name] = _new_hash(name, length)
+
+    @classmethod
+    def from_state(cls, state, track_length):
+        ret = cls([])
+        ret.state = state
+        ret.track_length = track_length
+
+    @classmethod
+    def from_file(cls, fobj, hash_names=DEFAULT_ALGORITHMS, length=None):
+        ret = cls(length=length, hash_names=hash_names)
+        while True:
+            chunk = fobj.read(HASH_BLOCK_SIZE)
+            if not chunk:
+                break
+            ret.update(chunk)
+        return ret
+
+    @classmethod
+    def from_path(cls, path, hash_names=DEFAULT_ALGORITHMS):
+        length = os.path.getsize(path)
+        with open(path, "rb") as f:
+            ret = cls.from_file(f, hash_names=hash_names, length=length)
+        return ret
+
+    @classmethod
+    def from_data(cls, data, hash_names=DEFAULT_ALGORITHMS):
+        length = len(data)
+        fobj = BytesIO(data)
+        return cls.from_file(fobj, hash_names=hash_names, length=length)
+
+    def update(self, chunk):
+        for name, h in self.state.items():
+            if name == "length":
+                continue
+            h.update(chunk)
+        if self.track_length:
+            self.state["length"] += len(chunk)
+
+    def digest(self):
+        return {
+            name: h.digest() if name != "length" else h
+            for name, h in self.state.items()
+        }
+
+    def hexdigest(self):
+        return {
+            name: h.hexdigest() if name != "length" else h
+            for name, h in self.state.items()
+        }
+
+    def bytehexdigest(self):
+        return {
+            name: hash_to_bytehex(h.digest()) if name != "length" else h
+            for name, h in self.state.items()
+        }
+
+    def copy(self):
+        copied_state = {
+            name: h.copy() if name != "length" else h for name, h in self.state.items()
+        }
+        return self.from_state(copied_state, self.track_length)
+
+
+def _new_blake2_hash(algo):
+    """Return a function that initializes a blake2 hash."""
+    if algo in _blake2_hash_cache:
+        return _blake2_hash_cache[algo]()
+
+    lalgo = algo.lower()
+    if not lalgo.startswith("blake2"):
+        raise ValueError("Algorithm %s is not a blake2 hash" % algo)
+
+    blake_family = lalgo[:7]
+
+    digest_size = None
+    if lalgo[7:]:
+        try:
+            digest_size, remainder = divmod(int(lalgo[7:]), 8)
+        except ValueError:
+            raise ValueError("Unknown digest size for algo %s" % algo) from None
+        if remainder:
+            raise ValueError(
+                "Digest size for algorithm %s must be a multiple of 8" % algo
+            )
+
+    blake2 = getattr(hashlib, blake_family)
+    _blake2_hash_cache[algo] = lambda: blake2(digest_size=digest_size)
+
+    return _blake2_hash_cache[algo]()
+
+
+def _new_hashlib_hash(algo):
+    """Initialize a digest object from hashlib.
+
+    Handle the swh-specific names for the blake2-related algorithms
+    """
+    if algo.startswith("blake2"):
+        return _new_blake2_hash(algo)
+    else:
+        return hashlib.new(algo)
+
+
+def git_object_header(git_type: str, length: int) -> bytes:
+    """Returns the header for a git object of the given type and length.
+
+    The header of a git object consists of:
      - The type of the object (encoded in ASCII)
      - One ASCII space (\x20)
      - The length of the object (decimal encoded in ASCII)
@@ -74,15 +218,26 @@ def _new_git_hash(base_algo, git_type, length):
     Returns:
         a hashutil.hash object
     """
+    git_object_types = {
+        "blob",
+        "tree",
+        "commit",
+        "tag",
+        "snapshot",
+        "raw_extrinsic_metadata",
+        "extid",
+    }
 
-    h = hashlib.new(base_algo)
-    git_header = '%s %d\0' % (git_type, length)
-    h.update(git_header.encode('ascii'))
+    if git_type not in git_object_types:
+        raise ValueError(
+            "Unexpected git object type %s, expected one of %s"
+            % (git_type, ", ".join(sorted(git_object_types)))
+        )
 
-    return h
+    return ("%s %d\0" % (git_type, length)).encode("ascii")
 
 
-def _new_hash(algo, length=None):
+def _new_hash(algo: str, length: Optional[int] = None):
     """Initialize a digest object (as returned by python's hashlib) for
     the requested algorithm. See the constant ALGORITHMS for the list
     of supported algorithms. If a git-specific hashing algorithm is
@@ -104,87 +259,22 @@ def _new_hash(algo, length=None):
     """
     if algo not in ALGORITHMS:
         raise ValueError(
-            'Unexpected hashing algorithm %s, expected one of %s' %
-            (algo, ', '.join(sorted(ALGORITHMS))))
+            "Unexpected hashing algorithm %s, expected one of %s"
+            % (algo, ", ".join(sorted(ALGORITHMS)))
+        )
 
-    if algo.endswith('_git'):
+    if algo.endswith("_git"):
         if length is None:
-            raise ValueError('Missing length for git hashing algorithm')
+            raise ValueError("Missing length for git hashing algorithm")
         base_algo = algo[:-4]
-        return _new_git_hash(base_algo, 'blob', length)
-
-    return hashlib.new(algo)
-
-
-def hash_file(fobj, length=None, algorithms=DEFAULT_ALGORITHMS, chunk_cb=None):
-    """Hash the contents of the given file object with the given algorithms.
-
-    Args:
-        fobj: a file-like object
-        length: the length of the contents of the file-like object (for the
-                git-specific algorithms)
-        algorithms: the hashing algorithms used
-
-    Returns: a dict mapping each algorithm to a bytes digest.
-
-    Raises:
-        ValueError if algorithms contains an unknown hash algorithm.
-    """
-    hashes = {algo: _new_hash(algo, length) for algo in algorithms}
-
-    while True:
-        chunk = fobj.read(HASH_BLOCK_SIZE)
-        if not chunk:
-            break
-        for hash in hashes.values():
-            hash.update(chunk)
-        if chunk_cb:
-            chunk_cb(chunk)
-
-    return {algo: hash.digest() for algo, hash in hashes.items()}
-
+        h = _new_hashlib_hash(base_algo)
+        h.update(git_object_header("blob", length))
+        return h
 
-def hash_path(path, algorithms=DEFAULT_ALGORITHMS, chunk_cb=None):
-    """Hash the contents of the file at the given path with the given
-       algorithms.
-
-    Args:
-        path: the path of the file to hash
-        algorithms: the hashing algorithms used
-        chunk_cb: a callback
-
-    Returns: a dict mapping each algorithm to a bytes digest.
-
-    Raises:
-        ValueError if algorithms contains an unknown hash algorithm.
-        OSError on file access error
-
-    """
-    length = os.path.getsize(path)
-    with open(path, 'rb') as fobj:
-        hash = hash_file(fobj, length, algorithms, chunk_cb)
-    hash['length'] = length
-    return hash
+    return _new_hashlib_hash(algo)
 
 
-def hash_data(data, algorithms=DEFAULT_ALGORITHMS):
-    """Hash the given binary blob with the given algorithms.
-
-    Args:
-        data: a bytes object
-        algorithms: the hashing algorithms used
-
-    Returns: a dict mapping each algorithm to a bytes digest
-
-    Raises:
-        TypeError if data does not support the buffer interface.
-        ValueError if algorithms contains an unknown hash algorithm.
-    """
-    fobj = BytesIO(data)
-    return hash_file(fobj, len(data), algorithms)
-
-
-def hash_git_data(data, git_type, base_algo='sha1'):
+def hash_git_data(data, git_type, base_algo="sha1"):
     """Hash the given data as a git object of type git_type.
 
     Args:
@@ -197,21 +287,15 @@ def hash_git_data(data, git_type, base_algo='sha1'):
     Raises:
         ValueError if the git_type is unexpected.
     """
-
-    git_object_types = {'blob', 'tree', 'commit', 'tag', 'snapshot'}
-
-    if git_type not in git_object_types:
-        raise ValueError('Unexpected git object type %s, expected one of %s' %
-                         (git_type, ', '.join(sorted(git_object_types))))
-
-    h = _new_git_hash(base_algo, git_type, len(data))
+    h = _new_hashlib_hash(base_algo)
+    h.update(git_object_header(git_type, len(data)))
     h.update(data)
 
     return h.digest()
 
 
 @functools.lru_cache()
-def hash_to_hex(hash):
+def hash_to_hex(hash: Union[str, bytes]) -> str:
     """Converts a hash (in hex or bytes form) to its hexadecimal ascii form
 
     Args:
@@ -223,11 +307,11 @@ def hash_to_hex(hash):
     """
     if isinstance(hash, str):
         return hash
-    return binascii.hexlify(hash).decode('ascii')
+    return binascii.hexlify(hash).decode("ascii")
 
 
 @functools.lru_cache()
-def hash_to_bytehex(hash):
+def hash_to_bytehex(hash: bytes) -> bytes:
     """Converts a hash to its hexadecimal bytes representation
 
     Args:
@@ -240,7 +324,7 @@ def hash_to_bytehex(hash):
 
 
 @functools.lru_cache()
-def hash_to_bytes(hash):
+def hash_to_bytes(hash: Union[str, bytes]) -> bytes:
     """Converts a hash (in hex or bytes form) to its raw bytes form
 
     Args:
@@ -256,7 +340,7 @@ def hash_to_bytes(hash):
 
 
 @functools.lru_cache()
-def bytehex_to_hash(hex):
+def bytehex_to_hash(hex: bytes) -> bytes:
     """Converts a hexadecimal bytes representation of a hash to that hash
 
     Args:

swh/model/merkle.py

-# Copyright (C) 2017 The Software Heritage developers
+# Copyright (C) 2017-2022 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
 
 """Merkle tree data structure"""
 
-import abc
-import collections
-
-
-def deep_update(left, right):
-    """Recursively update the left mapping with deeply nested values from the right
-    mapping.
-
-    This function is useful to merge the results of several calls to
-    :func:`MerkleNode.collect`.
-
-    Arguments:
-      left: a mapping (modified by the update operation)
-      right: a mapping
-
-    Returns:
-      the left mapping, updated with nested values from the right mapping
-
-    Example:
-        >>> a = {
-        ...     'key1': {
-        ...         'key2': {
-        ...              'key3': 'value1/2/3',
-        ...         },
-        ...     },
-        ... }
-        >>> deep_update(a, {
-        ...     'key1': {
-        ...         'key2': {
-        ...              'key4': 'value1/2/4',
-        ...         },
-        ...     },
-        ... }) == {
-        ...     'key1': {
-        ...         'key2': {
-        ...             'key3': 'value1/2/3',
-        ...             'key4': 'value1/2/4',
-        ...         },
-        ...     },
-        ... }
-        True
-        >>> deep_update(a, {
-        ...     'key1': {
-        ...         'key2': {
-        ...              'key3': 'newvalue1/2/3',
-        ...         },
-        ...     },
-        ... }) == {
-        ...     'key1': {
-        ...         'key2': {
-        ...             'key3': 'newvalue1/2/3',
-        ...             'key4': 'value1/2/4',
-        ...         },
-        ...     },
-        ... }
-        True
+from __future__ import annotations
 
-    """
-    for key, rvalue in right.items():
-        if isinstance(rvalue, collections.Mapping):
-            new_lvalue = deep_update(left.get(key, {}), rvalue)
-            left[key] = new_lvalue
-        else:
-            left[key] = rvalue
-    return left
+import abc
+from typing import Any, Dict, Iterator, List, Set
 
 
 class MerkleNode(dict, metaclass=abc.ABCMeta):
@@ -100,16 +39,18 @@ class MerkleNode(dict, metaclass=abc.ABCMeta):
     The collection of updated data from the tree is implemented through the
     :func:`collect` function and associated helpers.
 
-    Attributes:
-      data (dict): data associated to the current node
-      parents (list): known parents of the current node
-      collected (bool): whether the current node has been collected
-
     """
-    __slots__ = ['parents', 'data', '__hash', 'collected']
 
-    type = None
-    """Type of the current node (used as a classifier for :func:`collect`)"""
+    __slots__ = ["parents", "data", "__hash", "collected"]
+
+    data: Dict
+    """data associated to the current node"""
+
+    parents: List
+    """known parents of the current node"""
+
+    collected: bool
+    """whether the current node has been collected"""
 
     def __init__(self, data=None):
         super().__init__()
@@ -118,6 +59,16 @@ class MerkleNode(dict, metaclass=abc.ABCMeta):
         self.__hash = None
         self.collected = False
 
+    def __eq__(self, other):
+        return (
+            isinstance(other, MerkleNode)
+            and super().__eq__(other)
+            and self.data == other.data
+        )
+
+    def __ne__(self, other):
+        return not self.__eq__(other)
+
     def invalidate_hash(self):
         """Invalidate the cached hash of the current node."""
         if not self.__hash:
@@ -128,7 +79,7 @@ class MerkleNode(dict, metaclass=abc.ABCMeta):
         for parent in self.parents:
             parent.invalidate_hash()
 
-    def update_hash(self, *, force=False):
+    def update_hash(self, *, force=False) -> Any:
         """Recursively compute the hash of the current node.
 
         Args:
@@ -148,20 +99,23 @@ class MerkleNode(dict, metaclass=abc.ABCMeta):
         return self.__hash
 
     @property
-    def hash(self):
+    def hash(self) -> Any:
         """The hash of the current node, as calculated by
         :func:`compute_hash`.
         """
         return self.update_hash()
 
+    def __hash__(self):
+        return hash(self.hash)
+
     @abc.abstractmethod
-    def compute_hash(self):
+    def compute_hash(self) -> Any:
         """Compute the hash of the current node.
 
         The hash should depend on the data of the node, as well as on hashes
         of the children nodes.
         """
-        raise NotImplementedError('Must implement compute_hash method')
+        raise NotImplementedError("Must implement compute_hash method")
 
     def __setitem__(self, name, new_child):
         """Add a child, invalidating the current hash"""
@@ -210,47 +164,24 @@ class MerkleNode(dict, metaclass=abc.ABCMeta):
         """
         return self.data
 
-    def collect_node(self, **kwargs):
-        """Collect the data for the current node, for use by :func:`collect`.
-
-        Arguments:
-          kwargs: passed as-is to :func:`get_data`.
-
-        Returns:
-          A :class:`dict` compatible with :func:`collect`.
-        """
+    def collect_node(self) -> Set[MerkleNode]:
+        """Collect the current node if it has not been yet, for use by :func:`collect`."""
         if not self.collected:
             self.collected = True
-            return {self.type: {self.hash: self.get_data(**kwargs)}}
+            return {self}
         else:
-            return {}
-
-    def collect(self, **kwargs):
-        """Collect the data for all nodes in the subtree rooted at `self`.
+            return set()
 
-        The data is deduplicated by type and by hash.
-
-        Arguments:
-          kwargs: passed as-is to :func:`get_data`.
+    def collect(self) -> Set[MerkleNode]:
+        """Collect the added and modified nodes in the subtree rooted at `self`
+        since the last collect operation.
 
         Returns:
-           A :class:`dict` with the following structure::
-
-             {
-               'typeA': {
-                 node1.hash: node1.get_data(),
-                 node2.hash: node2.get_data(),
-               },
-               'typeB': {
-                 node3.hash: node3.get_data(),
-                 ...
-               },
-               ...
-             }
+           A :class:`set` of collected nodes
         """
-        ret = self.collect_node(**kwargs)
+        ret = self.collect_node()
         for child in self.values():
-            deep_update(ret, child.collect(**kwargs))
+            ret.update(child.collect())
 
         return ret
 
@@ -264,23 +195,39 @@ class MerkleNode(dict, metaclass=abc.ABCMeta):
         for child in self.values():
             child.reset_collect()
 
+    def iter_tree(self, dedup=True) -> Iterator[MerkleNode]:
+        """Yields all children nodes, recursively. Common nodes are deduplicated
+        by default (deduplication can be turned off setting the given argument
+        'dedup' to False).
+        """
+        yield from self._iter_tree(seen=set(), dedup=dedup)
+
+    def _iter_tree(self, seen: Set[bytes], dedup) -> Iterator[MerkleNode]:
+        if self.hash not in seen:
+            if dedup:
+                seen.add(self.hash)
+            yield self
+            for child in self.values():
+                yield from child._iter_tree(seen=seen, dedup=dedup)
+
 
 class MerkleLeaf(MerkleNode):
     """A leaf to a Merkle tree.
 
     A Merkle leaf is simply a Merkle node with children disabled.
     """
-    __slots__ = []
+
+    __slots__: List[str] = []
 
     def __setitem__(self, name, child):
-        raise ValueError('%s is a leaf' % self.__class__.__name__)
+        raise ValueError("%s is a leaf" % self.__class__.__name__)
 
     def __getitem__(self, name):
-        raise ValueError('%s is a leaf' % self.__class__.__name__)
+        raise ValueError("%s is a leaf" % self.__class__.__name__)
 
     def __delitem__(self, name):
-        raise ValueError('%s is a leaf' % self.__class__.__name__)
+        raise ValueError("%s is a leaf" % self.__class__.__name__)
 
     def update(self, new_children):
         """Children update operation. Disabled for leaves."""
-        raise ValueError('%s is a leaf' % self.__class__.__name__)
+        raise ValueError("%s is a leaf" % self.__class__.__name__)

swh/model/py.typed

+# Marker file for PEP 561.

swh/model/tests/generate_testdata.py

+# Copyright (C) 2019 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 datetime import datetime
+from random import choice, randint, random, shuffle
+from typing import Dict, List
+
+from pytz import all_timezones, timezone
+
+from swh.model.hashutil import MultiHash
+
+PROTOCOLS = ["git", "http", "https", "deb", "svn", "mock"]
+DOMAINS = ["example.com", "some.long.host.name", "xn--n28h.tld"]
+PATHS = [
+    "",
+    "/",
+    "/stuff",
+    "/stuff/",
+    "/path/to/resource",
+    "/path/with/anchor#id=42",
+    "/path/with/qargs?q=1&b",
+]
+CONTENT_STATUS = ["visible", "hidden", "absent"]
+MAX_DATE = 3e9  # around 2065
+
+
+def gen_all_origins():
+    for protocol in PROTOCOLS:
+        for domain in DOMAINS:
+            for urlpath in PATHS:
+                yield {"url": "%s://%s%s" % (protocol, domain, urlpath)}
+
+
+ORIGINS = list(gen_all_origins())
+
+
+def gen_origins(n: int = 100) -> List:
+    """Returns a list of n randomly generated origins suitable for using as
+    Storage.add_origin() argument.
+
+    """
+    origins = ORIGINS[:]
+    shuffle(origins)
+    return origins[:n]
+
+
+def gen_content():
+    size = randint(1, 10 * 1024)
+    data = bytes(randint(0, 255) for i in range(size))
+    status = choice(CONTENT_STATUS)
+    h = MultiHash.from_data(data)
+    ctime = datetime.fromtimestamp(random() * MAX_DATE, timezone(choice(all_timezones)))
+    content = {
+        "data": data,
+        "status": status,
+        "length": size,
+        "ctime": ctime,
+        **h.digest(),
+    }
+    if status == "absent":
+        content["reason"] = "why not"
+        content["data"] = None
+    return content
+
+
+def gen_contents(n=20) -> List[Dict]:
+    """Returns a list of n randomly generated content objects (as dict) suitable
+    for using as Storage.content_add() argument.
+    """
+    return [gen_content() for i in range(n)]

swh/model/tests/test_collections.py

+# 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 pytest
+
+from swh.model.collections import ImmutableDict
+
+
+def test_immutabledict_empty():
+    d = ImmutableDict()
+
+    assert d == {}
+    assert d != {"foo": "bar"}
+
+    assert list(d) == []
+    assert list(d.items()) == []
+
+
+def test_immutabledict_one_item():
+    d = ImmutableDict({"foo": "bar"})
+
+    assert d == {"foo": "bar"}
+    assert d != {}
+
+    assert d["foo"] == "bar"
+    with pytest.raises(KeyError, match="bar"):
+        d["bar"]
+
+    assert list(d) == ["foo"]
+    assert list(d.items()) == [("foo", "bar")]
+
+
+def test_immutabledict_from_iterable():
+    d1 = ImmutableDict()
+    d2 = ImmutableDict({"foo": "bar"})
+
+    assert ImmutableDict([]) == d1
+    assert ImmutableDict([("foo", "bar")]) == d2
+
+
+def test_immutabledict_from_immutabledict():
+    d1 = ImmutableDict()
+    d2 = ImmutableDict({"foo": "bar"})
+
+    assert ImmutableDict(d1) == d1
+    assert ImmutableDict(d2) == d2
+
+
+def test_immutabledict_immutable():
+    d = ImmutableDict({"foo": "bar"})
+
+    with pytest.raises(TypeError, match="item assignment"):
+        d["bar"] = "baz"
+
+    with pytest.raises(TypeError, match="item deletion"):
+        del d["foo"]
+
+
+def test_immutabledict_copy_pop():
+    d = ImmutableDict({"foo": "bar", "baz": "qux"})
+
+    assert d.copy_pop("foo") == ("bar", ImmutableDict({"baz": "qux"}))
+
+    assert d.copy_pop("not a key") == (None, d)
+
+
+def test_hash():
+    assert hash(ImmutableDict()) == hash(ImmutableDict({}))
+    assert hash(ImmutableDict({"foo": "bar"})) == hash(ImmutableDict({"foo": "bar"}))
+    assert hash(ImmutableDict({"foo": "bar", "baz": "qux"})) == hash(
+        ImmutableDict({"foo": "bar", "baz": "qux"})
+    )
+    assert hash(ImmutableDict({"foo": "bar", "baz": "qux"})) == hash(
+        ImmutableDict({"baz": "qux", "foo": "bar"})
+    )
+
+
+def test_equality_order():
+    assert ImmutableDict({"foo": "bar", "baz": "qux"}) == ImmutableDict(
+        {"foo": "bar", "baz": "qux"}
+    )
+    assert ImmutableDict({"foo": "bar", "baz": "qux"}) == ImmutableDict(
+        {"baz": "qux", "foo": "bar"}
+    )