# Patch transplanting extension for Mercurial

#

# Copyright 2006 Brendan Cully <brendan@kublai.com>

#

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

# of the GNU General Public License, incorporated herein by reference.

from mercurial.demandload import *

from mercurial.i18n import gettext as _

demandload(globals(), 'os tempfile')

demandload(globals(), 'mercurial:bundlerepo,cmdutil,commands,hg,merge,patch')

demandload(globals(), 'mercurial:revlog,util')

'''patch transplanting tool

This extension allows you to transplant patches from another branch.

Transplanted patches are recorded in .hg/transplant/transplants, as a map

from a changeset hash to its hash in the source repository.

'''

class transplantentry:

    def __init__(self, lnode, rnode):

        self.lnode = lnode

        self.rnode = rnode

class transplants:

    def __init__(self, path=None, transplantfile=None, opener=None):

        self.path = path

        self.transplantfile = transplantfile

        self.opener = opener

        if not opener:

            self.opener = util.opener(self.path)

        self.transplants = []

        self.dirty = False

        self.read()

    def read(self):

        abspath = os.path.join(self.path, self.transplantfile)

        if self.transplantfile and os.path.exists(abspath):

            for line in self.opener(self.transplantfile).read().splitlines():

                lnode, rnode = map(revlog.bin, line.split(':'))

                self.transplants.append(transplantentry(lnode, rnode))

    def write(self):

        if self.dirty and self.transplantfile:

            if not os.path.isdir(self.path):

                os.mkdir(self.path)

            fp = self.opener(self.transplantfile, 'w')

            for c in self.transplants:

                l, r = map(revlog.hex, (c.lnode, c.rnode))

                fp.write(l + ':' + r + '\n')

            fp.close()

        self.dirty = False

    def get(self, rnode):

        return [t for t in self.transplants if t.rnode == rnode]

    def set(self, lnode, rnode):

        self.transplants.append(transplantentry(lnode, rnode))

        self.dirty = True

    def remove(self, transplant):

        del self.transplants[self.transplants.index(transplant)]

        self.dirty = True

class transplanter:

    def __init__(self, ui, repo):

        self.ui = ui

        self.path = repo.join('transplant')

        self.opener = util.opener(self.path)

        self.transplants = transplants(self.path, 'transplants', opener=self.opener)

    def applied(self, repo, node, parent):

        '''returns True if a node is already an ancestor of parent

        or has already been transplanted'''

        if hasnode(repo, node):

            if node in repo.changelog.reachable(parent, stop=node):

                return True

        for t in self.transplants.get(node):

            # it might have been stripped

            if not hasnode(repo, t.lnode):

                self.transplants.remove(t)

                return False

            if t.lnode in repo.changelog.reachable(parent, stop=t.lnode):

                return True

        return False

    def apply(self, repo, source, revmap, merges, opts={}):

        '''apply the revisions in revmap one by one in revision order'''

        revs = revmap.keys()

        revs.sort()

        p1, p2 = repo.dirstate.parents()

        pulls = []

        diffopts = patch.diffopts(self.ui, opts)

        diffopts.git = True

        lock = repo.lock()

        wlock = repo.wlock()

        try:

            for rev in revs:

                node = revmap[rev]

                revstr = '%s:%s' % (rev, revlog.short(node))

                if self.applied(repo, node, p1):

                    self.ui.warn(_('skipping already applied revision %s\n') %

                                 revstr)

                    continue

                parents = source.changelog.parents(node)

                if not opts.get('filter'):

                    # If the changeset parent is the same as the wdir's parent,

                    # just pull it.

                    if parents[0] == p1:

                        pulls.append(node)

                        p1 = node

                        continue

                    if pulls:

                        if source != repo:

                            repo.pull(source, heads=pulls, lock=lock)

                        merge.update(repo, pulls[-1], wlock=wlock)

                        p1, p2 = repo.dirstate.parents()

                        pulls = []

                domerge = False

                if node in merges:

                    # pulling all the merge revs at once would mean we couldn't

                    # transplant after the latest even if transplants before them

                    # fail.

                    domerge = True

                    if not hasnode(repo, node):

                        repo.pull(source, heads=[node], lock=lock)

                if parents[1] != revlog.nullid:

                    self.ui.note(_('skipping merge changeset %s:%s\n')

                                 % (rev, revlog.short(node)))

                    patchfile = None

                else:

                    fd, patchfile = tempfile.mkstemp(prefix='hg-transplant-')

                    fp = os.fdopen(fd, 'w')

                    patch.export(source, [node], fp=fp, opts=diffopts)

                    fp.close()

                del revmap[rev]

                if patchfile or domerge:

                    try:

                        n = self.applyone(repo, node, source.changelog.read(node),

                                          patchfile, merge=domerge,

                                          log=opts.get('log'),

                                          filter=opts.get('filter'),

                                          lock=lock, wlock=wlock)

                        if domerge:

                            self.ui.status(_('%s merged at %s\n') % (revstr,

                                      revlog.short(n)))

                        else:

                            self.ui.status(_('%s transplanted to %s\n') % (revlog.short(node),

                                                                       revlog.short(n)))

                    finally:

                        if patchfile:

                            os.unlink(patchfile)

            if pulls:

                repo.pull(source, heads=pulls, lock=lock)

                merge.update(repo, pulls[-1], wlock=wlock)

        finally:

            self.saveseries(revmap, merges)

            self.transplants.write()

    def filter(self, filter, changelog, patchfile):

        '''arbitrarily rewrite changeset before applying it'''

        self.ui.status('filtering %s\n' % patchfile)

        util.system('%s %s' % (filter, util.shellquote(patchfile)),

                    environ={'HGUSER': changelog[1]},

                    onerr=util.Abort, errprefix=_('filter failed'))

    def applyone(self, repo, node, cl, patchfile, merge=False, log=False,

                 filter=None, lock=None, wlock=None):

        '''apply the patch in patchfile to the repository as a transplant'''

        (manifest, user, (time, timezone), files, message) = cl[:5]

        date = "%d %d" % (time, timezone)

        extra = {'transplant_source': node}

        if filter:

            self.filter(filter, cl, patchfile)

            patchfile, message, user, date = patch.extract(self.ui, file(patchfile))

        if log:

            message += '\n(transplanted from %s)' % revlog.hex(node)

        self.ui.status(_('applying %s\n') % revlog.short(node))

        self.ui.note('%s %s\n%s\n' % (user, date, message))

        if not patchfile and not merge:

            raise util.Abort(_('can only omit patchfile if merging'))

        if patchfile:

            try:

                files = {}

                try:

                    fuzz = patch.patch(patchfile, self.ui, cwd=repo.root,

                                       files=files)

                    if not files:

                        self.ui.warn(_('%s: empty changeset') % revlog.hex(node))

                        return

                finally:

                    files = patch.updatedir(self.ui, repo, files, wlock=wlock)

                if filter:

                    os.unlink(patchfile)

            except Exception, inst:

                if filter:

                    os.unlink(patchfile)

                p1 = repo.dirstate.parents()[0]

                p2 = node

                self.log(user, date, message, p1, p2, merge=merge)

                self.ui.write(str(inst) + '\n')

                raise util.Abort(_('Fix up the merge and run hg transplant --continue'))

        else:

            files = None

        if merge:

            p1, p2 = repo.dirstate.parents()

            repo.dirstate.setparents(p1, node)

        n = repo.commit(files, message, user, date, lock=lock, wlock=wlock,

                        extra=extra)

        if not merge:

            self.transplants.set(n, node)

        return n

    def resume(self, repo, source, opts=None):

        '''recover last transaction and apply remaining changesets'''

        if os.path.exists(os.path.join(self.path, 'journal')):

            n, node = self.recover(repo)

            self.ui.status(_('%s transplanted as %s\n') % (revlog.short(node),

                                                           revlog.short(n)))

        seriespath = os.path.join(self.path, 'series')

        if not os.path.exists(seriespath):

            return

        nodes, merges = self.readseries()

        revmap = {}

        for n in nodes:

            revmap[source.changelog.rev(n)] = n

        os.unlink(seriespath)

        self.apply(repo, source, revmap, merges, opts)

    def recover(self, repo):

        '''commit working directory using journal metadata'''

        node, user, date, message, parents = self.readlog()

        merge = len(parents) == 2

        if not user or not date or not message or not parents[0]:

            raise util.Abort(_('transplant log file is corrupt'))

        wlock = repo.wlock()

        p1, p2 = repo.dirstate.parents()

        if p1 != parents[0]:

            raise util.Abort(_('working dir not at transplant parent %s') %

                             revlog.hex(parents[0]))

        if merge:

            repo.dirstate.setparents(p1, parents[1])

        n = repo.commit(None, message, user, date, wlock=wlock)

        if not n:

            raise util.Abort(_('commit failed'))

        if not merge:

            self.transplants.set(n, node)

        self.unlog()

        return n, node

    def readseries(self):

        nodes = []

        merges = []

        cur = nodes

        for line in self.opener('series').read().splitlines():

            if line.startswith('# Merges'):

                cur = merges

                continue

            cur.append(revlog.bin(line))

        return (nodes, merges)

    def saveseries(self, revmap, merges):

        if not revmap:

            return

        if not os.path.isdir(self.path):

            os.mkdir(self.path)

        series = self.opener('series', 'w')

        revs = revmap.keys()

        revs.sort()

        for rev in revs:

            series.write(revlog.hex(revmap[rev]) + '\n')

        if merges:

            series.write('# Merges\n')

            for m in merges:

                series.write(revlog.hex(m) + '\n')

        series.close()

    def log(self, user, date, message, p1, p2, merge=False):

        '''journal changelog metadata for later recover'''

        if not os.path.isdir(self.path):

            os.mkdir(self.path)

        fp = self.opener('journal', 'w')

        fp.write('# User %s\n' % user)

        fp.write('# Date %s\n' % date)

        fp.write('# Node ID %s\n' % revlog.hex(p2))

        fp.write('# Parent ' + revlog.hex(p1) + '\n')

        if merge:

            fp.write('# Parent ' + revlog.hex(p2) + '\n')

        fp.write(message.rstrip() + '\n')

        fp.close()

    def readlog(self):

        parents = []

        message = []

        for line in self.opener('journal').read().splitlines():

            if line.startswith('# User '):

                user = line[7:]

            elif line.startswith('# Date '):

                date = line[7:]

            elif line.startswith('# Node ID '):

                node = revlog.bin(line[10:])

            elif line.startswith('# Parent '):

                parents.append(revlog.bin(line[9:]))

            else:

                message.append(line)

        return (node, user, date, '\n'.join(message), parents)

    def unlog(self):

        '''remove changelog journal'''

        absdst = os.path.join(self.path, 'journal')

        if os.path.exists(absdst):

            os.unlink(absdst)

    def transplantfilter(self, repo, source, root):

        def matchfn(node):

            if self.applied(repo, node, root):

                return False

            if source.changelog.parents(node)[1] != revlog.nullid:

                return False

            extra = source.changelog.read(node)[5]

            cnode = extra.get('transplant_source')

            if cnode and self.applied(repo, cnode, root):

                return False

            return True

        return matchfn

def hasnode(repo, node):

    try:

        return repo.changelog.rev(node) != None

    except revlog.RevlogError:

        return False

def browserevs(ui, repo, nodes, opts):

    '''interactively transplant changesets'''

    def browsehelp(ui):

        ui.write('y: transplant this changeset\n'

                 'n: skip this changeset\n'