1# coding=UTF-8 2 3from __future__ import absolute_import 4 5import base64 6import zlib 7 8from mercurial import ( 9 bundlecaches, 10 changegroup, 11 extensions, 12 revlog, 13 util, 14) 15from mercurial.revlogutils import flagutil 16from mercurial.interfaces import repository 17 18# Test only: These flags are defined here only in the context of testing the 19# behavior of the flag processor. The canonical way to add flags is to get in 20# touch with the community and make them known in revlog. 21REVIDX_NOOP = 1 << 3 22REVIDX_BASE64 = 1 << 2 23REVIDX_GZIP = 1 << 1 24REVIDX_FAIL = 1 25 26 27def validatehash(self, text): 28 return True 29 30 31def bypass(self, text): 32 return False 33 34 35def noopdonothing(self, text): 36 return (text, True) 37 38 39def noopdonothingread(self, text): 40 return (text, True) 41 42 43def b64encode(self, text): 44 return (base64.b64encode(text), False) 45 46 47def b64decode(self, text): 48 return (base64.b64decode(text), True) 49 50 51def gzipcompress(self, text): 52 return (zlib.compress(text), False) 53 54 55def gzipdecompress(self, text): 56 return (zlib.decompress(text), True) 57 58 59def supportedoutgoingversions(orig, repo): 60 versions = orig(repo) 61 versions.discard(b'01') 62 versions.discard(b'02') 63 versions.add(b'03') 64 return versions 65 66 67def allsupportedversions(orig, ui): 68 versions = orig(ui) 69 versions.add(b'03') 70 return versions 71 72 73def makewrappedfile(obj): 74 class wrappedfile(obj.__class__): 75 def addrevision( 76 self, 77 text, 78 transaction, 79 link, 80 p1, 81 p2, 82 cachedelta=None, 83 node=None, 84 flags=flagutil.REVIDX_DEFAULT_FLAGS, 85 ): 86 if b'[NOOP]' in text: 87 flags |= REVIDX_NOOP 88 89 if b'[BASE64]' in text: 90 flags |= REVIDX_BASE64 91 92 if b'[GZIP]' in text: 93 flags |= REVIDX_GZIP 94 95 # This addrevision wrapper is meant to add a flag we will not have 96 # transforms registered for, ensuring we handle this error case. 97 if b'[FAIL]' in text: 98 flags |= REVIDX_FAIL 99 100 return super(wrappedfile, self).addrevision( 101 text, 102 transaction, 103 link, 104 p1, 105 p2, 106 cachedelta=cachedelta, 107 node=node, 108 flags=flags, 109 ) 110 111 obj.__class__ = wrappedfile 112 113 114def reposetup(ui, repo): 115 class wrappingflagprocessorrepo(repo.__class__): 116 def file(self, f): 117 orig = super(wrappingflagprocessorrepo, self).file(f) 118 makewrappedfile(orig) 119 return orig 120 121 repo.__class__ = wrappingflagprocessorrepo 122 123 124def extsetup(ui): 125 # Enable changegroup3 for flags to be sent over the wire 126 wrapfunction = extensions.wrapfunction 127 wrapfunction( 128 changegroup, 'supportedoutgoingversions', supportedoutgoingversions 129 ) 130 wrapfunction(changegroup, 'allsupportedversions', allsupportedversions) 131 132 # Teach revlog about our test flags 133 flags = [REVIDX_NOOP, REVIDX_BASE64, REVIDX_GZIP, REVIDX_FAIL] 134 flagutil.REVIDX_KNOWN_FLAGS |= util.bitsfrom(flags) 135 repository.REVISION_FLAGS_KNOWN |= util.bitsfrom(flags) 136 revlog.REVIDX_FLAGS_ORDER.extend(flags) 137 138 # Teach exchange to use changegroup 3 139 for k in bundlecaches._bundlespeccontentopts.keys(): 140 bundlecaches._bundlespeccontentopts[k][b"cg.version"] = b"03" 141 142 # Register flag processors for each extension 143 flagutil.addflagprocessor( 144 REVIDX_NOOP, 145 ( 146 noopdonothingread, 147 noopdonothing, 148 validatehash, 149 ), 150 ) 151 flagutil.addflagprocessor( 152 REVIDX_BASE64, 153 ( 154 b64decode, 155 b64encode, 156 bypass, 157 ), 158 ) 159 flagutil.addflagprocessor( 160 REVIDX_GZIP, (gzipdecompress, gzipcompress, bypass) 161 ) 162