~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 codecs
import errno
import os
import re
import tempfile
import shutil
from subprocess import Popen, PIPE
import sys

import bzrlib
import bzrlib.errors
from bzrlib.errors import BzrCommandError
from bzrlib.bzrdir import BzrDir

def temp_tree():
    dirname = tempfile.mkdtemp("temp-branch")
    return BzrDir.create_standalone_workingtree(dirname)

def rm_tree(tree):
    shutil.rmtree(tree.basedir)

def is_clean(cur_tree):
    """
    Return true if no files are modifed or unknown
    >>> import bzrlib.add
    >>> tree = temp_tree()
    >>> is_clean(tree)
    (True, [])
    >>> fooname = os.path.join(tree.basedir, "foo")
    >>> file(fooname, "wb").write("bar")
    >>> is_clean(tree)
    (True, [u'foo'])
    >>> bzrlib.add.smart_add_tree(tree, [tree.basedir])
    ([u'foo'], {})
    >>> is_clean(tree)
    (False, [])
    >>> tree.commit("added file")
    >>> is_clean(tree)
    (True, [])
    >>> rm_tree(tree)
    """
    from bzrlib.diff import compare_trees
    old_tree = cur_tree.basis_tree()
    new_tree = cur_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)
    return not delta.has_changed(), non_source

def set_push_data(tree, location):
    push_file = file (tree._control_files.controlfilename("x-push-data"), "wb")
    push_file.write("%s\n" % location)

def get_push_data(tree):
    """
    >>> tree = temp_tree()
    >>> get_push_data(tree) is None
    True
    >>> set_push_data(tree, 'http://somewhere')
    >>> get_push_data(tree)
    'http://somewhere'
    >>> rm_tree(tree)
    """
    filename = tree._control_files.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"):
    """
    >>> new_dir = tempfile.mkdtemp()
    >>> old_dir = os.getcwd()
    >>> os.chdir(new_dir)
    >>> rsync("a", "b", silent=True)
    Traceback (most recent call last):
    RsyncNoFile: No such file...
    >>> rsync(tempdir + "/a", tempdir + "/b", excludes=("*.py",), silent=True) #doctest: +ELLIPSIS
    Traceback (most recent call last):
    RsyncNoFile: No such file...
    >>> rsync(tempdir + "/a", tempdir + "/b", excludes=("*.py",), silent=True, rsync_name="rsyncc")
    Traceback (most recent call last):
    NoRsync: rsyncc not found.
    >>> os.chdir(old_dir)
    >>> os.rmdir(new_dir)
    """
    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(tree, location=None, overwrite=False, working_tree=True):
    push_location = get_push_data(tree)
    if location is not None:
        if not location.endswith('/'):
            location += '/'
        push_location = location
    
    if push_location is None:
        if tree.branch.get_push_location() is None:
            raise BzrCommandError("No push location known or specified.")
        else:
            raise bzrlib.errors.MustUseDecorated

    if push_location.find('://') != -1:
        raise bzrlib.errors.MustUseDecorated

    if push_location.find(':') == -1:
        raise bzrlib.errors.MustUseDecorated

    clean, non_source = is_clean(tree)
    if not clean:
        print """Error: This tree has uncommitted changes or unknown (?) files.
Use "bzr status" to list them."""
        sys.exit(1)
    if working_tree:
        final_exclusions = non_source[:]
    else:
        wt = tree
        final_exclusions = []
        for path, status, kind, file_id, entry in wt.list_files():
            final_exclusions.append(path)

    final_exclusions.extend(exclusions)
    if not overwrite:
        try:
            if not history_subset(push_location, tree.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(tree.basedir+'/', push_location, ssh=True, 
          excludes=final_exclusions)

    set_push_data(tree, push_location)


def short_committer(committer):
    new_committer = re.sub('<.*>', '', committer).strip(' ')
    if len(new_committer) < 2:
        return committer
    return new_committer


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