hgext/narrow/narrowwirepeer.py

# narrowwirepeer.py - passes narrow spec with unbundle command
#
# Copyright 2017 Google, Inc.
#
# 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

from mercurial import (
    extensions,
    hg,
    wireprotov1server,
)

NARROWCAP = 'exp-narrow-1'
ELLIPSESCAP = 'exp-ellipses-1'

def uisetup():
    extensions.wrapfunction(wireprotov1server, '_capabilities', addnarrowcap)

def addnarrowcap(orig, repo, proto):
    """add the narrow capability to the server"""
    caps = orig(repo, proto)
    caps.append(NARROWCAP)
    if repo.ui.configbool('experimental', 'narrowservebrokenellipses'):
        caps.append(ELLIPSESCAP)
    return caps

def reposetup(repo):
    def wirereposetup(ui, peer):
        def wrapped(orig, cmd, *args, **kwargs):
            if cmd == 'unbundle':
                # TODO: don't blindly add include/exclude wireproto
                # arguments to unbundle.
                include, exclude = repo.narrowpats
                kwargs[r"includepats"] = ','.join(include)
                kwargs[r"excludepats"] = ','.join(exclude)
            return orig(cmd, *args, **kwargs)
        extensions.wrapfunction(peer, '_calltwowaystream', wrapped)
    hg.wirepeersetupfuncs.append(wirereposetup)