~abentley/bzrtools/bzrtools.dev

#!/usr/bin/python

import os

import sys

import subprocess

from datetime import datetime

from errors import CommandError, PatchFailed

from hunk_selector import ShelveHunkSelector, UnshelveHunkSelector

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):

        indexes = self.__list()

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

        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 = []

        for f in patches:

            if f.endswith('~'):

                continue # ignore backup files

            try:

                indexes.append(int(f))

            except ValueError:

                self.log("Warning: Ignoring junk file '%s' on shelf.\n" % f)

        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 unshelve(self, patch_source, all_hunks=False, force=False):

        self._check_upgrade()

        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 all_hunks:

            to_unshelve = hunks

            to_remain = []

        else:

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

        if len(to_unshelve) == 0:

            raise CommandError('Nothing to unshelve')

        message = self.get_patch_message(patch_name)

        if message is None:

            message = "No message saved with patch."

        self.log('Unshelving from %s/%s: "%s"\n' % \

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

        try:

            self._run_patch(to_unshelve, dry_run=True)

            self._run_patch(to_unshelve)

        except PatchFailed:

            try:

                self._run_patch(to_unshelve, strip=0, dry_run=True)

                self._run_patch(to_unshelve, strip=0)

            except PatchFailed:

                if force:

                    self.log('Warning: Unshelving failed, forcing as ' \

                             'requested. Shelf will not be modified.\n')

                    try:

                        self._run_patch(to_unshelve)

                    except PatchFailed:

                        pass

                    return

                raise CommandError("Your shelved patch no " \

                    "longer applies cleanly to the working tree!")

        # Backup the shelved patch

        os.rename(patch_name, '%s~' % patch_name)

        if len(to_remain) > 0:

            f = open(patch_name, 'w')

            for hunk in to_remain:

                f.write(str(hunk))

            f.close()

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

        self._check_upgrade()

        hunks = patch_source.readhunks()

        if all_hunks:

            to_shelve = hunks

        else:

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

        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()

        try:

            self._run_patch(to_shelve, reverse=True, dry_run=True)

            self._run_patch(to_shelve, reverse=True)

        except PatchFailed:

            try:

                self._run_patch(to_shelve, reverse=True, strip=0, dry_run=True)

                self._run_patch(to_shelve, reverse=True, strip=0)

            except PatchFailed:

                raise CommandError("Failed removing shelved changes from the"

                    "working tree!")

    def _run_patch(self, patches, strip=1, reverse=False, dry_run=False):

        args = ['patch', '-d', self.base, '-s', '-p%d' % strip, '-f']

        if reverse:

            args.append('-R')

        if dry_run:

            args.append('--dry-run')

            stdout = stderr = subprocess.PIPE

        else:

            stdout = stderr = None

        process = subprocess.Popen(args, stdin=subprocess.PIPE, stdout=stdout,

                        stderr=stderr)

        for patch in patches:

            process.stdin.write(str(patch))

        process.communicate()

        result = process.wait()

        if result != 0:

            raise PatchFailed()

        return result

    def _check_upgrade(self):

        if len(self._list_old_shelves()) > 0:

            raise CommandError("Old format shelves found, either upgrade " \

                    "or remove them!")

    def _list_old_shelves(self):

        import glob

        stem = os.path.join(self.base, '.bzr-shelf')

        patches = glob.glob(stem)

        patches.extend(glob.glob(stem + '-*[!~]'))

        if len(patches) == 0:

            return []

        def patch_index(name):

            if name == stem:

                return 0

            return int(name[len(stem) + 1:])

        # patches might not be sorted in the right order

        patch_ids = []

        for patch in patches:

            if patch == stem:

                patch_ids.append(0)

            else:

                patch_ids.append(int(patch[len(stem) + 1:]))

        patch_ids.sort()

        patches = []

        for id in patch_ids:

            if id == 0:

                patches.append(stem)

            else:

                patches.append('%s-%s' % (stem, id))

        return patches

    def upgrade(self):

        patches = self._list_old_shelves()

        if len(patches) == 0:

            self.log('No old-style shelves found to upgrade.\n')

            return

        for patch in patches:

            old_file = open(patch, 'r')

            new_path = self.next_patch()

            new_file = open(new_path, 'w')

            new_file.write(old_file.read())

            old_file.close()

            new_file.close()

            self.log('Copied %s to %s/%s\n' % (os.path.basename(patch),

                self.name, os.path.basename(new_path)))

            os.rename(patch, patch + '~')