~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.branch import Branch
from bzrlib.commands import Command
from errors import NoPyBaz
try:
    import pybaz
    import pybaz.errors
    from pybaz.backends.baz import null_cmd
except ImportError:
    raise NoPyBaz
import tempfile
import os
import os.path
import shutil
import bzrlib
from bzrlib.errors import BzrError, NotBranchError, BzrCommandError
import bzrlib.trace
import bzrlib.merge
import bzrlib.inventory
import bzrlib.osutils
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_revision()
    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 last_patch is None:
            raise NotPreviousImport(branch.base)
        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)
        if len(ancestors) > 0 and not ancestors[0].archive.is_registered():
            ancestors = ancestors[1:]
    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 NotPreviousImport(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:
                    show_progress(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(BzrCommandError):
    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
        """
        BzrCommandError.__init__(self, message)

class NotPreviousImport(UserError):
    def __init__(self, path):
        UserError.__init__(self, "%s is not the location of a previous import."
                           % path)


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 = Branch.initialize(revdir)
        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.working_tree().set_inventory(baz_inv)
            bzrlib.trace.silent = True
            wt = branch.working_tree()
            wt.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 str(log.revision) == str(revision), (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 arch_id, path in inv_iter:
        bzr_file_id = arch_id.replace('%', '%25').replace('/', '%2f')
        inv_map[path] = bzr_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), Branch)
    True
    """
    try:
        return Branch.open(path)
    except NotBranchError, e:
        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(from_branch)
            except pybaz.errors.NamespaceError:
                print "%s is not a valid Arch branch." % from_branch
                return 1
        import_version(to_location, from_branch)