~abentley/bzrtools/bzrtools.dev

# Copyright (C) 2005 Aaron Bentley

# <aaron.bentley@utoronto.ca>

#

#    This program is free software; you can redistribute it and/or modify

#    it under the terms of the GNU General Public License as published by

#    the Free Software Foundation; either version 2 of the License, or

#    (at your option) any later version.

#

#    This program is distributed in the hope that it will be useful,

#    but WITHOUT ANY WARRANTY; without even the implied warranty of

#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the

#    GNU General Public License for more details.

#

#    You should have received a copy of the GNU General Public License

#    along with this program; if not, write to the Free Software

#    Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA

import bzrlib

import bzrlib.errors

import os

import os.path

import sys

import tempfile

import shutil

import errno

from subprocess import Popen, PIPE

import codecs

def temp_branch():

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

    return bzrlib.branch.Branch.initialize(dirname)

def rm_branch(br):

    shutil.rmtree(br.base)

def is_clean(cur_branch):

    """

    Return true if no files are modifed or unknown

    >>> import bzrlib.add

    >>> br = temp_branch()

    >>> is_clean(br)

    (True, [])

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

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

    >>> is_clean(br)

    (True, [u'foo'])

    >>> bzrlib.add.smart_add_branch(br, [br.base])

    1

    >>> is_clean(br)

    (False, [])

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

    >>> is_clean(br)

    (True, [])

    >>> rm_branch(br)

    """

    from bzrlib.diff import compare_trees

    old_tree = cur_branch.basis_tree()

    new_tree = cur_branch.working_tree()

    non_source = []

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

        if file_class in ('?', 'I'):

            non_source.append(path)

    delta = 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, non_source

    return True, non_source 

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)

class RsyncUnknownStatus(Exception):

    def __init__(self, status):

        Exception.__init__(self, "Unknown status: %d" % status)

class NoRsync(Exception):

    def __init__(self, rsync_name):

        Exception.__init__(self, "%s not found." % rsync_name)

def rsync(source, target, ssh=False, excludes=(), silent=False, 

          rsync_name="rsync"):

    """

    >>> rsync("a", "b", silent=True)

    Traceback (most recent call last):

    RsyncNoFile: No such file a

    >>> rsync("a", "b", excludes=("*.py",), silent=True)

    Traceback (most recent call last):

    RsyncNoFile: No such file a

    >>> rsync("a", "b", excludes=("*.py",), silent=True, rsync_name="rsyncc")

    Traceback (most recent call last):

    NoRsync: rsyncc not found.

    """

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

    if ssh:

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

    if len(excludes) > 0:

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

    cmd.extend((source, target))

    if silent:

        stderr = PIPE

        stdout = PIPE

    else:

        stderr = None

        stdout = None

    try:

        proc = Popen(cmd, stdin=PIPE, stderr=stderr, stdout=stdout)

    except OSError, e:

        if e.errno == errno.ENOENT:

            raise NoRsync(rsync_name)

    proc.stdin.write('\n'.join(excludes)+'\n')

    proc.stdin.close()

    if silent:

        proc.stderr.read()

        proc.stderr.close()

        proc.stdout.read()

        proc.stdout.close()

    proc.wait()

    if proc.returncode == 12:

        raise RsyncStreamIO()

    elif proc.returncode == 23:

        raise RsyncNoFile(source)

    elif proc.returncode != 0:

        raise RsyncUnknownStatus(proc.returncode)

    return cmd

def rsync_ls(source, ssh=False, silent=True):

    cmd = ["rsync"]

    if ssh:

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

    cmd.append(source)

    if silent:

        stderr = PIPE

    else:

        stderr = None

    proc = Popen(cmd, stderr=stderr, stdout=PIPE)

    result = proc.stdout.read()

    proc.stdout.close()

    if silent:

        proc.stderr.read()

        proc.stderr.close()

    proc.wait()

    if proc.returncode == 12:

        raise RsyncStreamIO()

    elif proc.returncode == 23:

        raise RsyncNoFile(source)

    elif proc.returncode != 0:

        raise RsyncUnknownStatus(proc.returncode)

    return [l.split(' ')[-1].rstrip('\n') for l in result.splitlines(True)]

exclusions = ('.bzr/x-push-data', '.bzr/parent', '.bzr/x-pull-data', 

              '.bzr/x-pull', '.bzr/pull', '.bzr/stat-cache',

              '.bzr/x-rsync-data')

def read_revision_history(fname):

    return [l.rstrip('\r\n') for l in

            codecs.open(fname, 'rb', 'utf-8').readlines()]

class RsyncNoFile(Exception):

    def __init__(self, path):

        Exception.__init__(self, "No such file %s" % path)

class RsyncStreamIO(Exception):

    def __init__(self):

        Exception.__init__(self, "Error in rsync protocol data stream.")

def get_revision_history(location):

    tempdir = tempfile.mkdtemp('push')

    try:

        history_fname = os.path.join(tempdir, 'revision-history')

        cmd = rsync(location+'.bzr/revision-history', history_fname,

                    silent=True)

        history = read_revision_history(history_fname)

    finally:

        shutil.rmtree(tempdir)

    return history

def history_subset(location, branch):

    remote_history = get_revision_history(location)

    local_history = branch.revision_history()

    if len(remote_history) > len(local_history):

        return False

    for local, remote in zip(remote_history, local_history):

        if local != remote:

            return False 

    return True

def empty_or_absent(location):

    try:

        files = rsync_ls(location)

        return files == ['.']

    except RsyncNoFile:

        return True

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

    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)

    clean, non_source = is_clean(cur_branch)

    if not clean:

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

Use "bzr status" to list them."""

        sys.exit(1)

    non_source.extend(exclusions)

    if not overwrite:

        try:

            if not history_subset(push_location, cur_branch):

                raise bzrlib.errors.BzrCommandError("Local branch is not a"

                                                    " newer version of remote"

                                                    " branch.")

        except RsyncNoFile:

            if not empty_or_absent(push_location):

                raise bzrlib.errors.BzrCommandError("Remote location is not a"

                                                    " bzr branch (or empty"

                                                    " directory)")

        except RsyncStreamIO:

            raise bzrlib.errors.BzrCommandError("Rsync could not use the"

                " specified location.  Please ensure that"

                ' "%s" is of the form "machine:/path".' % push_location)

    print "Pushing to %s" % push_location

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

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