# sparse.py - allow sparse checkouts of the working directory

#

# Copyright 2014 Facebook, Inc.

#

# This software may be used and distributed according to the terms of the

# GNU General Public License version 2 or any later version.

"""allow sparse checkouts of the working directory (EXPERIMENTAL)

(This extension is not yet protected by backwards compatibility

guarantees. Any aspect may break in future releases until this

notice is removed.)

This extension allows the working directory to only consist of a

subset of files for the revision. This allows specific files or

directories to be explicitly included or excluded. Many repository

operations have performance proportional to the number of files in

the working directory. So only realizing a subset of files in the

working directory can improve performance.

"""

from __future__ import absolute_import

import collections

import hashlib

import os

from mercurial.i18n import _

from mercurial.node import nullid

from mercurial import (

    cmdutil,

    commands,

    context,

    dirstate,

    error,

    extensions,

    hg,

    localrepo,

    match as matchmod,

    merge as mergemod,

    registrar,

    util,

)

# Note for extension authors: ONLY specify testedwith = 'ships-with-hg-core' for

# extensions which SHIP WITH MERCURIAL. Non-mainline extensions should

# be specifying the version(s) of Mercurial they are tested with, or

# leave the attribute unspecified.

testedwith = 'ships-with-hg-core'

cmdtable = {}

command = registrar.command(cmdtable)

def uisetup(ui):

    _setupupdates(ui)

    _setupcommit(ui)

def extsetup(ui):

    _setupclone(ui)

    _setuplog(ui)

    _setupadd(ui)

    _setupdirstate(ui)

    # if fsmonitor is enabled, tell it to use our hash function

    try:

        fsmonitor = extensions.find('fsmonitor')

        def _hashignore(orig, ignore):

            return _hashmatcher(ignore)

        extensions.wrapfunction(fsmonitor, '_hashignore', _hashignore)

    except KeyError:

        pass

def reposetup(ui, repo):

    if not util.safehasattr(repo, 'dirstate'):

        return

    _wraprepo(ui, repo)

def replacefilecache(cls, propname, replacement):

    """Replace a filecache property with a new class. This allows changing the

    cache invalidation condition."""

    origcls = cls

    assert callable(replacement)

    while cls is not object:

        if propname in cls.__dict__:

            orig = cls.__dict__[propname]

            setattr(cls, propname, replacement(orig))

            break

        cls = cls.__bases__[0]

    if cls is object:

        raise AttributeError(_("type '%s' has no property '%s'") % (origcls,

                             propname))

def _setupupdates(ui):

    def _calculateupdates(orig, repo, wctx, mctx, ancestors, branchmerge, *arg,

                          **kwargs):

        """Filter updates to only lay out files that match the sparse rules.

        """

        actions, diverge, renamedelete = orig(repo, wctx, mctx, ancestors,

                                              branchmerge, *arg, **kwargs)

        if not util.safehasattr(repo, 'sparsematch'):

            return actions, diverge, renamedelete

        files = set()

        prunedactions = {}

        oldrevs = [pctx.rev() for pctx in wctx.parents()]

        oldsparsematch = repo.sparsematch(*oldrevs)

        if branchmerge:

            # If we're merging, use the wctx filter, since we're merging into

            # the wctx.

            sparsematch = repo.sparsematch(wctx.parents()[0].rev())

        else:

            # If we're updating, use the target context's filter, since we're

            # moving to the target context.

            sparsematch = repo.sparsematch(mctx.rev())

        temporaryfiles = []

        for file, action in actions.iteritems():

            type, args, msg = action

            files.add(file)

            if sparsematch(file):

                prunedactions[file] = action

            elif type == 'm':

                temporaryfiles.append(file)

                prunedactions[file] = action

            elif branchmerge:

                if type != 'k':

                    temporaryfiles.append(file)

                    prunedactions[file] = action

            elif type == 'f':

                prunedactions[file] = action

            elif file in wctx:

                prunedactions[file] = ('r', args, msg)

        if len(temporaryfiles) > 0:

            ui.status(_("temporarily included %d file(s) in the sparse checkout"

                " for merging\n") % len(temporaryfiles))

            repo.addtemporaryincludes(temporaryfiles)

            # Add the new files to the working copy so they can be merged, etc

            actions = []

            message = 'temporarily adding to sparse checkout'

            wctxmanifest = repo[None].manifest()

            for file in temporaryfiles:

                if file in wctxmanifest:

                    fctx = repo[None][file]

                    actions.append((file, (fctx.flags(), False), message))

            typeactions = collections.defaultdict(list)

            typeactions['g'] = actions

            mergemod.applyupdates(repo, typeactions, repo[None], repo['.'],

                                  False)

            dirstate = repo.dirstate

            for file, flags, msg in actions:

                dirstate.normal(file)

        profiles = repo.getactiveprofiles()

        changedprofiles = profiles & files

        # If an active profile changed during the update, refresh the checkout.

        # Don't do this during a branch merge, since all incoming changes should

        # have been handled by the temporary includes above.

        if changedprofiles and not branchmerge:

            mf = mctx.manifest()

            for file in mf:

                old = oldsparsematch(file)

                new = sparsematch(file)

                if not old and new:

                    flags = mf.flags(file)

                    prunedactions[file] = ('g', (flags, False), '')

                elif old and not new:

                    prunedactions[file] = ('r', [], '')

        return prunedactions, diverge, renamedelete

    extensions.wrapfunction(mergemod, 'calculateupdates', _calculateupdates)

    def _update(orig, repo, node, branchmerge, *args, **kwargs):

        results = orig(repo, node, branchmerge, *args, **kwargs)

        # If we're updating to a location, clean up any stale temporary includes

        # (ex: this happens during hg rebase --abort).

        if not branchmerge and util.safehasattr(repo, 'sparsematch'):

            repo.prunetemporaryincludes()

        return results

    extensions.wrapfunction(mergemod, 'update', _update)

def _setupcommit(ui):

    def _refreshoncommit(orig, self, node):

        """Refresh the checkout when commits touch .hgsparse

        """

        orig(self, node)

        repo = self._repo

        if util.safehasattr(repo, 'getsparsepatterns'):

            ctx = repo[node]

            _, _, profiles = repo.getsparsepatterns(ctx.rev())

            if set(profiles) & set(ctx.files()):

                origstatus = repo.status()

                origsparsematch = repo.sparsematch()

                _refresh(repo.ui, repo, origstatus, origsparsematch, True)

            repo.prunetemporaryincludes()

    extensions.wrapfunction(context.committablectx, 'markcommitted',

        _refreshoncommit)

def _setuplog(ui):

    entry = commands.table['^log|history']

    entry[1].append(('', 'sparse', None,

        "limit to changesets affecting the sparse checkout"))

    def _logrevs(orig, repo, opts):

        revs = orig(repo, opts)

        if opts.get('sparse'):

            sparsematch = repo.sparsematch()

            def ctxmatch(rev):

                ctx = repo[rev]

                return any(f for f in ctx.files() if sparsematch(f))

            revs = revs.filter(ctxmatch)

        return revs

    extensions.wrapfunction(cmdutil, '_logrevs', _logrevs)

def _clonesparsecmd(orig, ui, repo, *args, **opts):

    include_pat = opts.get('include')

    exclude_pat = opts.get('exclude')

    enableprofile_pat = opts.get('enable_profile')

    include = exclude = enableprofile = False

    if include_pat:

        pat = include_pat

        include = True

    if exclude_pat:

        pat = exclude_pat

        exclude = True

    if enableprofile_pat:

        pat = enableprofile_pat

        enableprofile = True

    if sum([include, exclude, enableprofile]) > 1:

        raise error.Abort(_("too many flags specified."))

    if include or exclude or enableprofile:

        def clonesparse(orig, self, node, overwrite, *args, **kwargs):

            _config(self.ui, self.unfiltered(), pat, {}, include=include,

                    exclude=exclude, enableprofile=enableprofile)

            return orig(self, node, overwrite, *args, **kwargs)

        extensions.wrapfunction(hg, 'updaterepo', clonesparse)

    return orig(ui, repo, *args, **opts)

def _setupclone(ui):

    entry = commands.table['^clone']

    entry[1].append(('', 'enable-profile', [],

                    'enable a sparse profile'))

    entry[1].append(('', 'include', [],

                    'include sparse pattern'))

    entry[1].append(('', 'exclude', [],

                    'exclude sparse pattern'))

    extensions.wrapcommand(commands.table, 'clone', _clonesparsecmd)

def _setupadd(ui):

    entry = commands.table['^add']

    entry[1].append(('s', 'sparse', None,

                    'also include directories of added files in sparse config'))

    def _add(orig, ui, repo, *pats, **opts):

        if opts.get('sparse'):

            dirs = set()

            for pat in pats:

                dirname, basename = util.split(pat)

                dirs.add(dirname)

            _config(ui, repo, list(dirs), opts, include=True)

        return orig(ui, repo, *pats, **opts)

    extensions.wrapcommand(commands.table, 'add', _add)

def _setupdirstate(ui):

    """Modify the dirstate to prevent stat'ing excluded files,

    and to prevent modifications to files outside the checkout.

    """

    def _dirstate(orig, repo):

        dirstate = orig(repo)

        dirstate.repo = repo

        return dirstate

    extensions.wrapfunction(

        localrepo.localrepository.dirstate, 'func', _dirstate)

    # The atrocity below is needed to wrap dirstate._ignore. It is a cached

    # property, which means normal function wrapping doesn't work.

    class ignorewrapper(object):

        def __init__(self, orig):

            self.orig = orig

            self.origignore = None

            self.func = None

            self.sparsematch = None

        def __get__(self, obj, type=None):

            repo = obj.repo

            origignore = self.orig.__get__(obj)

            if not util.safehasattr(repo, 'sparsematch'):

                return origignore

            sparsematch = repo.sparsematch()

            if self.sparsematch != sparsematch or self.origignore != origignore:

                self.func = unionmatcher([origignore,

                                          negatematcher(sparsematch)])

                self.sparsematch = sparsematch

                self.origignore = origignore

            return self.func

        def __set__(self, obj, value):

            return self.orig.__set__(obj, value)

        def __delete__(self, obj):

            return self.orig.__delete__(obj)

    replacefilecache(dirstate.dirstate, '_ignore', ignorewrapper)

    # dirstate.rebuild should not add non-matching files

    def _rebuild(orig, self, parent, allfiles, changedfiles=None):

        if util.safehasattr(self.repo, 'sparsematch'):

            matcher = self.repo.sparsematch()

            allfiles = allfiles.matches(matcher)

            if changedfiles:

                changedfiles = [f for f in changedfiles if matcher(f)]

            if changedfiles is not None:

                # In _rebuild, these files will be deleted from the dirstate

                # when they are not found to be in allfiles

                dirstatefilestoremove = set(f for f in self if not matcher(f))

                changedfiles = dirstatefilestoremove.union(changedfiles)

        return orig(self, parent, allfiles, changedfiles)

    extensions.wrapfunction(dirstate.dirstate, 'rebuild', _rebuild)

    # Prevent adding files that are outside the sparse checkout

    editfuncs = ['normal', 'add', 'normallookup', 'copy', 'remove', 'merge']

    hint = _('include file with `hg debugsparse --include <pattern>` or use ' +

             '`hg add -s <file>` to include file directory while adding')

    for func in editfuncs:

        def _wrapper(orig, self, *args):

            repo = self.repo

            if util.safehasattr(repo, 'sparsematch'):

                dirstate = repo.dirstate

                sparsematch = repo.sparsematch()

                for f in args:

                    if (f is not None and not sparsematch(f) and

                        f not in dirstate):

                        raise error.Abort(_("cannot add '%s' - it is outside "