~abentley/bzrtools/bzrtools.dev

from bzrlib.bzrdir import BzrDir

from bzrlib.errors import NotBranchError, NoSuchFile, BzrError

from bzrlib.rio import RioReader, rio_file, Stanza

def branch_mark(mark, location, delete=False):

    if location is None:

        location = '.'

    bzrdir = BzrDir.open_containing(location)[0]

    try:

        branch = bzrdir.open_branch()

        repository = branch.repository

    except NotBranchError:

        branch = None

        repository = bzrdir.open_repository()

    if delete is True:

        if mark is None:

            raise BzrError("Please specify a mark to delete")

        unset_mark(branch, mark)

    elif mark is None:

        print_marks(repository)

    else:

        set_mark(branch, mark)

def print_marks(repository):

    marks = get_marks(repository)

    for mark in sorted(marks.keys()):

        branches = marks[mark]

        assert len(branches) != 0

        print mark

        for branch in branches:

            print "  %s" % branch

MARK_FORMAT_1 = "Branch Marks Format 1"

def get_rio_stanzas(resource, filename, missing_to_empty=True):

    try:

        riofile = resource.get(filename)

    except NoSuchFile:

        if not missing_to_empty:

            raise

        return None, []

    header = riofile.next().rstrip('\n')

    return header, list(RioReader(riofile))

def get_marks(repository):

    header, stanzas = get_rio_stanzas(repository.control_files, 'branch-marks')

    if len(stanzas) == 0:

        return {}

    if header != MARK_FORMAT_1:

        raise BzrError("Unknown mark format: %s" % header)

    marks = {}

    for stanza in stanzas:

        mark = stanza['mark']

        if mark not in marks:

            marks[mark] = []

        marks[mark].append(stanza['branch'])

        marks[mark].sort()

    return marks

def relative_base(branch):

    return branch.repository.bzrdir.transport.clone('..').relpath(branch.base)

def set_mark(branch, mark):

    def add(marks):

        if mark not in marks:

            marks[mark] = []

        marks[mark].append(relative_base(branch))

        return marks

    return _set_mark(branch, mark, add)

def unset_mark(branch, mark):

    def remove(marks):

        try:

            marks[mark].remove(relative_base(branch))

            return marks

        except KeyError, ValueError:

            raise BzrError("Branch does not have this mark set.")

    return _set_mark(branch, mark, remove)

def _set_mark(branch, mark, mutate):

    branch.repository.lock_write()

    try:

        marks = get_marks(branch.repository)

        marks = mutate(marks)

        set_marks(branch.repository, marks)

    finally:

        branch.repository.unlock()

def set_marks(repository, marks):

    riofile = rio_file(mark_stanzas(marks), MARK_FORMAT_1)

    repository.control_files.put('branch-marks', riofile)

def mark_stanzas(marks):

    for mark,branches in marks.iteritems():

        for branch in branches:

            yield Stanza(mark=mark, branch=branch)