← Back to branch summary

~bzr-pqm/bzr/bzr.dev

« back to all changes in this revision

Viewing changes to bzrlib/pack.py

(vila) Make all transport put_bytes() raises TypeError when given unicode
 strings rather than bytes (Vincent Ladeuil)

Show diffs side-by-side

added added

removed removed

Lines of Context:
bzrlib/pack.py
1
 
# Copyright (C) 2007 Canonical Ltd
 
1
# Copyright (C) 2007, 2009, 2010 Canonical Ltd
2
2
#
3
3
# This program is free software; you can redistribute it and/or modify
4
4
# it under the terms of the GNU General Public License as published by
12
12
#
13
13
# You should have received a copy of the GNU General Public License
14
14
# along with this program; if not, write to the Free Software
15
 
# Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
 
15
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
16
16
 
17
17
"""Container format for Bazaar data.
18
18
 
19
 
"Containers" and "records" are described in doc/developers/container-format.txt.
 
19
"Containers" and "records" are described in
 
20
doc/developers/container-format.txt.
20
21
"""
21
22
 
 
23
from __future__ import absolute_import
 
24
 
22
25
from cStringIO import StringIO
23
26
import re
24
27
 
33
36
 
34
37
def _check_name(name):
35
38
    """Do some basic checking of 'name'.
36
 
    
 
39
 
37
40
    At the moment, this just checks that there are no whitespace characters in a
38
41
    name.
39
42
 
46
49
 
47
50
def _check_name_encoding(name):
48
51
    """Check that 'name' is valid UTF-8.
49
 
    
 
52
 
50
53
    This is separate from _check_name because UTF-8 decoding is relatively
51
54
    expensive, and we usually want to avoid it.
52
55
 
58
61
        raise errors.InvalidRecordError(str(e))
59
62
 
60
63
 
 
64
class ContainerSerialiser(object):
 
65
    """A helper class for serialising containers.
 
66
 
 
67
    It simply returns bytes from method calls to 'begin', 'end' and
 
68
    'bytes_record'.  You may find ContainerWriter to be a more convenient
 
69
    interface.
 
70
    """
 
71
 
 
72
    def begin(self):
 
73
        """Return the bytes to begin a container."""
 
74
        return FORMAT_ONE + "\n"
 
75
 
 
76
    def end(self):
 
77
        """Return the bytes to finish a container."""
 
78
        return "E"
 
79
 
 
80
    def bytes_header(self, length, names):
 
81
        """Return the header for a Bytes record."""
 
82
        # Kind marker
 
83
        byte_sections = ["B"]
 
84
        # Length
 
85
        byte_sections.append(str(length) + "\n")
 
86
        # Names
 
87
        for name_tuple in names:
 
88
            # Make sure we're writing valid names.  Note that we will leave a
 
89
            # half-written record if a name is bad!
 
90
            for name in name_tuple:
 
91
                _check_name(name)
 
92
            byte_sections.append('\x00'.join(name_tuple) + "\n")
 
93
        # End of headers
 
94
        byte_sections.append("\n")
 
95
        return ''.join(byte_sections)
 
96
 
 
97
    def bytes_record(self, bytes, names):
 
98
        """Return the bytes for a Bytes record with the given name and
 
99
        contents.
 
100
 
 
101
        If the content may be large, construct the header separately and then
 
102
        stream out the contents.
 
103
        """
 
104
        return self.bytes_header(len(bytes), names) + bytes
 
105
 
 
106
 
61
107
class ContainerWriter(object):
62
 
    """A class for writing containers.
 
108
    """A class for writing containers to a file.
63
109
 
64
110
    :attribute records_written: The number of user records added to the
65
111
        container. This does not count the prelude or suffix of the container
66
112
        introduced by the begin() and end() methods.
67
113
    """
68
114
 
 
115
    # Join up headers with the body if writing fewer than this many bytes:
 
116
    # trades off memory usage and copying to do less IO ops.
 
117
    _JOIN_WRITES_THRESHOLD = 100000
 
118
 
69
119
    def __init__(self, write_func):
70
120
        """Constructor.
71
121
 
75
125
        self._write_func = write_func
76
126
        self.current_offset = 0
77
127
        self.records_written = 0
 
128
        self._serialiser = ContainerSerialiser()
78
129
 
79
130
    def begin(self):
80
131
        """Begin writing a container."""
81
 
        self.write_func(FORMAT_ONE + "\n")
 
132
        self.write_func(self._serialiser.begin())
82
133
 
83
134
    def write_func(self, bytes):
84
135
        self._write_func(bytes)
86
137
 
87
138
    def end(self):
88
139
        """Finish writing a container."""
89
 
        self.write_func("E")
 
140
        self.write_func(self._serialiser.end())
90
141
 
91
142
    def add_bytes_record(self, bytes, names):
92
143
        """Add a Bytes record with the given names.
93
 
        
 
144
 
94
145
        :param bytes: The bytes to insert.
95
146
        :param names: The names to give the inserted bytes. Each name is
96
147
            a tuple of bytestrings. The bytestrings may not contain
103
154
            and thus are only suitable for use by a ContainerReader.
104
155
        """
105
156
        current_offset = self.current_offset
106
 
        # Kind marker
107
 
        self.write_func("B")
108
 
        # Length
109
 
        self.write_func(str(len(bytes)) + "\n")
110
 
        # Names
111
 
        for name_tuple in names:
112
 
            # Make sure we're writing valid names.  Note that we will leave a
113
 
            # half-written record if a name is bad!
114
 
            for name in name_tuple:
115
 
                _check_name(name)
116
 
            self.write_func('\x00'.join(name_tuple) + "\n")
117
 
        # End of headers
118
 
        self.write_func("\n")
119
 
        # Finally, the contents.
120
 
        self.write_func(bytes)
 
157
        length = len(bytes)
 
158
        if length < self._JOIN_WRITES_THRESHOLD:
 
159
            self.write_func(self._serialiser.bytes_header(length, names)
 
160
                + bytes)
 
161
        else:
 
162
            self.write_func(self._serialiser.bytes_header(length, names))
 
163
            self.write_func(bytes)
121
164
        self.records_written += 1
122
165
        # return a memo of where we wrote data to allow random access.
123
166
        return current_offset, self.current_offset - current_offset
124
167
 
125
168
 
126
169
class ReadVFile(object):
127
 
    """Adapt a readv result iterator to a file like protocol."""
 
170
    """Adapt a readv result iterator to a file like protocol.
 
171
    
 
172
    The readv result must support the iterator protocol returning (offset,
 
173
    data_bytes) pairs.
 
174
    """
 
175
 
 
176
    # XXX: This could be a generic transport class, as other code may want to
 
177
    # gradually consume the readv result.
128
178
 
129
179
    def __init__(self, readv_result):
 
180
        """Construct a new ReadVFile wrapper.
 
181
 
 
182
        :seealso: make_readv_reader
 
183
 
 
184
        :param readv_result: the most recent readv result - list or generator
 
185
        """
 
186
        # readv can return a sequence or an iterator, but we require an
 
187
        # iterator to know how much has been consumed.
 
188
        readv_result = iter(readv_result)
130
189
        self.readv_result = readv_result
131
 
        # the most recent readv result block
132
190
        self._string = None
133
191
 
134
192
    def _next(self):
135
193
        if (self._string is None or
136
194
            self._string.tell() == self._string_length):
137
 
            length, data = self.readv_result.next()
 
195
            offset, data = self.readv_result.next()
138
196
            self._string_length = len(data)
139
197
            self._string = StringIO(data)
140
198
 
142
200
        self._next()
143
201
        result = self._string.read(length)
144
202
        if len(result) < length:
145
 
            raise errors.BzrError('request for too much data from a readv hunk.')
 
203
            raise errors.BzrError('wanted %d bytes but next '
 
204
                'hunk only contains %d: %r...' %
 
205
                (length, len(result), result[:20]))
146
206
        return result
147
207
 
148
208
    def readline(self):
150
210
        self._next()
151
211
        result = self._string.readline()
152
212
        if self._string.tell() == self._string_length and result[-1] != '\n':
153
 
            raise errors.BzrError('short readline in the readvfile hunk.')
 
213
            raise errors.BzrError('short readline in the readvfile hunk: %r'
 
214
                % (result, ))
154
215
        return result
155
216
 
156
217
 
199
260
        is a ``list`` and bytes is a function that takes one argument,
200
261
        ``max_length``.
201
262
 
202
 
        You **must not** call the callable after advancing the interator to the
 
263
        You **must not** call the callable after advancing the iterator to the
203
264
        next record.  That is, this code is invalid::
204
265
 
205
266
            record_iter = container.iter_records()
206
267
            names1, callable1 = record_iter.next()
207
268
            names2, callable2 = record_iter.next()
208
269
            bytes1 = callable1(None)
209
 
        
 
270
 
210
271
        As it will give incorrect results and invalidate the state of the
211
272
        ContainerReader.
212
273
 
213
 
        :raises ContainerError: if any sort of containter corruption is
 
274
        :raises ContainerError: if any sort of container corruption is
214
275
            detected, e.g. UnknownContainerFormatError is the format of the
215
276
            container is unrecognised.
216
277
        :seealso: ContainerReader.read
217
278
        """
218
279
        self._read_format()
219
280
        return self._iter_records()
220
 
    
 
281
 
221
282
    def iter_record_objects(self):
222
283
        """Iterate over the container, yielding each record as it is read.
223
284
 
225
286
        methods.  Like with iter_records, it is not safe to use a record object
226
287
        after advancing the iterator to yield next record.
227
288
 
228
 
        :raises ContainerError: if any sort of containter corruption is
 
289
        :raises ContainerError: if any sort of container corruption is
229
290
            detected, e.g. UnknownContainerFormatError is the format of the
230
291
            container is unrecognised.
231
292
        :seealso: iter_records
232
293
        """
233
294
        self._read_format()
234
295
        return self._iter_record_objects()
235
 
    
 
296
 
236
297
    def _iter_records(self):
237
298
        for record in self._iter_record_objects():
238
299
            yield record.read()
280
341
                # risk that the same unicode string has been encoded two
281
342
                # different ways.
282
343
                if name_tuple in all_names:
283
 
                    raise errors.DuplicateRecordNameError(name_tuple)
 
344
                    raise errors.DuplicateRecordNameError(name_tuple[0])
284
345
                all_names.add(name_tuple)
285
346
        excess_bytes = self.reader_func(1)
286
347
        if excess_bytes != '':
307
368
        except ValueError:
308
369
            raise errors.InvalidRecordError(
309
370
                "%r is not a valid length." % (length_line,))
310
 
        
 
371
 
311
372
        # Read the list of names.
312
373
        names = []
313
374
        while True:
346
407
                _check_name_encoding(name)
347
408
        read_bytes(None)
348
409
 
 
410
 
 
411
class ContainerPushParser(object):
 
412
    """A "push" parser for container format 1.
 
413
 
 
414
    It accepts bytes via the ``accept_bytes`` method, and parses them into
 
415
    records which can be retrieved via the ``read_pending_records`` method.
 
416
    """
 
417
 
 
418
    def __init__(self):
 
419
        self._buffer = ''
 
420
        self._state_handler = self._state_expecting_format_line
 
421
        self._parsed_records = []
 
422
        self._reset_current_record()
 
423
        self.finished = False
 
424
 
 
425
    def _reset_current_record(self):
 
426
        self._current_record_length = None
 
427
        self._current_record_names = []
 
428
 
 
429
    def accept_bytes(self, bytes):
 
430
        self._buffer += bytes
 
431
        # Keep iterating the state machine until it stops consuming bytes from
 
432
        # the buffer.
 
433
        last_buffer_length = None
 
434
        cur_buffer_length = len(self._buffer)
 
435
        last_state_handler = None
 
436
        while (cur_buffer_length != last_buffer_length
 
437
               or last_state_handler != self._state_handler):
 
438
            last_buffer_length = cur_buffer_length
 
439
            last_state_handler = self._state_handler
 
440
            self._state_handler()
 
441
            cur_buffer_length = len(self._buffer)
 
442
 
 
443
    def read_pending_records(self, max=None):
 
444
        if max:
 
445
            records = self._parsed_records[:max]
 
446
            del self._parsed_records[:max]
 
447
            return records
 
448
        else:
 
449
            records = self._parsed_records
 
450
            self._parsed_records = []
 
451
            return records
 
452
 
 
453
    def _consume_line(self):
 
454
        """Take a line out of the buffer, and return the line.
 
455
 
 
456
        If a newline byte is not found in the buffer, the buffer is
 
457
        unchanged and this returns None instead.
 
458
        """
 
459
        newline_pos = self._buffer.find('\n')
 
460
        if newline_pos != -1:
 
461
            line = self._buffer[:newline_pos]
 
462
            self._buffer = self._buffer[newline_pos+1:]
 
463
            return line
 
464
        else:
 
465
            return None
 
466
 
 
467
    def _state_expecting_format_line(self):
 
468
        line = self._consume_line()
 
469
        if line is not None:
 
470
            if line != FORMAT_ONE:
 
471
                raise errors.UnknownContainerFormatError(line)
 
472
            self._state_handler = self._state_expecting_record_type
 
473
 
 
474
    def _state_expecting_record_type(self):
 
475
        if len(self._buffer) >= 1:
 
476
            record_type = self._buffer[0]
 
477
            self._buffer = self._buffer[1:]
 
478
            if record_type == 'B':
 
479
                self._state_handler = self._state_expecting_length
 
480
            elif record_type == 'E':
 
481
                self.finished = True
 
482
                self._state_handler = self._state_expecting_nothing
 
483
            else:
 
484
                raise errors.UnknownRecordTypeError(record_type)
 
485
 
 
486
    def _state_expecting_length(self):
 
487
        line = self._consume_line()
 
488
        if line is not None:
 
489
            try:
 
490
                self._current_record_length = int(line)
 
491
            except ValueError:
 
492
                raise errors.InvalidRecordError(
 
493
                    "%r is not a valid length." % (line,))
 
494
            self._state_handler = self._state_expecting_name
 
495
 
 
496
    def _state_expecting_name(self):
 
497
        encoded_name_parts = self._consume_line()
 
498
        if encoded_name_parts == '':
 
499
            self._state_handler = self._state_expecting_body
 
500
        elif encoded_name_parts:
 
501
            name_parts = tuple(encoded_name_parts.split('\x00'))
 
502
            for name_part in name_parts:
 
503
                _check_name(name_part)
 
504
            self._current_record_names.append(name_parts)
 
505
 
 
506
    def _state_expecting_body(self):
 
507
        if len(self._buffer) >= self._current_record_length:
 
508
            body_bytes = self._buffer[:self._current_record_length]
 
509
            self._buffer = self._buffer[self._current_record_length:]
 
510
            record = (self._current_record_names, body_bytes)
 
511
            self._parsed_records.append(record)
 
512
            self._reset_current_record()
 
513
            self._state_handler = self._state_expecting_record_type
 
514
 
 
515
    def _state_expecting_nothing(self):
 
516
        pass
 
517
 
 
518
    def read_size_hint(self):
 
519
        hint = 16384
 
520
        if self._state_handler == self._state_expecting_body:
 
521
            remaining = self._current_record_length - len(self._buffer)
 
522
            if remaining < 0:
 
523
                remaining = 0
 
524
            return max(hint, remaining)
 
525
        return hint
 
526
 
 
527
 
 
528
def iter_records_from_file(source_file):
 
529
    parser = ContainerPushParser()
 
530
    while True:
 
531
        bytes = source_file.read(parser.read_size_hint())
 
532
        parser.accept_bytes(bytes)
 
533
        for record in parser.read_pending_records():
 
534
            yield record
 
535
        if parser.finished:
 
536
            break
 
537