← 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-12-12 22:41:55 UTC
  • Revision ID: abentley@panoramicfeedback.com-20061212224155-uwyt16xhvekwwh3m
Add zip import support

Show diffs side-by-side

added added

removed removed

Lines of Context:
upstream_import.py
7
7
from StringIO import StringIO
8
8
import tarfile
9
9
from unittest import makeSuite
 
10
import zipfile
10
11
 
11
12
from bzrlib.bzrdir import BzrDir
12
13
from bzrlib.errors import NoSuchFile, BzrCommandError, NotBranchError
16
17
from bzrlib.transform import TreeTransform, resolve_conflicts, cook_conflicts
17
18
from bzrlib.workingtree import WorkingTree
18
19
 
 
20
class ZipFileWrapper(object):
 
21
 
 
22
    def __init__(self, zipfile):
 
23
        self.zipfile = zipfile
 
24
 
 
25
    def getmembers(self):
 
26
        for info in self.zipfile.infolist():
 
27
            yield ZipInfoWrapper(self.zipfile, info)
 
28
 
 
29
    def extractfile(self, infowrapper):
 
30
        return StringIO(self.zipfile.read(infowrapper.name))
 
31
 
 
32
 
 
33
class ZipInfoWrapper(object):
 
34
    
 
35
    def __init__(self, zipfile, info):
 
36
        self.info = info
 
37
        self.type = None
 
38
        self.name = info.filename
 
39
        self.zipfile = zipfile
 
40
 
 
41
    def isdir(self):
 
42
        # Really? Eeeew!
 
43
        return bool(self.name.endswith('/'))
 
44
 
 
45
    def isreg(self):
 
46
        # Really? Eeeew!
 
47
        return not self.isdir()
 
48
 
19
49
 
20
50
def top_directory(path):
21
51
    """Return the top directory given in a path."""
68
98
    The tarfile may be a gzipped stream.  File ids will be updated.
69
99
    """
70
100
    tar_file = tarfile.open('lala', 'r', tar_input)
71
 
    prefix = common_directory(names_of_files(tar_file))
 
101
    import_archive(tree, tar_file)
 
102
 
 
103
def import_zip(tree, zip_input):
 
104
    zip_file = ZipFileWrapper(zipfile.ZipFile(zip_input))
 
105
    import_archive(tree, zip_file)
 
106
 
 
107
def import_archive(tree, archive_file):
 
108
    prefix = common_directory(names_of_files(archive_file))
72
109
    tt = TreeTransform(tree)
73
110
 
74
111
    removed = set()
82
119
    added = set() 
83
120
    implied_parents = set()
84
121
    seen = set()
85
 
    for member in tar_file.getmembers():
 
122
    for member in archive_file.getmembers():
86
123
        if member.type == 'g':
87
124
            # type 'g' is a header
88
125
            continue
99
136
            tt.cancel_creation(trans_id)
100
137
        seen.add(member.name)
101
138
        if member.isreg():
102
 
            tt.create_file(file_iterator(tar_file.extractfile(member)), 
 
139
            tt.create_file(file_iterator(archive_file.extractfile(member)), 
103
140
                           trans_id)
104
141
        elif member.isdir():
105
142
            do_directory(tt, trans_id, tree, relative_path, path)
162
199
                import_tar(tree, tar_input)
163
200
            finally:
164
201
                tar_input.close()
 
202
        elif source.endswith('.zip'):
 
203
            import_zip(tree, open(source, 'rb'))
 
204
        else:
 
205
            raise BzrCommandError('Unhandled import source')
165
206
    finally:
166
207
        tree.unlock()
167
208
 
229
270
        result.seek(0)
230
271
        return result
231
272
 
 
273
    def make_zip(self, mode='w'):
 
274
        result = StringIO()
 
275
        zip_file = zipfile.ZipFile(result, mode)
 
276
        os.mkdir('project-0.1')
 
277
        zip_file.writestr('project-0.1/', '')
 
278
        os.mkdir('project-0.1/junk')
 
279
        zip_file.writestr('project-0.1/junk/', '')
 
280
        
 
281
        f = file('project-0.1/README', 'wb')
 
282
        f.write('What?')
 
283
        f.close()
 
284
        zip_file.write('project-0.1/README')
 
285
 
 
286
        f = file('project-0.1/FEEDME', 'wb')
 
287
        f.write('Hungry!!')
 
288
        f.close()
 
289
        zip_file.write('project-0.1/FEEDME')
 
290
 
 
291
        zip_file.close()
 
292
        rmtree('project-0.1')
 
293
        result.seek(0)
 
294
        return result
 
295
 
232
296
    def test_top_directory(self):
233
297
        self.assertEqual(top_directory('ab/b/c'), 'ab')
234
298
        self.assertEqual(top_directory('/etc'), '/')
269
333
        import_tar(tree, tar_file)
270
334
        self.assertTrue(tree.path2id('README') is not None) 
271
335
 
 
336
    def test_unzip(self):
 
337
        zip_file = self.make_zip()
 
338
        tree = BzrDir.create_standalone_workingtree('tree')
 
339
        import_zip(tree, zip_file)
 
340
        self.assertTrue(tree.path2id('README') is not None) 
 
341
        self.assertTrue(tree.path2id('FEEDME') is not None)
 
342
        self.assertTrue(os.path.isfile(tree.abspath('README')))
 
343
        self.assertEqual(tree.inventory[tree.path2id('README')].kind, 'file')
 
344
        self.assertEqual(tree.inventory[tree.path2id('FEEDME')].kind, 'file')
 
345
        
 
346
        f = file(tree.abspath('junk/food'), 'wb')
 
347
        f.write('I like food\n')
 
348
        f.close()
 
349
 
272
350
 
273
351
def test_suite():
274
352
    return makeSuite(TestImport)