← Back to branch summary

~abentley/bzrtools/bzrtools.dev

« back to all changes in this revision

Viewing changes to upstream_import.py

  • Committer: Aaron Bentley
  • Date: 2006-08-07 03:07:35 UTC
  • mto: This revision was merged to the branch mainline in revision 425.
  • Revision ID: aaron.bentley@utoronto.ca-20060807030735-0a9f8330ce1a836a
Add --no-color option to shelf

Show diffs side-by-side

added added

removed removed

Lines of Context:
upstream_import.py
1
1
"""Import upstream source into a branch"""
2
2
 
 
3
from bz2 import BZ2File
3
4
import errno
4
5
import os
5
 
import re
 
6
from shutil import rmtree
6
7
from StringIO import StringIO
7
 
import stat
8
8
import tarfile
9
 
import zipfile
 
9
from unittest import makeSuite
10
10
 
11
 
from bzrlib import generate_ids
12
11
from bzrlib.bzrdir import BzrDir
 
12
from bzrlib.delta import compare_trees
13
13
from bzrlib.errors import NoSuchFile, BzrCommandError, NotBranchError
14
 
from bzrlib.osutils import (pathjoin, isdir, file_iterator, basename,
15
 
                            file_kind, splitpath)
 
14
from bzrlib.osutils import pathjoin, isdir, file_iterator
 
15
from bzrlib.tests import TestCaseInTempDir
16
16
from bzrlib.trace import warning
17
17
from bzrlib.transform import TreeTransform, resolve_conflicts, cook_conflicts
18
18
from bzrlib.workingtree import WorkingTree
19
 
from bzrlib.plugins.bzrtools.bzrtools import open_from_url
20
 
from bzrlib.plugins.bzrtools import errors
21
 
 
22
 
class ZipFileWrapper(object):
23
 
 
24
 
    def __init__(self, fileobj, mode):
25
 
        self.zipfile = zipfile.ZipFile(fileobj, mode)
26
 
 
27
 
    def getmembers(self):
28
 
        for info in self.zipfile.infolist():
29
 
            yield ZipInfoWrapper(self.zipfile, info)
30
 
 
31
 
    def extractfile(self, infowrapper):
32
 
        return StringIO(self.zipfile.read(infowrapper.name))
33
 
 
34
 
    def add(self, filename):
35
 
        if isdir(filename):
36
 
            self.zipfile.writestr(filename+'/', '')
37
 
        else:
38
 
            self.zipfile.write(filename)
39
 
 
40
 
    def close(self):
41
 
        self.zipfile.close()
42
 
 
43
 
 
44
 
class ZipInfoWrapper(object):
45
 
 
46
 
    def __init__(self, zipfile, info):
47
 
        self.info = info
48
 
        self.type = None
49
 
        self.name = info.filename
50
 
        self.zipfile = zipfile
51
 
        self.mode = 0666
52
 
 
53
 
    def isdir(self):
54
 
        # Really? Eeeew!
55
 
        return bool(self.name.endswith('/'))
56
 
 
57
 
    def isreg(self):
58
 
        # Really? Eeeew!
59
 
        return not self.isdir()
60
 
 
61
 
 
62
 
class DirWrapper(object):
63
 
    def __init__(self, fileobj, mode='r'):
64
 
        assert mode == 'r', mode
65
 
        self.root = os.path.realpath(fileobj.read())
66
 
 
67
 
    def __repr__(self):
68
 
        return 'DirWrapper(%r)' % self.root
69
 
 
70
 
    def getmembers(self, subdir=None):
71
 
        if subdir is not None:
72
 
            mydir = pathjoin(self.root, subdir)
73
 
        else:
74
 
            mydir = self.root
75
 
        for child in os.listdir(mydir):
76
 
            if subdir is not None:
77
 
                child = pathjoin(subdir, child)
78
 
            fi = FileInfo(self.root, child)
79
 
            yield fi
80
 
            if fi.isdir():
81
 
                for v in self.getmembers(child):
82
 
                    yield v
83
 
 
84
 
    def extractfile(self, member):
85
 
        return open(member.fullpath)
86
 
 
87
 
 
88
 
class FileInfo(object):
89
 
 
90
 
    def __init__(self, root, filepath):
91
 
        self.fullpath = pathjoin(root, filepath)
92
 
        self.root = root
93
 
        if filepath != '':
94
 
            self.name = pathjoin(basename(root), filepath)
95
 
        else:
96
 
            print 'root %r' % root
97
 
            self.name = basename(root)
98
 
        self.type = None
99
 
        stat = os.lstat(self.fullpath)
100
 
        self.mode = stat.st_mode
101
 
        if self.isdir():
102
 
            self.name += '/'
103
 
 
104
 
    def __repr__(self):
105
 
        return 'FileInfo(%r)' % self.name
106
 
 
107
 
    def isreg(self):
108
 
        return stat.S_ISREG(self.mode)
109
 
 
110
 
    def isdir(self):
111
 
        return stat.S_ISDIR(self.mode)
112
 
 
113
 
    def issym(self):
114
 
        if stat.S_ISLNK(self.mode):
115
 
            self.linkname = os.readlink(self.fullpath)
116
 
            return True
117
 
        else:
118
 
            return False
119
 
 
120
 
 
121
 
def top_path(path):
 
19
 
 
20
 
 
21
def top_directory(path):
122
22
    """Return the top directory given in a path."""
123
 
    components = splitpath(path)
124
 
    if len(components) > 0:
125
 
        return components[0]
126
 
    else:
127
 
        return ''
 
23
    dirname = os.path.dirname(path)
 
24
    last_dirname = dirname
 
25
    while True:
 
26
        dirname = os.path.dirname(dirname)
 
27
        if dirname == '' or dirname == last_dirname:
 
28
            return last_dirname
 
29
        last_dirname = dirname
128
30
 
129
31
 
130
32
def common_directory(names):
131
33
    """Determine a single directory prefix from a list of names"""
132
34
    possible_prefix = None
133
35
    for name in names:
134
 
        name_top = top_path(name)
135
 
        if name_top == '':
136
 
            return None
 
36
        name_top = top_directory(name)
137
37
        if possible_prefix is None:
138
38
            possible_prefix = name_top
139
39
        else:
164
64
            yield member.name
165
65
 
166
66
 
167
 
def should_ignore(relative_path):
168
 
    return top_path(relative_path) == '.bzr'
169
 
 
170
 
 
171
67
def import_tar(tree, tar_input):
172
68
    """Replace the contents of a working directory with tarfile contents.
173
69
    The tarfile may be a gzipped stream.  File ids will be updated.
174
70
    """
175
71
    tar_file = tarfile.open('lala', 'r', tar_input)
176
 
    import_archive(tree, tar_file)
177
 
 
178
 
def import_zip(tree, zip_input):
179
 
    zip_file = ZipFileWrapper(zip_input, 'r')
180
 
    import_archive(tree, zip_file)
181
 
 
182
 
def import_dir(tree, dir_input):
183
 
    dir_file = DirWrapper(dir_input)
184
 
    import_archive(tree, dir_file)
185
 
 
186
 
 
187
 
def import_archive(tree, archive_file):
 
72
    prefix = common_directory(names_of_files(tar_file))
188
73
    tt = TreeTransform(tree)
189
 
    try:
190
 
        import_archive_to_transform(tree, archive_file, tt)
191
 
        tt.apply()
192
 
    finally:
193
 
        tt.finalize()
194
 
 
195
 
 
196
 
def import_archive_to_transform(tree, archive_file, tt):
197
 
    prefix = common_directory(names_of_files(archive_file))
 
74
 
198
75
    removed = set()
199
 
    for path, entry in tree.iter_entries_by_dir():
200
 
        if entry.parent_id is None:
201
 
            continue
 
76
    for path, entry in tree.inventory.iter_entries():
202
77
        trans_id = tt.trans_id_tree_path(path)
203
78
        tt.delete_contents(trans_id)
204
79
        removed.add(path)
205
80
 
206
 
    added = set()
 
81
    added = set() 
207
82
    implied_parents = set()
208
83
    seen = set()
209
 
    for member in archive_file.getmembers():
 
84
    for member in tar_file.getmembers():
210
85
        if member.type == 'g':
211
86
            # type 'g' is a header
212
87
            continue
213
 
        # Inverse functionality in bzr uses utf-8.  We could also
214
 
        # interpret relative to fs encoding, which would match native
215
 
        # behaviour better.
216
 
        relative_path = member.name.decode('utf-8')
 
88
        relative_path = member.name 
217
89
        if prefix is not None:
218
90
            relative_path = relative_path[len(prefix)+1:]
219
 
            relative_path = relative_path.rstrip('/')
220
91
        if relative_path == '':
221
92
            continue
222
 
        if should_ignore(relative_path):
223
 
            continue
224
93
        add_implied_parents(implied_parents, relative_path)
225
94
        trans_id = tt.trans_id_tree_path(relative_path)
226
95
        added.add(relative_path.rstrip('/'))
227
96
        path = tree.abspath(relative_path)
228
97
        if member.name in seen:
229
 
            if tt.final_kind(trans_id) == 'file':
230
 
                tt.set_executability(None, trans_id)
231
98
            tt.cancel_creation(trans_id)
232
99
        seen.add(member.name)
233
100
        if member.isreg():
234
 
            tt.create_file(file_iterator(archive_file.extractfile(member)),
 
101
            tt.create_file(file_iterator(tar_file.extractfile(member)), 
235
102
                           trans_id)
236
 
            executable = (member.mode & 0111) != 0
237
 
            tt.set_executability(executable, trans_id)
238
103
        elif member.isdir():
239
104
            do_directory(tt, trans_id, tree, relative_path, path)
240
105
        elif member.issym():
241
106
            tt.create_symlink(member.linkname, trans_id)
242
 
        else:
243
 
            continue
244
 
        if tt.tree_file_id(trans_id) is None:
245
 
            name = basename(member.name.rstrip('/'))
246
 
            file_id = generate_ids.gen_file_id(name)
247
 
            tt.version_file(file_id, trans_id)
248
107
 
249
108
    for relative_path in implied_parents.difference(added):
250
109
        if relative_path == "":
252
111
        trans_id = tt.trans_id_tree_path(relative_path)
253
112
        path = tree.abspath(relative_path)
254
113
        do_directory(tt, trans_id, tree, relative_path, path)
255
 
        if tt.tree_file_id(trans_id) is None:
256
 
            tt.version_file(trans_id, trans_id)
257
114
        added.add(relative_path)
258
115
 
259
 
    for path in removed.difference(added):
260
 
        tt.unversion_file(tt.trans_id_tree_path(path))
261
 
 
262
116
    for conflict in cook_conflicts(resolve_conflicts(tt), tt):
263
117
        warning(conflict)
 
118
    tt.apply()
 
119
    update_ids(tree, added, removed)
 
120
 
 
121
 
 
122
def update_ids(tree, added, removed):
 
123
    """Make sure that all present files files have file_ids.
 
124
    """
 
125
    # XXX detect renames
 
126
    new = added.difference(removed)
 
127
    deleted = removed.difference(added)
 
128
    tree.add(sorted(new))
 
129
    tree.remove(sorted(deleted, reverse=True))
264
130
 
265
131
 
266
132
def do_import(source, tree_directory=None):
277
143
        tree = WorkingTree.open_containing('.')[0]
278
144
    tree.lock_write()
279
145
    try:
280
 
        if tree.changes_from(tree.basis_tree()).has_changed():
 
146
        if compare_trees(tree, tree.basis_tree()).has_changed():
281
147
            raise BzrCommandError("Working tree has uncommitted changes.")
282
148
 
283
 
        try:
284
 
            archive, external_compressor = get_archive_type(source)
285
 
        except errors.NotArchiveType:
286
 
            if file_kind(source) == 'directory':
287
 
                s = StringIO(source)
288
 
                s.seek(0)
289
 
                import_dir(tree, s)
290
 
            else:
291
 
                raise BzrCommandError('Unhandled import source')
292
 
        else:
293
 
            if archive == 'zip':
294
 
                import_zip(tree, open_from_url(source))
295
 
            elif archive == 'tar':
296
 
                try:
297
 
                    tar_input = open_from_url(source)
298
 
                    if external_compressor == 'bz2':
299
 
                        import bz2
300
 
                        tar_input = StringIO(bz2.decompress(tar_input.read()))
301
 
                    elif external_compressor == 'lzma':
302
 
                        import lzma
303
 
                        tar_input = StringIO(lzma.decompress(tar_input.read()))
304
 
                except IOError, e:
305
 
                    if e.errno == errno.ENOENT:
306
 
                        raise NoSuchFile(source)
307
 
                try:
308
 
                    import_tar(tree, tar_input)
309
 
                finally:
310
 
                    tar_input.close()
 
149
        if (source.endswith('.tar') or source.endswith('.tar.gz') or 
 
150
            source.endswith('.tar.bz2')) or source.endswith('.tgz'):
 
151
            try:
 
152
                if source.endswith('.bz2'):
 
153
                    tar_input = BZ2File(source, 'r')
 
154
                    tar_input = StringIO(tar_input.read())
 
155
                else:
 
156
                    tar_input = file(source, 'rb')
 
157
            except IOError, e:
 
158
                if e.errno == errno.ENOENT:
 
159
                    raise NoSuchFile(source)
 
160
            try:
 
161
                import_tar(tree, tar_input)
 
162
            finally:
 
163
                tar_input.close()
311
164
    finally:
312
165
        tree.unlock()
313
166
 
314
 
 
315
 
def get_archive_type(path):
316
 
    """Return the type of archive and compressor indicated by path name.
317
 
 
318
 
    Only external compressors are returned, so zip files are only
319
 
    ('zip', None).  .tgz is treated as ('tar', 'gz') and '.tar.xz' is treated
320
 
    as ('tar', 'lzma').
321
 
    """
322
 
    matches = re.match(r'.*\.(zip|tgz|tar(.(gz|bz2|lzma|xz))?)$', path)
323
 
    if not matches:
324
 
        raise errors.NotArchiveType(path)
325
 
    external_compressor = None
326
 
    if matches.group(3) is not None:
327
 
        archive = 'tar'
328
 
        external_compressor = matches.group(3)
329
 
        if external_compressor == 'xz':
330
 
            external_compressor = 'lzma'
331
 
    elif matches.group(1) == 'tgz':
332
 
        return 'tar', 'gz'
333
 
    else:
334
 
        archive = matches.group(1)
335
 
    return archive, external_compressor
 
167
class TestImport(TestCaseInTempDir):
 
168
 
 
169
    def make_tar(self, mode='w'):
 
170
        result = StringIO()
 
171
        tar_file = tarfile.open('project-0.1.tar', mode, result)
 
172
        os.mkdir('project-0.1')
 
173
        tar_file.add('project-0.1')
 
174
        os.mkdir('project-0.1/junk')
 
175
        tar_file.add('project-0.1/junk')
 
176
        
 
177
        f = file('project-0.1/README', 'wb')
 
178
        f.write('What?')
 
179
        f.close()
 
180
        tar_file.add('project-0.1/README')
 
181
 
 
182
        f = file('project-0.1/FEEDME', 'wb')
 
183
        f.write('Hungry!!')
 
184
        f.close()
 
185
        tar_file.add('project-0.1/FEEDME')
 
186
 
 
187
        tar_file.close()
 
188
        rmtree('project-0.1')
 
189
        result.seek(0)
 
190
        return result
 
191
 
 
192
    def make_tar2(self):
 
193
        result = StringIO()
 
194
        tar_file = tarfile.open('project-0.2.tar', 'w', result)
 
195
        os.mkdir('project-0.2')
 
196
        tar_file.add('project-0.2')
 
197
        
 
198
        os.mkdir('project-0.2/junk')
 
199
        tar_file.add('project-0.2/junk')
 
200
 
 
201
        f = file('project-0.2/README', 'wb')
 
202
        f.write('Now?')
 
203
        f.close()
 
204
        tar_file.add('project-0.2/README')
 
205
        tar_file.close()
 
206
 
 
207
        tar_file = tarfile.open('project-0.2.tar', 'a', result)
 
208
        tar_file.add('project-0.2/README')
 
209
 
 
210
        rmtree('project-0.2')
 
211
        return result
 
212
 
 
213
    def make_messed_tar(self):
 
214
        result = StringIO()
 
215
        tar_file = tarfile.open('project-0.1.tar', 'w', result)
 
216
        os.mkdir('project-0.1')
 
217
        tar_file.add('project-0.1')
 
218
 
 
219
        os.mkdir('project-0.2')
 
220
        tar_file.add('project-0.2')
 
221
        
 
222
        f = file('project-0.1/README', 'wb')
 
223
        f.write('What?')
 
224
        f.close()
 
225
        tar_file.add('project-0.1/README')
 
226
        tar_file.close()
 
227
        rmtree('project-0.1')
 
228
        result.seek(0)
 
229
        return result
 
230
 
 
231
    def test_top_directory(self):
 
232
        self.assertEqual(top_directory('ab/b/c'), 'ab')
 
233
        self.assertEqual(top_directory('/etc'), '/')
 
234
 
 
235
    def test_common_directory(self):
 
236
        self.assertEqual(common_directory(['ab/c/d', 'ab/c/e']), 'ab')
 
237
        self.assertIs(common_directory(['ab/c/d', 'ac/c/e']), None)
 
238
 
 
239
    def test_untar(self):
 
240
        tar_file = self.make_tar()
 
241
        tree = BzrDir.create_standalone_workingtree('tree')
 
242
        import_tar(tree, tar_file)
 
243
        self.assertTrue(tree.path2id('README') is not None) 
 
244
        self.assertTrue(tree.path2id('FEEDME') is not None)
 
245
        self.assertTrue(os.path.isfile(tree.abspath('README')))
 
246
        self.assertEqual(tree.inventory[tree.path2id('README')].kind, 'file')
 
247
        self.assertEqual(tree.inventory[tree.path2id('FEEDME')].kind, 'file')
 
248
        
 
249
        f = file(tree.abspath('junk/food'), 'wb')
 
250
        f.write('I like food\n')
 
251
        f.close()
 
252
 
 
253
        tar_file = self.make_tar2()
 
254
        import_tar(tree, tar_file)
 
255
        self.assertTrue(tree.path2id('README') is not None) 
 
256
        self.assertTrue(not os.path.exists(tree.abspath('FEEDME')))
 
257
 
 
258
 
 
259
    def test_untar2(self):
 
260
        tar_file = self.make_messed_tar()
 
261
        tree = BzrDir.create_standalone_workingtree('tree')
 
262
        import_tar(tree, tar_file)
 
263
        self.assertTrue(tree.path2id('project-0.1/README') is not None) 
 
264
 
 
265
    def test_untar_gzip(self):
 
266
        tar_file = self.make_tar(mode='w:gz')
 
267
        tree = BzrDir.create_standalone_workingtree('tree')
 
268
        import_tar(tree, tar_file)
 
269
        self.assertTrue(tree.path2id('README') is not None) 
 
270
 
 
271
 
 
272
def test_suite():
 
273
    return makeSuite(TestImport)