hgext3rd/topic/discovery.py
branchmercurial-4.6
changeset 5368 844b1ad5b34b
parent 5276 8431bb224862
--- a/hgext3rd/topic/discovery.py	Wed Apr 08 01:17:01 2020 +0800
+++ b/hgext3rd/topic/discovery.py	Fri May 08 22:50:09 2020 +0800
@@ -17,15 +17,9 @@
     compat,
 )
 
-# hg <= 4.5 (b4d85bc122bd)
-try:
-    from mercurial import wireproto  # pytype: disable=import-error
-    wireproto.branchmap
-except (AttributeError, ImportError):
-    from mercurial import wireprotov1server as wireproto
+from mercurial import wireprotov1server
 
 def _headssummary(orig, pushop, *args, **kwargs):
-    # In mercurial > 4.3, we receive the pushop as arguments
     repo = pushop.repo.unfiltered()
     remote = pushop.remote
 
@@ -205,12 +199,13 @@
         return
     tr._prepushheads = _nbheads(op.repo)
     reporef = weakref.ref(op.repo)
-    if util.safehasattr(tr, 'validator'): # hg <= 4.7
+    if util.safehasattr(tr, 'validator'): # hg <= 4.7 (ebbba3ba3f66)
         oldvalidator = tr.validator
-    else:
+    elif util.safehasattr(tr, '_validator'):
+        # hg <= 5.3 (36f08ae87ef6)
         oldvalidator = tr._validator
 
-    def validator(tr):
+    def _validate(tr):
         repo = reporef()
         if repo is not None:
             repo.invalidatecaches()
@@ -225,11 +220,19 @@
                     msg = _(b'push create more than 1 head on new branch "%s"'
                             % branch)
                     raise error.Abort(msg)
+
+    def validator(tr):
+        _validate(tr)
         return oldvalidator(tr)
-    if util.safehasattr(tr, 'validator'): # hg <= 4.7
+
+    if util.safehasattr(tr, 'validator'): # hg <= 4.7 (ebbba3ba3f66)
         tr.validator = validator
+    elif util.safehasattr(tr, '_validator'):
+        # hg <= 5.3 (36f08ae87ef6)
+        tr._validator = validator
     else:
-        tr._validator = validator
+        tr.addvalidator(b'000-new-head-check', _validate)
+
 handlecheckheads.params = frozenset()
 
 def _pushb2phases(orig, pushop, bundler):
@@ -249,8 +252,8 @@
 def modsetup(ui):
     """run at uisetup time to install all destinations wrapping"""
     extensions.wrapfunction(discovery, '_headssummary', _headssummary)
-    extensions.wrapfunction(wireproto, 'branchmap', wireprotobranchmap)
-    extensions.wrapfunction(wireproto, '_capabilities', wireprotocaps)
+    extensions.wrapfunction(wireprotov1server, 'branchmap', wireprotobranchmap)
+    extensions.wrapfunction(wireprotov1server, '_capabilities', wireprotocaps)
     # we need a proper wrap b2 part stuff
     extensions.wrapfunction(bundle2, 'handlecheckheads', handlecheckheads)
     bundle2.handlecheckheads.params = frozenset()