~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

import zipfile

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

class ZipFileWrapper(object):

    def __init__(self, zipfile):

        self.zipfile = zipfile

    def getmembers(self):

        for info in self.zipfile.infolist():

            yield ZipInfoWrapper(self.zipfile, info)

    def extractfile(self, infowrapper):

        return StringIO(self.zipfile.read(infowrapper.name))

    def add(self, filename):

        if isdir(filename):

            self.zipfile.writestr(filename+'/', '')

        else:

            self.zipfile.write(filename)

    def close(self):

        self.zipfile.close()

class ZipInfoWrapper(object):

    def __init__(self, zipfile, info):

        self.info = info

        self.type = None

        self.name = info.filename

        self.zipfile = zipfile

    def isdir(self):

        # Really? Eeeew!

        return bool(self.name.endswith('/'))

    def isreg(self):

        # Really? Eeeew!

        return not self.isdir()

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)

    import_archive(tree, tar_file)

def import_zip(tree, zip_input):

    zip_file = ZipFileWrapper(zipfile.ZipFile(zip_input))

    import_archive(tree, zip_file)

def import_archive(tree, archive_file):

    prefix = common_directory(names_of_files(archive_file))

    tt = TreeTransform(tree)

    removed = set()

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

        if entry.parent_id is None:

            continue

        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 archive_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(archive_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()

        elif source.endswith('.zip'):

            import_zip(tree, open(source, 'rb'))

        else:

            raise BzrCommandError('Unhandled import source')

    finally:

        tree.unlock()

class TestImport(TestCaseInTempDir):

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

        def maker(fileobj):

            return tarfile.open('project-0.1.tar', mode, fileobj)

        return self.make_archive(maker)

    def make_archive(self, maker):

        result = StringIO()

        archive_file = maker(result)

        os.mkdir('project-0.1')

        archive_file.add('project-0.1')

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

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

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

        f.write('What?')

        f.close()

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

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

        f.write('Hungry!!')

        f.close()

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

        archive_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 make_zip(self):

        def maker(fileobj):

            return ZipFileWrapper(zipfile.ZipFile(fileobj, 'w'))

        return self.make_archive(maker)

    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_unzip(self):

        zip_file = self.make_zip()

        tree = BzrDir.create_standalone_workingtree('tree')

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

def test_suite():

    return makeSuite(TestImport)