changeset 1944:daad8249d5cf

discovery: move all setup into a 'modsetup' function

The various wrappings are moved into a function that live in the module and the
function is now properly called during 'uisetup'.

--- a/hgext3rd/topic/__init__.py	Wed Mar 30 03:22:50 2016 -0700
+++ b/hgext3rd/topic/__init__.py	Wed Mar 30 21:53:55 2016 -0700
@@ -18,14 +18,11 @@
 from mercurial.i18n import _
 from mercurial import (
     branchmap,
-    bundle2,
     changegroup,
     cmdutil,
     commands,
     context,
-    discovery as discoverymod,
     error,
-    exchange,
     extensions,
     localrepo,
     lock,
@@ -36,7 +33,6 @@
     patch,
     phases,
     util,
-    wireproto,
 )
 
 from . import (
@@ -93,6 +89,7 @@
 def uisetup(ui):
     destination.modsetup(ui)
     topicrevset.modsetup(ui)
+    discovery.modsetup(ui)
 
 @contextlib.contextmanager
 def usetopicmap(repo):
@@ -356,16 +353,7 @@
 
 extensions.wrapfunction(cmdutil, 'buildcommittext', committextwrap)
 extensions.wrapfunction(merge, 'update', mergeupdatewrap)
-extensions.wrapfunction(discoverymod, '_headssummary', discovery._headssummary)
-extensions.wrapfunction(wireproto, 'branchmap', discovery.wireprotobranchmap)
-extensions.wrapfunction(wireproto, '_capabilities', discovery.wireprotocaps)
-extensions.wrapfunction(bundle2, 'handlecheckheads', discovery.handlecheckheads)
-# we need a proper wrape b2 part stuff
-bundle2.handlecheckheads.params = frozenset()
-bundle2.parthandlermapping['check:heads'] = bundle2.handlecheckheads
-extensions.wrapfunction(exchange, '_pushb2phases', discovery._pushb2phases)
 extensions.wrapfunction(changegroup.cg1unpacker, 'apply', cgapply)
-exchange.b2partsgenmapping['phase'] = exchange._pushb2phases
 cmdutil.summaryhooks.add('topic', summaryhook)
 
 if util.safehasattr(cmdutil, 'extraexport'):

--- a/hgext3rd/topic/discovery.py	Wed Mar 30 03:22:50 2016 -0700
+++ b/hgext3rd/topic/discovery.py	Wed Mar 30 21:53:55 2016 -0700
@@ -5,8 +5,12 @@
 from mercurial.i18n import _
 from mercurial import (
     branchmap,
+    bundle2,
+    discovery,
     error,
     exchange,
+    extensions,
+    wireproto,
 )
 
 from . import topicmap
@@ -114,3 +118,15 @@
     if repo.peer().capable('topics'):
         caps.append('topics')
     return caps
+
+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(bundle2, 'handlecheckheads', handlecheckheads)
+    # we need a proper wrap b2 part stuff
+    bundle2.handlecheckheads.params = frozenset()
+    bundle2.parthandlermapping['check:heads'] = bundle2.handlecheckheads
+    extensions.wrapfunction(exchange, '_pushb2phases', _pushb2phases)
+    exchange.b2partsgenmapping['phase'] = exchange._pushb2phases