~abentley/bzrtools/bzrtools.dev

"""Import upstream source into a branch"""

from bz2 import BZ2File

import errno

import os

from shutil import rmtree

from StringIO import StringIO

import tarfile

from unittest import makeSuite

from bzrlib.bzrdir import BzrDir

from bzrlib.errors import NoSuchFile, BzrCommandError, NotBranchError

from bzrlib.osutils import pathjoin, isdir, file_iterator

from bzrlib.tests import TestCaseInTempDir

from bzrlib.trace import warning

from bzrlib.transform import TreeTransform, resolve_conflicts, cook_conflicts

from bzrlib.workingtree import WorkingTree

def top_directory(path):

    """Return the top directory given in a path."""

    dirname = os.path.dirname(path)

    last_dirname = dirname

    while True:

        dirname = os.path.dirname(dirname)

        if dirname == '' or dirname == last_dirname:

            return last_dirname

        last_dirname = dirname

def common_directory(names):

    """Determine a single directory prefix from a list of names"""

    possible_prefix = None

    for name in names:

        name_top = top_directory(name)

        if possible_prefix is None:

            possible_prefix = name_top

        else:

            if name_top != possible_prefix:

                return None

    return possible_prefix

def do_directory(tt, trans_id, tree, relative_path, path):

    if isdir(path) and tree.path2id(relative_path) is not None:

        tt.cancel_deletion(trans_id)

    else:

        tt.create_directory(trans_id)

def add_implied_parents(implied_parents, path):

    """Update the set of implied parents from a path"""

    parent = os.path.dirname(path)

    if parent in implied_parents:

        return

    implied_parents.add(parent)

    add_implied_parents(implied_parents, parent)

def names_of_files(tar_file):

    for member in tar_file.getmembers():

        if member.type != "g":

            yield member.name

def import_tar(tree, tar_input):

    """Replace the contents of a working directory with tarfile contents.

    The tarfile may be a gzipped stream.  File ids will be updated.

    """

    tar_file = tarfile.open('lala', 'r', tar_input)

    prefix = common_directory(names_of_files(tar_file))

    tt = TreeTransform(tree)

    removed = set()

    for path, entry in tree.inventory.iter_entries():

        trans_id = tt.trans_id_tree_path(path)

        tt.delete_contents(trans_id)

        removed.add(path)

    added = set() 

    implied_parents = set()

    seen = set()

    for member in tar_file.getmembers():

        if member.type == 'g':

            # type 'g' is a header

            continue

        relative_path = member.name 

        if prefix is not None:

            relative_path = relative_path[len(prefix)+1:]

        if relative_path == '':

            continue

        add_implied_parents(implied_parents, relative_path)

        trans_id = tt.trans_id_tree_path(relative_path)

        added.add(relative_path.rstrip('/'))

        path = tree.abspath(relative_path)

        if member.name in seen:

            tt.cancel_creation(trans_id)

        seen.add(member.name)

        if member.isreg():

            tt.create_file(file_iterator(tar_file.extractfile(member)), 

                           trans_id)

        elif member.isdir():

            do_directory(tt, trans_id, tree, relative_path, path)

        elif member.issym():

            tt.create_symlink(member.linkname, trans_id)

    for relative_path in implied_parents.difference(added):

        if relative_path == "":

            continue

        trans_id = tt.trans_id_tree_path(relative_path)

        path = tree.abspath(relative_path)

        do_directory(tt, trans_id, tree, relative_path, path)

        added.add(relative_path)

    for conflict in cook_conflicts(resolve_conflicts(tt), tt):

        warning(conflict)

    tt.apply()

    update_ids(tree, added, removed)

def update_ids(tree, added, removed):

    """Make sure that all present files files have file_ids.

    """

    # XXX detect renames

    new = added.difference(removed)

    deleted = removed.difference(added)

    tree.add(sorted(new))

    tree.remove(sorted(deleted, reverse=True))

def do_import(source, tree_directory=None):

    """Implementation of import command.  Intended for UI only"""

    if tree_directory is not None:

        try:

            tree = WorkingTree.open(tree_directory)

        except NotBranchError:

            if not os.path.exists(tree_directory):

                os.mkdir(tree_directory)

            branch = BzrDir.create_branch_convenience(tree_directory)

            tree = branch.bzrdir.open_workingtree()

    else:

        tree = WorkingTree.open_containing('.')[0]

    tree.lock_write()

    try:

        if tree.changes_from(tree.basis_tree()).has_changed():

            raise BzrCommandError("Working tree has uncommitted changes.")

        if (source.endswith('.tar') or source.endswith('.tar.gz') or 

            source.endswith('.tar.bz2')) or source.endswith('.tgz'):

            try:

                if source.endswith('.bz2'):

                    tar_input = BZ2File(source, 'r')

                    tar_input = StringIO(tar_input.read())

                else:

                    tar_input = file(source, 'rb')

            except IOError, e:

                if e.errno == errno.ENOENT:

                    raise NoSuchFile(source)

            try:

                import_tar(tree, tar_input)

            finally:

                tar_input.close()

    finally:

        tree.unlock()

class TestImport(TestCaseInTempDir):

    def make_tar(self, mode='w'):

        result = StringIO()

        tar_file = tarfile.open('project-0.1.tar', mode, result)

        os.mkdir('project-0.1')

        tar_file.add('project-0.1')

        os.mkdir('project-0.1/junk')

        tar_file.add('project-0.1/junk')

        f = file('project-0.1/README', 'wb')

        f.write('What?')

        f.close()

        tar_file.add('project-0.1/README')

        f = file('project-0.1/FEEDME', 'wb')

        f.write('Hungry!!')

        f.close()

        tar_file.add('project-0.1/FEEDME')

        tar_file.close()

        rmtree('project-0.1')

        result.seek(0)

        return result

    def make_tar2(self):

        result = StringIO()

        tar_file = tarfile.open('project-0.2.tar', 'w', result)

        os.mkdir('project-0.2')

        tar_file.add('project-0.2')

        os.mkdir('project-0.2/junk')

        tar_file.add('project-0.2/junk')

        f = file('project-0.2/README', 'wb')

        f.write('Now?')

        f.close()

        tar_file.add('project-0.2/README')

        tar_file.close()

        tar_file = tarfile.open('project-0.2.tar', 'a', result)

        tar_file.add('project-0.2/README')

        rmtree('project-0.2')

        return result

    def make_messed_tar(self):

        result = StringIO()

        tar_file = tarfile.open('project-0.1.tar', 'w', result)

        os.mkdir('project-0.1')

        tar_file.add('project-0.1')

        os.mkdir('project-0.2')

        tar_file.add('project-0.2')

        f = file('project-0.1/README', 'wb')

        f.write('What?')

        f.close()

        tar_file.add('project-0.1/README')

        tar_file.close()

        rmtree('project-0.1')

        result.seek(0)

        return result

    def test_top_directory(self):

        self.assertEqual(top_directory('ab/b/c'), 'ab')

        self.assertEqual(top_directory('/etc'), '/')

    def test_common_directory(self):

        self.assertEqual(common_directory(['ab/c/d', 'ab/c/e']), 'ab')

        self.assertIs(common_directory(['ab/c/d', 'ac/c/e']), None)

    def test_untar(self):

        tar_file = self.make_tar()

        tree = BzrDir.create_standalone_workingtree('tree')

        import_tar(tree, tar_file)

        self.assertTrue(tree.path2id('README') is not None) 

        self.assertTrue(tree.path2id('FEEDME') is not None)

        self.assertTrue(os.path.isfile(tree.abspath('README')))

        self.assertEqual(tree.inventory[tree.path2id('README')].kind, 'file')

        self.assertEqual(tree.inventory[tree.path2id('FEEDME')].kind, 'file')

        f = file(tree.abspath('junk/food'), 'wb')

        f.write('I like food\n')

        f.close()

        tar_file = self.make_tar2()

        import_tar(tree, tar_file)

        self.assertTrue(tree.path2id('README') is not None) 

        self.assertTrue(not os.path.exists(tree.abspath('FEEDME')))

    def test_untar2(self):

        tar_file = self.make_messed_tar()

        tree = BzrDir.create_standalone_workingtree('tree')

        import_tar(tree, tar_file)

        self.assertTrue(tree.path2id('project-0.1/README') is not None) 

    def test_untar_gzip(self):

        tar_file = self.make_tar(mode='w:gz')

        tree = BzrDir.create_standalone_workingtree('tree')

        import_tar(tree, tar_file)

        self.assertTrue(tree.path2id('README') is not None) 

def test_suite():

    return makeSuite(TestImport)