tests/testlib/ext-sidedata-5.py

dirstate: deal with read-race for python code using rust object

If we cannot read the dirstate data, this is probably because a writing process
wrote it under our feet. So refresh the docket and try again a handful of time.

# coding: utf8
# ext-sidedata-5.py - small extension to test (differently still) the sidedata
# logic
#
# Simulates a server for a simple sidedata exchange.
#
# Copyright 2021 Raphaël Gomès <[e-mail]>
#
# This software may be used and distributed according to the terms of the
# GNU General Public License version 2 or any later version.


import hashlib
import struct

from mercurial import (
    extensions,
    revlog,
)


from mercurial.revlogutils import sidedata as sidedatamod
from mercurial.revlogutils import constants

NO_FLAGS = (0, 0)


def compute_sidedata_1(repo, revlog, rev, sidedata, text=None):
    sidedata = sidedata.copy()
    if text is None:
        text = revlog.revision(rev)
    sidedata[sidedatamod.SD_TEST1] = struct.pack('>I', len(text))
    return sidedata, NO_FLAGS


def compute_sidedata_2(repo, revlog, rev, sidedata, text=None):
    sidedata = sidedata.copy()
    if text is None:
        text = revlog.revision(rev)
    sha256 = hashlib.sha256(text).digest()
    sidedata[sidedatamod.SD_TEST2] = struct.pack('>32s', sha256)
    return sidedata, NO_FLAGS


def reposetup(ui, repo):
    # Sidedata keys happen to be the same as the categories, easier for testing.
    for kind in constants.ALL_KINDS:
        repo.register_sidedata_computer(
            kind,
            sidedatamod.SD_TEST1,
            (sidedatamod.SD_TEST1,),
            compute_sidedata_1,
            0,
        )
        repo.register_sidedata_computer(
            kind,
            sidedatamod.SD_TEST2,
            (sidedatamod.SD_TEST2,),
            compute_sidedata_2,
            0,
        )

    # We don't register sidedata computers because we don't care within these
    # tests
    repo.register_wanted_sidedata(sidedatamod.SD_TEST1)
    repo.register_wanted_sidedata(sidedatamod.SD_TEST2)


def wrapaddrevision(
    orig, self, text, transaction, link, p1, p2, *args, **kwargs
):
    if kwargs.get('sidedata') is None:
        kwargs['sidedata'] = {}
    sd = kwargs['sidedata']
    ## let's store some arbitrary data just for testing
    # text length
    sd[sidedatamod.SD_TEST1] = struct.pack('>I', len(text))
    # and sha2 hashes
    sha256 = hashlib.sha256(text).digest()
    sd[sidedatamod.SD_TEST2] = struct.pack('>32s', sha256)
    return orig(self, text, transaction, link, p1, p2, *args, **kwargs)


def extsetup(ui):
    extensions.wrapfunction(revlog.revlog, 'addrevision', wrapaddrevision)