~abentley/bzrtools/bzrtools.dev

# Copyright (C) 2005, 2006, 2007 Aaron Bentley <aaron@aaronbentley.com>
# Copyright (C) 2007 John Arbash Meinel
#
#    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
from bzrlib import revision as _mod_revision, trace, urlutils
from bzrlib.errors import (
    BzrCommandError,
    BzrError,
    NotBranchError,
    NoSuchFile,
    )
from bzrlib.bzrdir import BzrDir
from bzrlib.transport import get_transport

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
    """
    old_tree = cur_tree.basis_tree()
    new_tree = cur_tree
    non_source = []
    new_tree.lock_read()
    try:
        for path, file_class, kind, file_id, entry in new_tree.list_files():
            if file_class in ('?', 'I'):
                non_source.append(path)
        delta = new_tree.changes_from(old_tree, want_unchanged=False)
    finally:
        new_tree.unlock()
    return not delta.has_changed(), non_source

def set_push_data(tree, location):
    tree.branch._transport.put_bytes("x-push-data", "%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)
    u'http://somewhere'
    >>> rm_tree(tree)
    """
    try:
        location = tree.branch._transport.get('x-push-data').read()
    except NoSuchFile:
        return None
    location = location.decode('utf-8')
    return location.rstrip('\n')

"""
>>> 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"):
    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/branch/x-push/data', '.bzr/parent',
              '.bzr/branch/parent', '.bzr/x-pull-data', '.bzr/x-pull',
              '.bzr/pull', '.bzr/stat-cache', '.bzr/x-rsync-data',
              '.bzr/basis-inventory', '.bzr/inventory.backup.weave')


def read_revision_history(fname):
    return [l.rstrip('\r\n') for l in
            codecs.open(fname, 'rb', 'utf-8').readlines()]


def read_revision_info(path):
    """Parse a last_revision file to determine revision_info"""
    line = open(path, 'rb').readlines()[0].strip('\n')
    revno, revision_id = line.split(' ', 1)
    revno = int(revno)
    return revno, revision_id


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.")


class NotStandalone(BzrError):

    _fmt = '%(location)s is not a standalone tree.'
    _internal = False

    def __init__(self, location):
        BzrError.__init__(self, location=location)


def get_revision_history(location, _rsync):
    tempdir = tempfile.mkdtemp('push')
    my_rsync = _rsync
    if my_rsync is None:
        my_rsync = rsync
    try:
        history_fname = os.path.join(tempdir, 'revision-history')
        try:
            cmd = my_rsync(location+'.bzr/revision-history', history_fname,
                        silent=True)
        except RsyncNoFile:
            cmd = rsync(location+'.bzr/branch/revision-history', history_fname,
                        silent=True)
        history = read_revision_history(history_fname)
    finally:
        shutil.rmtree(tempdir)
    return history


def get_revision_info(location, _rsync):
    """Get the revsision_info for an rsync-able branch"""
    tempdir = tempfile.mkdtemp('push')
    my_rsync = _rsync
    if my_rsync is None:
        my_rsync = rsync
    try:
        info_fname = os.path.join(tempdir, 'last-revision')
        cmd = rsync(location+'.bzr/branch/last-revision', info_fname,
                    silent=True)
        return read_revision_info(info_fname)
    finally:
        shutil.rmtree(tempdir)


def history_subset(location, branch, _rsync=None):
    local_history = branch.revision_history()
    try:
        remote_history = get_revision_history(location, _rsync)
    except RsyncNoFile:
        revno, revision_id = get_revision_info(location, _rsync)
        if revision_id == _mod_revision.NULL_REVISION:
            return True
        return bool(revision_id.decode('utf-8') in local_history)
    else:
        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 rspush(tree, location=None, overwrite=False, working_tree=True,
    _rsync=None):
    tree.lock_write()
    try:
        my_rsync = _rsync
        if my_rsync is None:
            my_rsync = rsync
        if (tree.bzrdir.root_transport.base !=
            tree.branch.bzrdir.root_transport.base):
            raise NotStandalone(tree.bzrdir.root_transport.base)
        if (tree.branch.get_bound_location() is not None):
            raise NotStandalone(tree.bzrdir.root_transport.base)
        if (tree.branch.repository.is_shared()):
            raise NotStandalone(tree.bzrdir.root_transport.base)
        push_location = get_push_data(tree)
        if location is not None:
            if not location.endswith('/'):
                location += '/'
            push_location = location

        if push_location is None:
            raise BzrCommandError("No rspush location known or specified.")

        if (push_location.find('::') != -1):
            usessh=False
        else:
            usessh=True

        if (push_location.find('://') != -1 or
            push_location.find(':') == -1):
            raise BzrCommandError("Invalid rsync path %r." % push_location)

        if working_tree:
            clean, non_source = is_clean(tree)
            if not clean:
                raise BzrCommandError(
                    'This tree has uncommitted changes or unknown'
                    ' (?) files.  Use "bzr status" to list them.')
                sys.exit(1)
            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,
                                      _rsync=my_rsync):
                    raise BzrCommandError(
                        "Local branch is not a newer version of remote"
                        " branch.")
            except RsyncNoFile:
                if not empty_or_absent(push_location):
                    raise BzrCommandError(
                        "Remote location is not a bzr branch (or empty"
                        " directory)")
            except RsyncStreamIO:
                raise BzrCommandError("Rsync could not use the"
                    " specified location.  Please ensure that"
                    ' "%s" is of the form "machine:/path".' % push_location)
        trace.note("Pushing to %s", push_location)
        my_rsync(tree.basedir+'/', push_location, ssh=usessh,
                 excludes=final_exclusions)

        set_push_data(tree, push_location)
    finally:
        tree.unlock()


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


def apache_ls(t):
    """Screen-scrape Apache listings"""
    apache_dir = '<img border="0" src="/icons/folder.gif" alt="[dir]">'\
        ' <a href="'
    t = t.clone()
    t._remote_path = lambda x: t.base
    try:
        lines = t.get('')
    except NoSuchFile:
        return
    expr = re.compile('<a[^>]*href="([^>]*)\/"[^>]*>', flags=re.I)
    for line in lines:
        match = expr.search(line)
        if match is None:
            continue
        url = match.group(1)
        if url.startswith('http://') or url.startswith('/') or '../' in url:
            continue
        if '?' in url:
            continue
        yield url.rstrip('/')


def list_branches(t):
    def is_inside(branch):
        return bool(branch.base.startswith(t.base))

    if t.base.startswith('http://'):
        def evaluate(bzrdir):
            try:
                branch = bzrdir.open_branch()
                if is_inside(branch):
                    return True, branch
                else:
                    return True, None
            except NotBranchError:
                return True, None
        return [b for b in BzrDir.find_bzrdirs(t, list_current=apache_ls,
                evaluate=evaluate) if b is not None]
    elif not t.listable():
        raise BzrCommandError("Can't list this type of location.")
    return [b for b in BzrDir.find_branches(t) if is_inside(b)]


def evaluate_branch_tree(bzrdir):
    try:
        tree, branch = bzrdir._get_tree_branch()
    except NotBranchError:
        return True, None
    else:
        return True, (branch, tree)


def iter_branch_tree(t, lister=None):
    return (x for x in BzrDir.find_bzrdirs(t, evaluate=evaluate_branch_tree,
            list_current=lister) if x is not None)


def open_from_url(location):
    location = urlutils.normalize_url(location)
    dirname, basename = urlutils.split(location)
    if location.endswith('/') and not basename.endswith('/'):
        basename += '/'
    return get_transport(dirname).get(basename)


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