~abentley/bzrtools/bzrtools.dev

import bzrlib

import os

import os.path

import sys

import tempfile

import shutil

def temp_branch():

    dirname = tempfile.mkdtemp("temp-branch")

    return bzrlib.Branch(dirname, init=True)

def rm_branch(br):

    shutil.rmtree(br.base)

def is_clean(cur_branch):

    """

    Return true if no files are modifed or unknown

    >>> br = temp_branch()

    >>> is_clean(br)

    True

    >>> fooname = os.path.join(br.base, "foo")

    >>> file(fooname, "wb").write("bar")

    >>> is_clean(br)

    False

    >>> bzrlib.add.smart_add([fooname])

    >>> is_clean(br)

    False

    >>> br.commit("added file")

    >>> is_clean(br)

    True

    >>> rm_branch(br)

    """

    old_tree = cur_branch.basis_tree()

    new_tree = cur_branch.working_tree()

    for path, file_class, kind, file_id in new_tree.list_files():

        if file_class == '?':

            return False

    delta = bzrlib.compare_trees(old_tree, new_tree, want_unchanged=False)

    if len(delta.added) > 0 or len(delta.removed) > 0 or \

        len(delta.modified) > 0:

        return False

    return True

def set_pull_data(br, location, rev_id):

    pull_file = file (br.controlfilename("x-pull-data"), "wb")

    pull_file.write("%s\n%s\n" % (location, rev_id))

def get_pull_data(br):

    """

    >>> br = temp_branch()

    >>> get_pull_data(br)

    (None, None)

    >>> set_pull_data(br, 'http://somewhere', '888-777')

    >>> get_pull_data(br)

    ('http://somewhere', '888-777')

    >>> rm_branch(br)

    """

    filename = br.controlfilename("x-pull-data")

    if not os.path.exists(filename):

        return (None, None)

    pull_file = file (filename, "rb")

    location, rev_id = [f.rstrip('\n') for f in pull_file]

    return location, rev_id

def set_push_data(br, location):

    push_file = file (br.controlfilename("x-push-data"), "wb")

    push_file.write("%s\n" % location)

def get_push_data(br):

    """

    >>> br = temp_branch()

    >>> get_push_data(br) is None

    True

    >>> set_push_data(br, 'http://somewhere')

    >>> get_push_data(br)

    'http://somewhere'

    >>> rm_branch(br)

    """

    filename = br.controlfilename("x-push-data")

    if not os.path.exists(filename):

        return None

    push_file = file (filename, "rb")

    (location,) = [f.rstrip('\n') for f in push_file]

    return location

"""

>>> shell_escape('hello')

'\h\e\l\l\o'

"""

def shell_escape(arg):

    return "".join(['\\'+c for c in arg])

def safe_system(args):

    """

    >>> real_system = os.system

    >>> os.system = sys.stdout.write

    >>> safe_system(['a', 'b', 'cd'])

    \\a \\b \\c\\d

    >>> os.system = real_system

    """

    arg_str = " ".join([shell_escape(a) for a in args])

    return os.system(arg_str)

def rsync(source, target, ssh=False, exclude_globs=()):

    """

    >>> real_system = os.system

    >>> os.system = sys.stdout.write

    >>> rsync("a", "b")

    \\r\\s\\y\\n\\c \\-\\a\\v \\-\\-\\d\\e\\l\\e\\t\\e \\a \\b

    >>> rsync("a", "b", exclude_globs=("*.py",))

    \\r\\s\\y\\n\\c \\-\\a\\v \\-\\-\\d\\e\\l\\e\\t\\e\

 \\-\\-\\e\\x\\c\\l\\u\\d\\e \\*\\.\\p\\y \\a \\b

    >>> os.system = real_system

    """

    cmd = ["rsync", "-av", "--delete"]

    if ssh:

        cmd.extend(('-e', 'ssh'))

    for exclude in exclude_globs:

        cmd.extend(('--exclude', exclude))

    cmd.extend((source, target))

    safe_system(cmd)

exclusions = ('x-push-data', 'x-pull-data')

def pull(cur_branch, location=None, overwrite=False):

    pull_location, pull_revision = get_pull_data(cur_branch)

    if pull_location is not None:

        if not overwrite and cur_branch.last_patch() != pull_revision:

            print "Aborting: This branch has had commits, so pull would lose data."

            sys.exit(1)

    if location is not None:

        pull_location = location

        if not pull_location.endswith('/'):

            pull_location+='/'

    if pull_location is None:

        print "No pull location saved.  Please specify one on the command line."

        sys.exit(1)

    if not is_clean(cur_branch):

        print "Error: This tree has uncommitted changes or unknown (?) files."

        sys.exit(1)

    print "Synchronizing with %s" % pull_location

    rsync (pull_location, cur_branch.base+'/', exclude_globs=exclusions)

    set_pull_data(cur_branch, pull_location, cur_branch.last_patch())

def push(cur_branch, location=None):

    push_location = get_push_data(cur_branch)

    if location is not None:

        if not location.endswith('/'):

            location += '/'

        push_location = location

    if push_location is None:

        print "No push location saved.  Please specify one on the command line."

        sys.exit(1)

    if not is_clean(cur_branch):

        print "Error: This tree has uncommitted changes or unknown (?) files."

        sys.exit(1)

    print "Pushing to %s" % push_location

    rsync(cur_branch.base+'/', push_location, ssh=True,

          exclude_globs=exclusions)

    set_push_data(cur_branch, push_location)

def run_tests():

    import doctest

    result = doctest.testmod()

    if result[1] > 0:

        if result[0] == 0:

            print "All tests passed"

    else:

        print "No tests to run"

if __name__ == "__main__":

    run_tests()