1
from bzrlib.bzrdir import BzrDir
2
from bzrlib.errors import NotBranchError, NoSuchFile, BzrError
3
from bzrlib.rio import RioReader, rio_file, Stanza
5
def branch_mark(mark, location, delete=False):
8
bzrdir = BzrDir.open_containing(location)[0]
10
branch = bzrdir.open_branch()
11
repository = branch.repository
12
except NotBranchError:
14
repository = bzrdir.open_repository()
17
raise BzrError("Please specify a mark to delete")
18
unset_mark(branch, mark)
20
print_marks(repository)
22
set_mark(branch, mark)
24
def print_marks(repository):
25
marks = get_marks(repository)
26
for mark in sorted(marks.keys()):
27
branches = marks[mark]
28
assert len(branches) != 0
30
for branch in branches:
34
MARK_FORMAT_1 = "Branch Marks Format 1"
36
def get_rio_stanzas(resource, filename, missing_to_empty=True):
38
riofile = resource.get(filename)
40
if not missing_to_empty:
43
header = riofile.next().rstrip('\n')
44
return header, list(RioReader(riofile))
47
def get_marks(repository):
48
header, stanzas = get_rio_stanzas(repository.control_files, 'branch-marks')
51
if header != MARK_FORMAT_1:
52
raise BzrError("Unknown mark format: %s" % header)
54
for stanza in stanzas:
58
marks[mark].append(stanza['branch'])
62
def relative_base(branch):
63
return branch.repository.bzrdir.transport.clone('..').relpath(branch.base)
65
def set_mark(branch, mark):
69
marks[mark].append(relative_base(branch))
71
return _set_mark(branch, mark, add)
73
def unset_mark(branch, mark):
76
marks[mark].remove(relative_base(branch))
78
except KeyError, ValueError:
79
raise BzrError("Branch does not have this mark set.")
80
return _set_mark(branch, mark, remove)
82
def _set_mark(branch, mark, mutate):
83
branch.repository.lock_write()
85
marks = get_marks(branch.repository)
87
set_marks(branch.repository, marks)
89
branch.repository.unlock()
91
def set_marks(repository, marks):
92
riofile = rio_file(mark_stanzas(marks), MARK_FORMAT_1)
93
repository.control_files.put('branch-marks', riofile)
95
def mark_stanzas(marks):
96
for mark,branches in marks.iteritems():
97
for branch in branches:
98
yield Stanza(mark=mark, branch=branch)