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

    """

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

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