~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.delta import compare_trees

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 compare_trees(tree, 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)