~abentley/bzrtools/bzrtools.dev

# Copyright (C) 2005, 2006, 2007 Aaron Bentley <aaron.bentley@utoronto.ca>
# 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 trace, urlutils
import bzrlib.errors
from bzrlib.errors import (
    BzrCommandError,
    BzrError,
    ConnectionError,
    NotBranchError,
    NoSuchFile,
    NoWorkingTree,
    PermissionDenied,
    UnsupportedFormatError,
    TransportError,
    )
from bzrlib.bzrdir import BzrDir, BzrDirFormat
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.control_files.put_utf8("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.control_files.get_utf8('x-push-data').read()
    except NoSuchFile:
        return None
    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"):
    """
    >>> 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(new_dir + "/a", new_dir + "/b", excludes=("*.py",), silent=True)
    Traceback (most recent call last):
    RsyncNoFile: No such file...
    >>> rsync(new_dir + "/a", new_dir + "/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/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()]

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

    _format = '%(location) 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 history_subset(location, branch, _rsync=None):
    remote_history = get_revision_history(location, _rsync)
    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 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 bzrlib.errors.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 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)
        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 bzrlib.errors.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 errors.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)
    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()