~abentley/bzrtools/bzrtools.dev

# Copyright (C) 2005 by Aaron Bentley

# 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
from bzrlib import Branch
from bzrlib.commands import Command
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:
                progress_bar.clear()
            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

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)

class cmd_baz_import(Command):
    """Import an Arch or Baz branch into a bzr branch"""
    takes_args = ['to_location', 'from_branch?']
    takes_options = ['verbose']

    def run(self, to_location, from_branch=None, skip_symlinks=False, 
            fast=False, max_count=None, verbose=False, dry_run=False):
        to_location = os.path.realpath(str(to_location))
        if from_branch is not None:
            try:
                from_branch = pybaz.Version(version)
            except pybaz.errors.NamespaceError:
                print "%s is not a valid Arch branch." % from_branch
                return 1
        import_version(to_location, from_branch)