~abentley/bzrtools/bzrtools.dev

#!/usr/bin/python

from userinteractor import UserInteractor, UserOption

class HunkSelector:

    def __init__(self, patches):

        self.standard_options = [

            UserOption('y', self.shelve, 'shelve this change.', default=True),

            UserOption('n', self.keep, 'keep this change in your tree.'),

            UserOption('d', UserInteractor.FINISH, 'done, skip to the end.'),

            UserOption('i', self.invert,

                'invert the current selection status of all hunks.'),

            UserOption('s', self.status, 'show selection status of all hunks.'),

            UserOption('q', UserInteractor.QUIT, 'quit')

        ]

        self.end_options = [

            UserOption('y', UserInteractor.FINISH,

                'shelve selected changes.', default=True),

            UserOption('r', UserInteractor.RESTART,

                'restart the hunk selection loop.'),

            UserOption('s', self.status, 'show selection status of all hunks.'),

            UserOption('i', self.invert,

                'invert the current selection status of all hunks.'),

            UserOption('q', UserInteractor.QUIT, 'quit')

        ]

        self.patches = patches

        self.total_hunks = 0

        self.last_printed = -1

        self.interactor = UserInteractor()

        self.interactor.set_item_callback(self.hunk_callback)

        self.interactor.set_start_callback(self.start_callback)

        self.interactor.set_end_callback(self.end_callback)

        for patch in patches:

            for hunk in patch.hunks:

                # everything's shelved by default

                hunk.selected = True

                self.total_hunks += 1

                # we need a back pointer in the callbacks

                hunk.patch = patch

                self.interactor.add_item(hunk)

    # Called at the start of the main loop

    def start_callback(self):

        self.interactor.set_prompt('Shelve this change? ' \

            '(%(count)d of %(total)d)')

        self.interactor.set_options(self.standard_options)

    # Called at the end of the item loop, return False to indicate that the

    # interaction isn't finished and the confirmation prompt should be displayed

    def end_callback(self):

        self.status()

        self.interactor.set_prompt("Shelve these changes?")

        self.interactor.set_options(self.end_options)

        return False

    # Called once for each hunk

    def hunk_callback(self, hunk, count):

        if self.last_printed != count:

            print hunk.patch.get_header(), hunk

            self.last_printed = count

        if hunk.selected:

            self.interactor.get_option('y').default = True

            self.interactor.get_option('n').default = False

        else:

            self.interactor.get_option('y').default = False

            self.interactor.get_option('n').default = True

    # The user chooses to shelve a hunk

    def shelve(self, hunk):

        hunk.selected = True

        return True

    # The user chooses to keep a hunk

    def keep(self, hunk):

        hunk.selected = False

        return True

    # The user chooses to invert the selection

    def invert(self, hunk):

        for patch in self.patches:

            for hunk in patch.hunks:

                if hunk.__dict__.has_key('selected'):

                    hunk.selected = not hunk.selected

                else:

                    hunk.selected = True

        self.status()

        return False

    # The user wants to see the status

    def status(self, hunk=None):

        print '\nStatus:'

        for patch in self.patches:

            print '  %s' % patch.oldname

            shelve = 0

            keep = 0

            for hunk in patch.hunks:

                if hunk.selected:

                    shelve += 1

                else:

                    keep += 1

            print '    %d hunks to be shelved' % shelve

            print '    %d hunks to be kept' % keep

            print

        # Tell the interactor we're not done with this item

        return False

    def select(self):

        if self.total_hunks == 0:

            return []

        if not self.interactor.interact():

            # False from interact means they quit

            return []

        for patch in self.patches:

            tmp = []

            for hunk in patch.hunks:

                if hunk.selected:

                    tmp.append(hunk)

            patch.hunks = tmp

        tmp = []

        for patch in self.patches:

            if len(patch.hunks):

                tmp.append(patch)

        self.patches = tmp

        return self.patches