~abentley/bzrtools/bzrtools.dev

#!/usr/bin/env python

# 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

try:

    import pybaz

    import pybaz.errors

except ImportError:

    print "This command requires PyBaz.  Please ensure that it is installed."

    import sys

    sys.exit(1)

from pybaz.backends.baz import null_cmd

import tempfile

import os

import os.path

import shutil

import bzrlib

from bzrlib.errors import BzrError

import bzrlib.trace

import bzrlib.merge

import sys

import email.Utils

from progress import *

def add_id(files, id=None):

    """Adds an explicit id to a list of files.

    :param files: the name of the file to add an id to

    :type files: list of str

    :param id: tag one file using the specified id, instead of generating id

    :type id: str

    """

    args = ["add-id"]

    if id is not None:

        args.extend(["--id", id])

    args.extend(files)

    return null_cmd(args)

def test_environ():

    """

    >>> q = test_environ()

    >>> os.path.exists(q)

    True

    >>> os.path.exists(os.path.join(q, "home", ".arch-params"))

    True

    >>> teardown_environ(q)

    >>> os.path.exists(q)

    False

    """

    tdir = tempfile.mkdtemp(prefix="testdir-")

    os.environ["HOME"] = os.path.join(tdir, "home")

    os.mkdir(os.environ["HOME"])

    arch_dir = os.path.join(tdir, "archive_dir")

    pybaz.make_archive("test@example.com", arch_dir)

    work_dir = os.path.join(tdir, "work_dir")

    os.mkdir(work_dir)

    os.chdir(work_dir)

    pybaz.init_tree(work_dir, "test@example.com/test--test--0")

    lib_dir = os.path.join(tdir, "lib_dir")

    os.mkdir(lib_dir)

    pybaz.register_revision_library(lib_dir)

    pybaz.set_my_id("Test User<test@example.org>")

    return tdir

def add_file(path, text, id):

    """

    >>> q = test_environ()

    >>> add_file("path with space", "text", "lalala")

    >>> tree = pybaz.tree_root(".")

    >>> inv = list(tree.iter_inventory_ids(source=True, both=True))

    >>> ("x_lalala", "path with space") in inv

    True

    >>> teardown_environ(q)

    """

    file(path, "wb").write(text)

    add_id([path], id)

def add_dir(path, id):

    """

    >>> q = test_environ()

    >>> add_dir("path with\(sp) space", "lalala")

    >>> tree = pybaz.tree_root(".")

    >>> inv = list(tree.iter_inventory_ids(source=True, both=True))

    >>> ("x_lalala", "path with\(sp) space") in inv

    True

    >>> teardown_environ(q)

    """

    os.mkdir(path)

    add_id([path], id)

def teardown_environ(tdir):

    os.chdir("/")

    shutil.rmtree(tdir)

def timport(tree, summary):

    msg = tree.log_message()

    msg["summary"] = summary

    tree.import_(msg)

def commit(tree, summary):

    """

    >>> q = test_environ()

    >>> tree = pybaz.tree_root(".")

    >>> timport(tree, "import")

    >>> commit(tree, "commit")

    >>> logs = [str(l.revision) for l in tree.iter_logs()]

    >>> len(logs)

    2

    >>> logs[0]

    'test@example.com/test--test--0--base-0'

    >>> logs[1]

    'test@example.com/test--test--0--patch-1'

    >>> teardown_environ(q)

    """

    msg = tree.log_message()

    msg["summary"] = summary

    tree.commit(msg)

def commit_test_revisions():

    """

    >>> q = test_environ()

    >>> commit_test_revisions()

    >>> a = pybaz.Archive("test@example.com")

    >>> revisions = list(a.iter_revisions("test--test--0"))

    >>> len(revisions)

    3

    >>> str(revisions[2])

    'test@example.com/test--test--0--base-0'

    >>> str(revisions[1])

    'test@example.com/test--test--0--patch-1'

    >>> str(revisions[0])

    'test@example.com/test--test--0--patch-2'

    >>> teardown_environ(q)

    """

    tree = pybaz.tree_root(".")

    add_file("mainfile", "void main(void){}", "mainfile by aaron")

    timport(tree, "Created mainfile")

    file("mainfile", "wb").write("or something like that")

    commit(tree, "altered mainfile")

    add_file("ofile", "this is another file", "ofile by aaron")

    commit(tree, "altered mainfile")

def commit_more_test_revisions():

    """

    >>> q = test_environ()

    >>> commit_test_revisions()

    >>> commit_more_test_revisions()

    >>> a = pybaz.Archive("test@example.com")

    >>> revisions = list(a.iter_revisions("test--test--0"))

    >>> len(revisions)

    4

    >>> str(revisions[0])

    'test@example.com/test--test--0--patch-3'

    >>> teardown_environ(q)

    """

    tree = pybaz.tree_root(".")

    add_file("trainfile", "void train(void){}", "trainfile by aaron")

    commit(tree, "altered trainfile")

class NoSuchVersion(Exception):

    def __init__(self, version):

        Exception.__init__(self, "The version %s does not exist." % version)

        self.version = version

def version_ancestry(version):

    """

    >>> q = test_environ()

    >>> commit_test_revisions()

    >>> version = pybaz.Version("test@example.com/test--test--0")

    >>> ancestors = version_ancestry(version)

    >>> str(ancestors[0])

    'test@example.com/test--test--0--base-0'

    >>> str(ancestors[1])

    'test@example.com/test--test--0--patch-1'

    >>> version = pybaz.Version("test@example.com/test--test--0.5")

    >>> ancestors = version_ancestry(version)

    Traceback (most recent call last):

    NoSuchVersion: The version test@example.com/test--test--0.5 does not exist.

    >>> teardown_environ(q)

    """

    try:

        revision = version.iter_revisions(reverse=True).next()

    except:

        if not version.exists():

            raise NoSuchVersion(version)

        else:

            raise

    ancestors = list(revision.iter_ancestors(metoo=True))

    ancestors.reverse()

    return ancestors

def get_last_revision(branch):

    last_patch = branch.last_patch()

    try:

        return arch_revision(last_patch)

    except NotArchRevision:

        raise UserError(

            "Directory \"%s\" already exists, and the last revision is not"

            " an Arch revision (%s)" % (output_dir, last_patch))

def get_remaining_revisions(output_dir, version):

    last_patch = None

    old_revno = None

    if os.path.exists(output_dir):

        # We are starting from an existing directory, figure out what

        # the current version is

        branch = find_branch(output_dir)

        last_patch = get_last_revision(branch)

        if version is None:

            version = last_patch.version

    elif version is None:

        raise UserError("No version specified, and directory does not exist.")

    try:

        ancestors = version_ancestry(version)

    except NoSuchVersion, e:

        raise UserError(e)

    if last_patch:

        for i in range(len(ancestors)):

            if ancestors[i] == last_patch:

                break

        else:

            raise UserError("Directory \"%s\" already exists, and the last "

                "revision (%s) is not in the ancestry of %s" % 

                (output_dir, last_patch, version))

        # Strip off all of the ancestors which are already present

        # And get a directory starting with the latest ancestor

        latest_ancestor = ancestors[i]

        old_revno = find_branch(output_dir).revno()

        ancestors = ancestors[i+1:]

    return ancestors, old_revno

def import_version(output_dir, version, fancy=True, fast=False, verbose=False, 

                   dry_run=False, max_count=None, skip_symlinks=False):

    """

    >>> q = test_environ()

    >>> result_path = os.path.join(q, "result")

    >>> commit_test_revisions()

    >>> version = pybaz.Version("test@example.com/test--test--0.1")

    >>> import_version('/', version, fancy=False, dry_run=True)

    Traceback (most recent call last):

    UserError: / exists, but is not a bzr branch.

    >>> import_version(result_path, version, fancy=False, dry_run=True)

    Traceback (most recent call last):

    UserError: The version test@example.com/test--test--0.1 does not exist.

    >>> version = pybaz.Version("test@example.com/test--test--0")

    >>> import_version(result_path, version, fancy=False, dry_run=True)

    not fancy

    ....

    Dry run, not modifying output_dir

    Cleaning up

    >>> import_version(result_path, version, fancy=False)

    not fancy

    ....

    Cleaning up

    Import complete.

    >>> import_version(result_path, version, fancy=False)

    Tree is up-to-date with test@example.com/test--test--0--patch-2

    >>> commit_more_test_revisions()

    >>> import_version(result_path, version, fancy=False)

    not fancy

    ..

    Cleaning up

    Import complete.

    >>> teardown_environ(q)

    """

    try:

        ancestors, old_revno = get_remaining_revisions(output_dir, version)

    except NotInABranch, e:

        raise UserError("%s exists, but is not a bzr branch." % e.path)

    if len(ancestors) == 0:

        last_revision = get_last_revision(find_branch(output_dir))

        print 'Tree is up-to-date with %s' % last_revision

        return

    progress_bar = ProgressBar()

    tempdir = tempfile.mkdtemp(prefix="baz2bzr-",

                               dir=os.path.dirname(output_dir))

    try:

        if not fancy:

            print "not fancy"

        try:

            for result in iter_import_version(output_dir, ancestors, tempdir,

                    fast=fast, verbose=verbose, dry_run=dry_run, 

                    max_count=max_count, skip_symlinks=skip_symlinks):

                if fancy:

                    progress_bar(result)

                else:

                    sys.stdout.write('.')

        finally:

            if fancy:

                clear_progress_bar()

            else:

                sys.stdout.write('\n')

        if dry_run:

            print 'Dry run, not modifying output_dir'

            return

        if os.path.exists(output_dir):

            # Move the bzr control directory back, and update the working tree

            tmp_bzr_dir = os.path.join(tempdir, '.bzr')

            bzr_dir = os.path.join(output_dir, '.bzr')

            new_bzr_dir = os.path.join(tempdir, "rd", '.bzr')

            os.rename(bzr_dir, tmp_bzr_dir) # Move the original bzr out of the way

            os.rename(new_bzr_dir, bzr_dir)

            try:

                bzrlib.merge.merge((output_dir, -1), (output_dir, old_revno), 

                                   check_clean=False, this_dir=output_dir, 

                                   ignore_zero=True)

            except:

                # If something failed, move back the original bzr directory

                os.rename(bzr_dir, new_bzr_dir)

                os.rename(tmp_bzr_dir, bzr_dir)

                raise

        else:

            revdir = os.path.join(tempdir, "rd")

            os.rename(revdir, output_dir)

    finally:

        print 'Cleaning up'

        shutil.rmtree(tempdir)

    print "Import complete."

class UserError(Exception):

    def __init__(self, message):

        """Exception to throw when a user makes an impossible request

        :param message: The message to emit when printing this exception

        :type message: string

        """

        Exception.__init__(self, message)

def revision_id(arch_revision):

    """

    Generate a Bzr revision id from an Arch revision id.  'x' in the id

    designates a revision imported with an experimental algorithm.  A number

    would indicate a particular standardized version.

    :param arch_revision: The Arch revision to generate an ID for.

    >>> revision_id(pybaz.Revision("you@example.com/cat--br--0--base-0"))

    'Arch-x:you@example.com%cat--br--0--base-0'

    """

    return "Arch-x:%s" % str(arch_revision).replace('/', '%')

class NotArchRevision(Exception):

    def __init__(self, revision_id):

        msg = "The revision id %s does not look like it came from Arch."\

            % revision_id

        Exception.__init__(self, msg)

def arch_revision(revision_id):

    """

    >>> str(arch_revision("Arch-x:jrandom@example.com%test--test--0"))

    Traceback (most recent call last):

    NotArchRevision: The revision id Arch-x:jrandom@example.com%test--test--0 does not look like it came from Arch.

    >>> str(arch_revision("Arch-x:jrandom@example.com%test--test--0--base-5"))

    Traceback (most recent call last):

    NotArchRevision: The revision id Arch-x:jrandom@example.com%test--test--0--base-5 does not look like it came from Arch.

    >>> str(arch_revision("Arch-x:jrandom@example.com%test--test--0--patch-5"))

    'jrandom@example.com/test--test--0--patch-5'

    """

    if revision_id is None:

        return None

    if revision_id[:7] != 'Arch-x:':

        raise NotArchRevision(revision_id)

    else:

        try:

            return pybaz.Revision(revision_id[7:].replace('%', '/'))

        except pybaz.errors.NamespaceError, e:

            raise NotArchRevision(revision_id)

def iter_import_version(output_dir, ancestors, tempdir, fast=False,

                        verbose=False, dry_run=False, max_count=None,

                        skip_symlinks=False):

    revdir = None

    # Uncomment this for testing, it basically just has baz2bzr only update

    # 5 patches at a time

    if max_count:

        ancestors = ancestors[:max_count]

    # Not sure if I want this output. basically it tells you ahead of time

    # what it is going to do, but then later it tells you as it is doing it.

    # what probably would be best would be to collapse it into ranges, so that

    # this gives the simple view, and then later it gives the blow by blow.

    #if verbose:

    #    print 'Adding the following revisions:'

    #    for a in ancestors:

    #        print '\t%s' % a

    previous_version=None

    for i in range(len(ancestors)):

        revision = ancestors[i]

        if verbose:

            version = str(revision.version)

            if version != previous_version:

                clear_progress_bar()

                print '\rOn version: %s' % version

            yield Progress(str(revision.patchlevel), i, len(ancestors))

            previous_version = version

        else:

            yield Progress("revisions", i, len(ancestors))

        if revdir is None:

            revdir = os.path.join(tempdir, "rd")

            baz_inv, log = get_revision(revdir, revision, 

                                        skip_symlinks=skip_symlinks)

            if os.path.exists(output_dir):

                bzr_dir = os.path.join(output_dir, '.bzr')

                new_bzr_dir = os.path.join(tempdir, "rd", '.bzr')

                # This would be much faster with a simple os.rename(), but if

                # we fail, we have corrupted the original .bzr directory.  Is

                # that a big problem, as we can just back out the last

                # revisions in .bzr/revision_history I don't really know

                shutil.copytree(bzr_dir, new_bzr_dir)

                # Now revdir should have a tree with the latest .bzr, and the

                # next revision of the baz tree

                branch = find_branch(revdir)

            else:

                branch = bzrlib.Branch(revdir, init=True)

        else:

            old = os.path.join(revdir, ".bzr")

            new = os.path.join(tempdir, ".bzr")

            os.rename(old, new)

            baz_inv, log = apply_revision(revdir, revision, 

                                          skip_symlinks=skip_symlinks)

            os.rename(new, old)

            branch = find_branch(revdir)

        timestamp = email.Utils.mktime_tz(log.date + (0,))

        rev_id = revision_id(revision)

        branch.lock_write()

        try:

            branch.set_inventory(baz_inv)

            bzrlib.trace.silent = True

            branch.commit(log.summary, verbose=False, committer=log.creator,

                          timestamp=timestamp, timezone=0, rev_id=rev_id)

        finally:

            bzrlib.trace.silent = False   

            branch.unlock()

    yield Progress("revisions", len(ancestors), len(ancestors))

    unlink_unversioned(branch, revdir)

def unlink_unversioned(branch, revdir):

    for unversioned in branch.working_tree().extras():

        path = os.path.join(revdir, unversioned)

        if os.path.isdir(path):

            shutil.rmtree(path)

        else:

            os.unlink(path)

def get_log(tree, revision):

    log = tree.iter_logs(version=revision.version, reverse=True).next()

    assert log.revision == revision

    return log

def get_revision(revdir, revision, skip_symlinks=False):

    revision.get(revdir)

    tree = pybaz.tree_root(revdir)

    log = get_log(tree, revision)

    try:

        return bzr_inventory_data(tree, skip_symlinks=skip_symlinks), log 

    except BadFileKind, e:

        raise UserError("Cannot convert %s because %s is a %s" % (revision,e.path, e.kind) )

def apply_revision(revdir, revision, skip_symlinks=False):

    tree = pybaz.tree_root(revdir)

    revision.apply(tree)

    log = get_log(tree, revision)

    try:

        return bzr_inventory_data(tree, skip_symlinks=skip_symlinks), log

    except BadFileKind, e:

        raise UserError("Cannot convert %s because %s is a %s" % (revision,e.path, e.kind) )

class BadFileKind(Exception):

    """The file kind is not permitted in bzr inventories"""

    def __init__(self, tree_root, path, kind):

        self.tree_root = tree_root

        self.path = path

        self.kind = kind

        Exception.__init__(self, "File %s is of forbidden type %s" %

                           (os.path.join(tree_root, path), kind))

def bzr_inventory_data(tree, skip_symlinks=False):

    inv_iter = tree.iter_inventory_ids(source=True, both=True)

    inv_map = {}

    for file_id, path in inv_iter:

        inv_map[path] = file_id 

    bzr_inv = []

    for path, file_id in inv_map.iteritems():

        full_path = os.path.join(tree, path)

        kind = bzrlib.osutils.file_kind(full_path)

        if skip_symlinks and kind == "symlink":

            continue

        if kind not in ("file", "directory"):

            raise BadFileKind(tree, path, kind)

        parent_dir = os.path.dirname(path)

        if parent_dir != "":

            parent_id = inv_map[parent_dir]

        else:

            parent_id = bzrlib.inventory.ROOT_ID

        bzr_inv.append((path, file_id, parent_id, kind))

    bzr_inv.sort()

    return bzr_inv

def main(args):

    """Just the main() function for this script.

    By separating it into a function, this can be called as a child from some other

    script.

    :param args: The arguments to this script. Essentially sys.argv[1:]

    """

    import optparse

    parser = optparse.OptionParser(usage='%prog [options] [VERSION] OUTDIR'

        '\n  VERSION is the arch version to import.'

        '\n  OUTDIR can be an existing directory to be updated'

        '\n         or a new directory which will be created from scratch.')

    parser.add_option('--verbose', action='store_true'

        , help='Get chatty')

    parser.add_option('--skip-symlinks', action="store_true", 

                      dest="skip_symlinks", 

                      help="Ignore any symlinks present in the Arch tree.")

    g = optparse.OptionGroup(parser, 'Test options', 'Options useful while testing process.')

    g.add_option('--test', action='store_true'

        , help='Run the self-tests and exit.')

    g.add_option('--dry-run', action='store_true'

        , help='Do the update, but don\'t copy the result to OUTDIR')

    g.add_option('--max-count', type='int', metavar='COUNT', default=None

        , help='At most, add COUNT patches.')

    g.add_option('--safe', action='store_false', dest='fast')

    g.add_option('--fast', action='store_true', default=False

        , help='By default the .bzr control directory will be copied, so that an error'

        ' does not modify the original. --fast allows the directory to be renamed instead.')

    parser.add_option_group(g)

    (opts, args) = parser.parse_args(args)

    if opts.test:

        print "Running tests"

        import doctest

        nfail, ntests = doctest.testmod(verbose=opts.verbose)

        if nfail > 0:

            return 1

        else:

            return 0

    if len(args) == 2:

        output_dir = os.path.realpath(args[1])

        try:

            version = pybaz.Version(args[0])

        except pybaz.errors.NamespaceError:

            print "%s is not a valid Arch branch." % args[0]

            return 1

    elif len(args) == 1:

        output_dir = os.path.realpath(args[0])

        version = None

    else:

        print 'Invalid number of arguments, try --help for more info'

        return 1

    try:

        import_version(output_dir, version,

            verbose=opts.verbose, fast=opts.fast,

            dry_run=opts.dry_run, max_count=opts.max_count,

            skip_symlinks=opts.skip_symlinks)

        return 0

    except UserError, e:

        print e

        return 1

    except KeyboardInterrupt:

        print "Aborted."

        return 1

class NotInABranch(Exception):

    def __init__(self, path):

        Exception.__init__(self, "%s is not in a branch." % path)

        self.path = path

def find_branch(path):

    """

    >>> find_branch('/')

    Traceback (most recent call last):

    NotInABranch: / is not in a branch.

    >>> sb = bzrlib.ScratchBranch()

    >>> isinstance(find_branch(sb.base), bzrlib.Branch)

    True

    """

    try:

        return bzrlib.Branch(path)

    except BzrError, e:

        if e.args[0].endswith("' is not in a branch"):

            raise NotInABranch(path)

if __name__ == '__main__':

    sys.exit(main(sys.argv[1:]))