~abentley/bzrtools/bzrtools.dev

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
"""Import upstream source into a branch"""

from bz2 import BZ2File
import errno
import os
from StringIO import StringIO
import tarfile
import zipfile

from bzrlib.bzrdir import BzrDir
from bzrlib.errors import NoSuchFile, BzrCommandError, NotBranchError
from bzrlib.osutils import pathjoin, isdir, file_iterator
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, fileobj, mode):
        self.zipfile = zipfile.ZipFile(fileobj, mode)

    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(zip_input, 'r')
    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()