/brz/remove-bazaar

To get this branch, use:
bzr branch http://gegoxaren.bato24.eu/bzr/brz/remove-bazaar

« back to all changes in this revision

Viewing changes to breezy/_bencode_pyx.pyx

[merge] robertc's integration, updated tests to check for retcode=3

Show diffs side-by-side

added added

removed removed

Lines of Context:
breezy/_bencode_pyx.pyx
1
 
# Copyright (C) 2007, 2009, 2010 Canonical Ltd
2
 
#
3
 
# This program is free software; you can redistribute it and/or modify
4
 
# it under the terms of the GNU General Public License as published by
5
 
# the Free Software Foundation; either version 2 of the License, or
6
 
# (at your option) any later version.
7
 
#
8
 
# This program is distributed in the hope that it will be useful,
9
 
# but WITHOUT ANY WARRANTY; without even the implied warranty of
10
 
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
11
 
# GNU General Public License for more details.
12
 
#
13
 
# You should have received a copy of the GNU General Public License
14
 
# along with this program; if not, write to the Free Software
15
 
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
16
 
 
17
 
"""Pyrex implementation for bencode coder/decoder"""
18
 
 
19
 
from __future__ import absolute_import
20
 
 
21
 
from cpython.bool cimport (
22
 
    PyBool_Check,
23
 
    )
24
 
from cpython.bytes cimport (
25
 
    PyBytes_CheckExact,
26
 
    PyBytes_FromStringAndSize,
27
 
    PyBytes_AS_STRING,
28
 
    PyBytes_GET_SIZE,
29
 
    )
30
 
from cpython.dict cimport (
31
 
    PyDict_CheckExact,
32
 
    )
33
 
from cpython.int cimport (
34
 
    PyInt_CheckExact,
35
 
    PyInt_FromString,
36
 
    )
37
 
from cpython.list cimport (
38
 
    PyList_CheckExact,
39
 
    PyList_Append,
40
 
    )
41
 
from cpython.long cimport (
42
 
    PyLong_CheckExact,
43
 
    )
44
 
from cpython.mem cimport (
45
 
    PyMem_Free,
46
 
    PyMem_Malloc,
47
 
    PyMem_Realloc,
48
 
    )
49
 
from cpython.tuple cimport (
50
 
    PyTuple_CheckExact,
51
 
    )
52
 
 
53
 
from libc.stdlib cimport (
54
 
    strtol,
55
 
    )
56
 
from libc.string cimport (
57
 
    memcpy,
58
 
    )
59
 
 
60
 
cdef extern from "python-compat.h":
61
 
    int snprintf(char* buffer, size_t nsize, char* fmt, ...)
62
 
    # Use wrapper with inverted error return so Cython can propogate
63
 
    int BrzPy_EnterRecursiveCall(char *) except 0
64
 
 
65
 
cdef extern from "Python.h":
66
 
    void Py_LeaveRecursiveCall()
67
 
 
68
 
cdef class Decoder
69
 
cdef class Encoder
70
 
 
71
 
cdef extern from "_bencode_pyx.h":
72
 
    void D_UPDATE_TAIL(Decoder, int n)
73
 
    void E_UPDATE_TAIL(Encoder, int n)
74
 
 
75
 
from ._static_tuple_c cimport StaticTuple, StaticTuple_CheckExact, \
76
 
    import_static_tuple_c
77
 
 
78
 
import_static_tuple_c()
79
 
 
80
 
 
81
 
cdef class Decoder:
82
 
    """Bencode decoder"""
83
 
 
84
 
    cdef readonly char *tail
85
 
    cdef readonly int size
86
 
    cdef readonly int _yield_tuples
87
 
    cdef object text
88
 
 
89
 
    def __init__(self, s, yield_tuples=0):
90
 
        """Initialize decoder engine.
91
 
        @param  s:  Python string.
92
 
        """
93
 
        if not PyBytes_CheckExact(s):
94
 
            raise TypeError("bytes required")
95
 
 
96
 
        self.text = s
97
 
        self.tail = PyBytes_AS_STRING(s)
98
 
        self.size = PyBytes_GET_SIZE(s)
99
 
        self._yield_tuples = int(yield_tuples)
100
 
 
101
 
    def decode(self):
102
 
        result = self._decode_object()
103
 
        if self.size != 0:
104
 
            raise ValueError('junk in stream')
105
 
        return result
106
 
 
107
 
    def decode_object(self):
108
 
        return self._decode_object()
109
 
 
110
 
    cdef object _decode_object(self):
111
 
        cdef char ch
112
 
 
113
 
        if 0 == self.size:
114
 
            raise ValueError('stream underflow')
115
 
 
116
 
        BrzPy_EnterRecursiveCall(" while bencode decoding")
117
 
        try:
118
 
            ch = self.tail[0]
119
 
            if c'0' <= ch <= c'9':
120
 
                return self._decode_string()
121
 
            elif ch == c'l':
122
 
                D_UPDATE_TAIL(self, 1)
123
 
                return self._decode_list()
124
 
            elif ch == c'i':
125
 
                D_UPDATE_TAIL(self, 1)
126
 
                return self._decode_int()
127
 
            elif ch == c'd':
128
 
                D_UPDATE_TAIL(self, 1)
129
 
                return self._decode_dict()
130
 
        finally:
131
 
            Py_LeaveRecursiveCall()
132
 
        raise ValueError('unknown object type identifier %r' % ch)
133
 
 
134
 
    cdef int _read_digits(self, char stop_char) except -1:
135
 
        cdef int i
136
 
        i = 0
137
 
        while ((self.tail[i] >= c'0' and self.tail[i] <= c'9') or
138
 
               self.tail[i] == c'-') and i < self.size:
139
 
            i = i + 1
140
 
 
141
 
        if self.tail[i] != stop_char:
142
 
            raise ValueError("Stop character %c not found: %c" % 
143
 
                (stop_char, self.tail[i]))
144
 
        if (self.tail[0] == c'0' or 
145
 
                (self.tail[0] == c'-' and self.tail[1] == c'0')):
146
 
            if i == 1:
147
 
                return i
148
 
            else:
149
 
                raise ValueError # leading zeroes are not allowed
150
 
        return i
151
 
 
152
 
    cdef object _decode_int(self):
153
 
        cdef int i
154
 
        i = self._read_digits(c'e')
155
 
        self.tail[i] = 0
156
 
        try:
157
 
            ret = PyInt_FromString(self.tail, NULL, 10)
158
 
        finally:
159
 
            self.tail[i] = c'e'
160
 
        D_UPDATE_TAIL(self, i+1)
161
 
        return ret
162
 
 
163
 
    cdef object _decode_string(self):
164
 
        cdef int n
165
 
        cdef char *next_tail
166
 
        # strtol allows leading whitespace, negatives, and leading zeros
167
 
        # however, all callers have already checked that '0' <= tail[0] <= '9'
168
 
        # or they wouldn't have called _decode_string
169
 
        # strtol will stop at trailing whitespace, etc
170
 
        n = strtol(self.tail, &next_tail, 10)
171
 
        if next_tail == NULL or next_tail[0] != c':':
172
 
            raise ValueError('string len not terminated by ":"')
173
 
        # strtol allows leading zeros, so validate that we don't have that
174
 
        if (self.tail[0] == c'0'
175
 
            and (n != 0 or (next_tail - self.tail != 1))):
176
 
            raise ValueError('leading zeros are not allowed')
177
 
        D_UPDATE_TAIL(self, next_tail - self.tail + 1)
178
 
        if n == 0:
179
 
            return b''
180
 
        if n > self.size:
181
 
            raise ValueError('stream underflow')
182
 
        if n < 0:
183
 
            raise ValueError('string size below zero: %d' % n)
184
 
 
185
 
        result = PyBytes_FromStringAndSize(self.tail, n)
186
 
        D_UPDATE_TAIL(self, n)
187
 
        return result
188
 
 
189
 
    cdef object _decode_list(self):
190
 
        result = []
191
 
 
192
 
        while self.size > 0:
193
 
            if self.tail[0] == c'e':
194
 
                D_UPDATE_TAIL(self, 1)
195
 
                if self._yield_tuples:
196
 
                    return tuple(result)
197
 
                else:
198
 
                    return result
199
 
            else:
200
 
                # As a quick shortcut, check to see if the next object is a
201
 
                # string, since we know that won't be creating recursion
202
 
                # if self.tail[0] >= c'0' and self.tail[0] <= c'9':
203
 
                PyList_Append(result, self._decode_object())
204
 
 
205
 
        raise ValueError('malformed list')
206
 
 
207
 
    cdef object _decode_dict(self):
208
 
        cdef char ch
209
 
 
210
 
        result = {}
211
 
        lastkey = None
212
 
 
213
 
        while self.size > 0:
214
 
            ch = self.tail[0]
215
 
            if ch == c'e':
216
 
                D_UPDATE_TAIL(self, 1)
217
 
                return result
218
 
            else:
219
 
                # keys should be strings only
220
 
                if self.tail[0] < c'0' or self.tail[0] > c'9':
221
 
                    raise ValueError('key was not a simple string.')
222
 
                key = self._decode_string()
223
 
                if lastkey is not None and lastkey >= key:
224
 
                    raise ValueError('dict keys disordered')
225
 
                else:
226
 
                    lastkey = key
227
 
                value = self._decode_object()
228
 
                result[key] = value
229
 
 
230
 
        raise ValueError('malformed dict')
231
 
 
232
 
 
233
 
def bdecode(object s):
234
 
    """Decode string x to Python object"""
235
 
    return Decoder(s).decode()
236
 
 
237
 
 
238
 
def bdecode_as_tuple(object s):
239
 
    """Decode string x to Python object, using tuples rather than lists."""
240
 
    return Decoder(s, True).decode()
241
 
 
242
 
 
243
 
class Bencached(object):
244
 
    __slots__ = ['bencoded']
245
 
 
246
 
    def __init__(self, s):
247
 
        self.bencoded = s
248
 
 
249
 
 
250
 
cdef enum:
251
 
    INITSIZE = 1024     # initial size for encoder buffer
252
 
    INT_BUF_SIZE = 32
253
 
 
254
 
 
255
 
cdef class Encoder:
256
 
    """Bencode encoder"""
257
 
 
258
 
    cdef readonly char *tail
259
 
    cdef readonly int size
260
 
    cdef readonly char *buffer
261
 
    cdef readonly int maxsize
262
 
 
263
 
    def __init__(self, int maxsize=INITSIZE):
264
 
        """Initialize encoder engine
265
 
        @param  maxsize:    initial size of internal char buffer
266
 
        """
267
 
        cdef char *p
268
 
 
269
 
        self.maxsize = 0
270
 
        self.size = 0
271
 
        self.tail = NULL
272
 
 
273
 
        p = <char*>PyMem_Malloc(maxsize)
274
 
        if p == NULL:
275
 
            raise MemoryError('Not enough memory to allocate buffer '
276
 
                              'for encoder')
277
 
        self.buffer = p
278
 
        self.maxsize = maxsize
279
 
        self.tail = p
280
 
 
281
 
    def __dealloc__(self):
282
 
        PyMem_Free(self.buffer)
283
 
        self.buffer = NULL
284
 
        self.maxsize = 0
285
 
 
286
 
    def to_bytes(self):
287
 
        if self.buffer != NULL and self.size != 0:
288
 
            return PyBytes_FromStringAndSize(self.buffer, self.size)
289
 
        return b''
290
 
 
291
 
    cdef int _ensure_buffer(self, int required) except 0:
292
 
        """Ensure that tail of CharTail buffer has enough size.
293
 
        If buffer is not big enough then function try to
294
 
        realloc buffer.
295
 
        """
296
 
        cdef char *new_buffer
297
 
        cdef int   new_size
298
 
 
299
 
        if self.size + required < self.maxsize:
300
 
            return 1
301
 
 
302
 
        new_size = self.maxsize
303
 
        while new_size < self.size + required:
304
 
            new_size = new_size * 2
305
 
        new_buffer = <char*>PyMem_Realloc(self.buffer, <size_t>new_size)
306
 
        if new_buffer == NULL:
307
 
            raise MemoryError('Cannot realloc buffer for encoder')
308
 
 
309
 
        self.buffer = new_buffer
310
 
        self.maxsize = new_size
311
 
        self.tail = &new_buffer[self.size]
312
 
        return 1
313
 
 
314
 
    cdef int _encode_int(self, int x) except 0:
315
 
        """Encode int to bencode string iNNNe
316
 
        @param  x:  value to encode
317
 
        """
318
 
        cdef int n
319
 
        self._ensure_buffer(INT_BUF_SIZE)
320
 
        n = snprintf(self.tail, INT_BUF_SIZE, b"i%de", x)
321
 
        if n < 0:
322
 
            raise MemoryError('int %d too big to encode' % x)
323
 
        E_UPDATE_TAIL(self, n)
324
 
        return 1
325
 
 
326
 
    cdef int _encode_long(self, x) except 0:
327
 
        return self._append_string(b'i%de' % x)
328
 
 
329
 
    cdef int _append_string(self, s) except 0:
330
 
        cdef Py_ssize_t n
331
 
        n = PyBytes_GET_SIZE(s)
332
 
        self._ensure_buffer(n)
333
 
        memcpy(self.tail, PyBytes_AS_STRING(s), n)
334
 
        E_UPDATE_TAIL(self, n)
335
 
        return 1
336
 
 
337
 
    cdef int _encode_string(self, x) except 0:
338
 
        cdef int n
339
 
        cdef Py_ssize_t x_len
340
 
        x_len = PyBytes_GET_SIZE(x)
341
 
        self._ensure_buffer(x_len + INT_BUF_SIZE)
342
 
        n = snprintf(self.tail, INT_BUF_SIZE, b'%ld:', x_len)
343
 
        if n < 0:
344
 
            raise MemoryError('string %s too big to encode' % x)
345
 
        memcpy(<void *>(self.tail+n), PyBytes_AS_STRING(x), x_len)
346
 
        E_UPDATE_TAIL(self, n + x_len)
347
 
        return 1
348
 
 
349
 
    cdef int _encode_list(self, x) except 0:
350
 
        self._ensure_buffer(1)
351
 
        self.tail[0] = c'l'
352
 
        E_UPDATE_TAIL(self, 1)
353
 
 
354
 
        for i in x:
355
 
            self.process(i)
356
 
 
357
 
        self._ensure_buffer(1)
358
 
        self.tail[0] = c'e'
359
 
        E_UPDATE_TAIL(self, 1)
360
 
        return 1
361
 
 
362
 
    cdef int _encode_dict(self, x) except 0:
363
 
        self._ensure_buffer(1)
364
 
        self.tail[0] = c'd'
365
 
        E_UPDATE_TAIL(self, 1)
366
 
 
367
 
        for k in sorted(x):
368
 
            if not PyBytes_CheckExact(k):
369
 
                raise TypeError('key in dict should be string')
370
 
            self._encode_string(k)
371
 
            self.process(x[k])
372
 
 
373
 
        self._ensure_buffer(1)
374
 
        self.tail[0] = c'e'
375
 
        E_UPDATE_TAIL(self, 1)
376
 
        return 1
377
 
 
378
 
    cpdef object process(self, object x):
379
 
        BrzPy_EnterRecursiveCall(" while bencode encoding")
380
 
        try:
381
 
            if PyBytes_CheckExact(x):
382
 
                self._encode_string(x)
383
 
            elif PyInt_CheckExact(x) and x.bit_length() < 32:
384
 
                self._encode_int(x)
385
 
            elif PyLong_CheckExact(x):
386
 
                self._encode_long(x)
387
 
            elif (PyList_CheckExact(x) or PyTuple_CheckExact(x)
388
 
                  or isinstance(x, StaticTuple)):
389
 
                self._encode_list(x)
390
 
            elif PyDict_CheckExact(x):
391
 
                self._encode_dict(x)
392
 
            elif PyBool_Check(x):
393
 
                self._encode_int(int(x))
394
 
            elif isinstance(x, Bencached):
395
 
                self._append_string(x.bencoded)
396
 
            else:
397
 
                raise TypeError('unsupported type %r' % x)
398
 
        finally:
399
 
            Py_LeaveRecursiveCall()
400
 
 
401
 
 
402
 
def bencode(x):
403
 
    """Encode Python object x to string"""
404
 
    encoder = Encoder()
405
 
    encoder.process(x)
406
 
    return encoder.to_bytes()