ext-sidedata.py
105 lines
| 3.5 KiB
| text/x-python
|
PythonLexer
r43308 | # ext-sidedata.py - small extension to test the sidedata logic | |||
# | ||||
|
r47452 | # Copyright 2019 Pierre-Yves David <pierre-yves.david@octobus.net> | ||
r43308 | # | |||
# This software may be used and distributed according to the terms of the | ||||
# GNU General Public License version 2 or any later version. | ||||
from __future__ import absolute_import | ||||
import hashlib | ||||
import struct | ||||
|
r47771 | from mercurial.node import nullrev | ||
r43308 | from mercurial import ( | |||
extensions, | ||||
|
r45933 | requirements, | ||
r43308 | revlog, | |||
) | ||||
r46661 | from mercurial.upgrade_utils import engine as upgrade_engine | |||
|
r47847 | from mercurial.revlogutils import constants | ||
|
r43346 | from mercurial.revlogutils import sidedata | ||
r43308 | ||||
|
r43346 | def wrapaddrevision( | ||
orig, self, text, transaction, link, p1, p2, *args, **kwargs | ||||
): | ||||
r43308 | if kwargs.get('sidedata') is None: | |||
kwargs['sidedata'] = {} | ||||
sd = kwargs['sidedata'] | ||||
## let's store some arbitrary data just for testing | ||||
# text length | ||||
sd[sidedata.SD_TEST1] = struct.pack('>I', len(text)) | ||||
# and sha2 hashes | ||||
sha256 = hashlib.sha256(text).digest() | ||||
sd[sidedata.SD_TEST2] = struct.pack('>32s', sha256) | ||||
return orig(self, text, transaction, link, p1, p2, *args, **kwargs) | ||||
|
r43346 | |||
|
r47443 | def wrap_revisiondata(orig, self, nodeorrev, *args, **kwargs): | ||
text, sd = orig(self, nodeorrev, *args, **kwargs) | ||||
r43404 | if getattr(self, 'sidedatanocheck', False): | |||
|
r47443 | return text, sd | ||
if self.version & 0xFFFF != 2: | ||||
return text, sd | ||||
|
r47771 | if nodeorrev != nullrev and nodeorrev != self.nullid: | ||
|
r47452 | cat1 = sd.get(sidedata.SD_TEST1) | ||
if cat1 is not None and len(text) != struct.unpack('>I', cat1)[0]: | ||||
r43310 | raise RuntimeError('text size mismatch') | |||
|
r47452 | expected = sd.get(sidedata.SD_TEST2) | ||
r43310 | got = hashlib.sha256(text).digest() | |||
|
r47452 | if expected is not None and got != expected: | ||
r43310 | raise RuntimeError('sha256 mismatch') | |||
|
r47443 | return text, sd | ||
r43310 | ||||
|
r43346 | |||
|
r47847 | def wrapget_sidedata_helpers(orig, srcrepo, dstrepo): | ||
repo, computers, removers = orig(srcrepo, dstrepo) | ||||
assert not computers and not removers # deal with composition later | ||||
r43404 | addedreqs = dstrepo.requirements - srcrepo.requirements | |||
|
r47847 | |||
|
r45933 | if requirements.SIDEDATA_REQUIREMENT in addedreqs: | ||
r43404 | ||||
|
r47847 | def computer(repo, revlog, rev, old_sidedata): | ||
assert not old_sidedata # not supported yet | ||||
r43404 | update = {} | |||
revlog.sidedatanocheck = True | ||||
try: | ||||
text = revlog.revision(rev) | ||||
finally: | ||||
del revlog.sidedatanocheck | ||||
## let's store some arbitrary data just for testing | ||||
# text length | ||||
update[sidedata.SD_TEST1] = struct.pack('>I', len(text)) | ||||
# and sha2 hashes | ||||
sha256 = hashlib.sha256(text).digest() | ||||
update[sidedata.SD_TEST2] = struct.pack('>32s', sha256) | ||||
|
r47847 | return update, (0, 0) | ||
r43404 | ||||
|
r47847 | srcrepo.register_sidedata_computer( | ||
constants.KIND_CHANGELOG, | ||||
b"whatever", | ||||
(sidedata.SD_TEST1, sidedata.SD_TEST2), | ||||
computer, | ||||
0, | ||||
) | ||||
dstrepo.register_wanted_sidedata(b"whatever") | ||||
|
r47848 | return sidedata.get_sidedata_helpers(srcrepo, dstrepo._wanted_sidedata) | ||
r43404 | ||||
r43308 | def extsetup(ui): | |||
extensions.wrapfunction(revlog.revlog, 'addrevision', wrapaddrevision) | ||||
|
r47443 | extensions.wrapfunction(revlog.revlog, '_revisiondata', wrap_revisiondata) | ||
r43404 | extensions.wrapfunction( | |||
|
r47847 | upgrade_engine, 'get_sidedata_helpers', wrapget_sidedata_helpers | ||
r43404 | ) | |||
|
r47452 | |||
def reposetup(ui, repo): | ||||
# We don't register sidedata computers because we don't care within these | ||||
# tests | ||||
repo.register_wanted_sidedata(sidedata.SD_TEST1) | ||||
repo.register_wanted_sidedata(sidedata.SD_TEST2) | ||||