~abentley/bzrtools/bzrtools.dev

#!/usr/bin/python

import os

import sys

from errors import CommandError

from hunk_selector import ShelveHunkSelector, UnshelveHunkSelector

from diffstat import DiffStat

from patchsource import PatchSource, FilePatchSource

class Shelf(object):

    MESSAGE_PREFIX = "# Shelved patch: "

    _paths = {

        'base'          : '.shelf',

        'shelves'       : '.shelf/shelves',

        'current-shelf' : '.shelf/current-shelf',

    }

    def __init__(self, base, name=None):

        self.base = base

        self.__setup()

        if name is None:

            current = os.path.join(self.base, self._paths['current-shelf'])

            name = open(current).read().strip()

        assert '\n' not in name

        self.name = name

        self.dir = os.path.join(self.base, self._paths['shelves'], name)

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

            os.mkdir(self.dir)

    def __setup(self):

        # Create required directories etc.

        for dir in [self._paths['base'], self._paths['shelves']]:

            dir = os.path.join(self.base, dir)

            if not os.path.isdir(dir):

                os.mkdir(dir)

        current = os.path.join(self.base, self._paths['current-shelf'])

        if not os.path.exists(current):

            f = open(current, 'w')

            f.write('default')

            f.close()

    def make_default(self):

        f = open(os.path.join(self.base, self._paths['current-shelf']), 'w')

        f.write(self.name)

        f.close()

        self.log("Default shelf is now '%s'\n" % self.name)

    def log(self, msg):

        sys.stderr.write(msg)

    def delete(self, patch):

        path = self.__path_from_user(patch)

        os.remove(path)

    def display(self, patch):

        path = self.__path_from_user(patch)

        sys.stdout.write(open(path).read())

    def list(self):

        self.log("Patches on shelf '%s':" % self.name)

        indexes = self.__list()

        if len(indexes) == 0:

            self.log(' None\n')

            return

        self.log('\n')

        for index in indexes:

            msg = self.get_patch_message(self.__path(index))

            if msg is None:

                msg = "No message saved with patch."

            self.log(' %.2d: %s\n' % (index, msg))

    def __path_from_user(self, patch_id):

        try:

            patch_index = int(patch_id)

        except TypeError:

            raise CommandError("Invalid patch name '%s'" % patch_id)

        path = self.__path(patch_index)

        if not os.path.exists(path):

            raise CommandError("Patch '%s' doesn't exist on shelf %s!" % \

                        (patch_id, self.name))

        return path

    def __path(self, index):

        return os.path.join(self.dir, '%.2d' % index)

    def next_patch(self):

        indexes = self.__list()

        if len(indexes) == 0:

            next = 0

        else:

            next = indexes[-1] + 1

        return self.__path(next)

    def __list(self):

        patches = os.listdir(self.dir)

        indexes = [int(f) for f in patches]

        indexes.sort()

        return indexes

    def last_patch(self):

        indexes = self.__list()

        if len(indexes) == 0:

            return None

        return self.__path(indexes[-1])

    def get_patch_message(self, patch_path):

        patch = open(patch_path, 'r').read()

        if not patch.startswith(self.MESSAGE_PREFIX):

            return None

        return patch[len(self.MESSAGE_PREFIX):patch.index('\n')]

    def __show_status(self, source):

        if source.can_live_update():

            diff_stat = str(DiffStat(source.readlines()))

            if len(diff_stat) > 0:

                self.log('Diff status is now:\n' + diff_stat + '\n')

            else:

                self.log('No changes left in working tree.\n')

    def unshelve(self, patch_source, pick_hunks=False):

        patch_name = self.last_patch()

        if patch_name is None:

            raise CommandError("No patch found on shelf %s" % self.name)

        hunks = FilePatchSource(patch_name).readhunks()

        if pick_hunks:

            to_unshelve, to_remain = UnshelveHunkSelector(hunks).select()

        else:

            to_unshelve = hunks

            to_remain = []

        if len(to_unshelve) == 0:

            raise CommandError('Nothing to unshelve')

        message = self.get_patch_message(patch_name)

        if message is None:

            message = ""

        self.log('Reapplying shelved patches "%s"\n' % message)

        pipe = os.popen('patch -d %s -s -p0' % self.base, 'w')

        for hunk in to_unshelve:

            pipe.write(str(hunk))

        pipe.flush()

        if pipe.close() is not None:

            raise CommandError("Failed running patch!")

        if len(to_remain) == 0:

            os.remove(patch_name)

        else:

            f = open(patch_name, 'w')

            for hunk in to_remain:

                f.write(str(hunk))

            f.close()

        self.__show_status(patch_source)

    def shelve(self, patch_source, pick_hunks=False, message=None):

        from datetime import datetime

        hunks = patch_source.readhunks()

        if pick_hunks:

            to_shelve = ShelveHunkSelector(hunks).select()[0]

        else:

            to_shelve = hunks

        if len(to_shelve) == 0:

            raise CommandError('Nothing to shelve')

        if message is None:

            timestamp = datetime.now().strftime('%Y-%m-%d %H:%M:%S')

            message = "Changes shelved on %s" % timestamp

        patch_name = self.next_patch()

        self.log('Shelving to %s/%s: "%s"\n' % \

                (self.name, os.path.basename(patch_name), message))

        patch = open(patch_name, 'a')

        assert '\n' not in message

        patch.write("%s%s\n" % (self.MESSAGE_PREFIX, message))

        for hunk in to_shelve:

            patch.write(str(hunk))

        patch.flush()

        os.fsync(patch.fileno())

        patch.close()

        pipe = os.popen('patch -d %s -sR -p0' % self.base, 'w')

        for hunk in to_shelve:

            pipe.write(str(hunk))

        pipe.flush()

        if pipe.close() is not None:

            raise CommandError("Failed running patch!")

        self.__show_status(patch_source)