~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, NotBranchError, NoSuchFile,
                           UnsupportedFormatError, TransportError, 
                           NoWorkingTree, PermissionDenied)
from bzrlib.bzrdir import BzrDir, BzrDirFormat

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", rev_id='commit-id')
    'commit-id'
    >>> is_clean(tree)
    (True, [])
    >>> rm_tree(tree)
    """
    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 = new_tree.changes_from(old_tree, want_unchanged=False)
    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.")

def get_revision_history(location):
    tempdir = tempfile.mkdtemp('push')
    try:
        history_fname = os.path.join(tempdir, 'revision-history')
        try:
            cmd = 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):
    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 rspush(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:
        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:
            print """Error: 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):
                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=usessh, 
          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 apache_ls(t):
    """Screen-scrape Apache listings"""
    apache_dir = '<img border="0" src="/icons/folder.gif" alt="[dir]">'\
        ' <a href="'
    lines = t.get('.')
    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 iter_branches(t, lister=None):
    """Iterate through all the branches under a transport"""
    for bzrdir in iter_bzrdirs(t, lister):
        try:
            branch = bzrdir.open_branch()
            if branch.bzrdir is bzrdir:
                yield branch
        except (NotBranchError, UnsupportedFormatError):
            pass


def iter_branch_tree(t, lister=None):
    for bzrdir in iter_bzrdirs(t, lister):
        try:
            wt = bzrdir.open_workingtree()
            yield wt.branch, wt
        except NoWorkingTree, UnsupportedFormatError:
            try:
                branch = bzrdir.open_branch()
                if branch.bzrdir is bzrdir:
                    yield branch, None
            except (NotBranchError, UnsupportedFormatError):
                continue


def iter_bzrdirs(t, lister=None):
    if lister is None:
        def lister(t):
            return t.list_dir('.')
    try:
        bzrdir = bzrdir_from_transport(t)
        yield bzrdir
    except (NotBranchError, UnsupportedFormatError, TransportError,
            PermissionDenied):
        pass
    try:
        for directory in lister(t):
            if directory == ".bzr":
                continue
            try:
                subt = t.clone(directory)
            except UnicodeDecodeError:
                continue
            for bzrdir in iter_bzrdirs(subt, lister):
                yield bzrdir
    except (NoSuchFile, PermissionDenied, TransportError):
        pass

    
def bzrdir_from_transport(t):
    """Open a bzrdir from a transport (not a location)"""
    format = BzrDirFormat.find_format(t)
    BzrDir._check_supported(format, False)
    return format.open(t)


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