← Back to branch summary

~bzr-pqm/bzr/bzr.dev

« back to all changes in this revision

Viewing changes to bzrlib/_bencode_c.pyx

  • Committer: John Arbash Meinel
  • Date: 2009-06-03 01:37:27 UTC
  • mfrom: (2694.5.22 bencode-pyrex)
  • mto: This revision was merged to the branch mainline in revision 4410.
  • Revision ID: john@arbash-meinel.com-20090603013727-wtxochxrd0zqf1vg
Merge the bencode implementation.

Show diffs side-by-side

added added

removed removed

Lines of Context:
bzrlib/_bencode_c.pyx
 
1
# Copyright (C) 2007,2009 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., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
 
16
 
 
17
"""Pyrex implementation for bencode coder/decoder"""
 
18
 
 
19
 
 
20
cdef extern from "stddef.h":
 
21
    ctypedef unsigned int size_t
 
22
 
 
23
cdef extern from "Python.h":
 
24
    ctypedef int  Py_ssize_t
 
25
    int PyInt_CheckExact(object o)
 
26
    int PyLong_CheckExact(object o)
 
27
    int PyString_CheckExact(object o)
 
28
    int PyTuple_CheckExact(object o)
 
29
    int PyList_CheckExact(object o)
 
30
    int PyDict_CheckExact(object o)
 
31
    int PyBool_Check(object o)
 
32
    object PyString_FromStringAndSize(char *v, Py_ssize_t len)
 
33
    char *PyString_AS_STRING(object o) except NULL
 
34
    Py_ssize_t PyString_GET_SIZE(object o) except -1
 
35
    object PyInt_FromString(char *str, char **pend, int base)
 
36
    int Py_GetRecursionLimit()
 
37
    int Py_EnterRecursiveCall(char *)
 
38
    void Py_LeaveRecursiveCall()
 
39
 
 
40
cdef extern from "stdlib.h":
 
41
    void free(void *memblock)
 
42
    void *malloc(size_t size)
 
43
    void *realloc(void *memblock, size_t size)
 
44
    long strtol(char *, char **, int)
 
45
 
 
46
cdef extern from "string.h":
 
47
    void *memcpy(void *dest, void *src, size_t count)
 
48
 
 
49
cdef extern from "python-compat.h":
 
50
    int snprintf(char* buffer, size_t nsize, char* fmt, ...)
 
51
 
 
52
 
 
53
cdef class Decoder:
 
54
    """Bencode decoder"""
 
55
 
 
56
    cdef readonly char *tail
 
57
    cdef readonly int size
 
58
    cdef readonly int _yield_tuples
 
59
    cdef object text
 
60
 
 
61
    def __init__(self, s, yield_tuples=0):
 
62
        """Initialize decoder engine.
 
63
        @param  s:  Python string.
 
64
        """
 
65
        if not PyString_CheckExact(s):
 
66
            raise TypeError("String required")
 
67
 
 
68
        self.text = s
 
69
        self.tail = PyString_AS_STRING(s)
 
70
        self.size = PyString_GET_SIZE(s)
 
71
        self._yield_tuples = int(yield_tuples)
 
72
 
 
73
    def decode(self):
 
74
        result = self.decode_object()
 
75
        if self.size != 0:
 
76
            raise ValueError('junk in stream')
 
77
        return result
 
78
 
 
79
    def decode_object(self):
 
80
        cdef char ch
 
81
 
 
82
        if 0 == self.size:
 
83
            raise ValueError('stream underflow')
 
84
 
 
85
        if Py_EnterRecursiveCall("decode_object"):
 
86
            raise RuntimeError("too deeply nested")
 
87
        try:
 
88
            ch = self.tail[0]
 
89
            if ch == c'i':
 
90
                self._update_tail(1)
 
91
                return self._decode_int()
 
92
            elif c'0' <= ch <= c'9':
 
93
                return self._decode_string()
 
94
            elif ch == c'l':
 
95
                self._update_tail(1)
 
96
                return self._decode_list()
 
97
            elif ch == c'd':
 
98
                self._update_tail(1)
 
99
                return self._decode_dict()
 
100
            else:
 
101
                raise ValueError('unknown object type identifier %r' % ch)
 
102
        finally:
 
103
            Py_LeaveRecursiveCall()
 
104
 
 
105
    cdef void _update_tail(self, int n):
 
106
        """Update tail pointer and resulting size by n characters"""
 
107
        self.size = self.size - n
 
108
        self.tail = &self.tail[n]
 
109
 
 
110
    cdef int _read_digits(self, char stop_char) except -1:
 
111
        cdef int i
 
112
        i = 0
 
113
        while ((self.tail[i] >= c'0' and self.tail[i] <= c'9') or 
 
114
               self.tail[i] == c'-') and i < self.size:
 
115
            i = i + 1
 
116
 
 
117
        if self.tail[i] != stop_char:
 
118
            raise ValueError("Stop character %c not found: %c" % 
 
119
                (stop_char, self.tail[i]))
 
120
        if (self.tail[0] == c'0' or 
 
121
                (self.tail[0] == c'-' and self.tail[1] == c'0')):
 
122
            if i == 1:
 
123
                return i
 
124
            else:
 
125
                raise ValueError # leading zeroes are not allowed
 
126
        return i
 
127
 
 
128
    cdef object _decode_int(self):
 
129
        cdef int i
 
130
        i = self._read_digits(c'e')
 
131
        self.tail[i] = 0
 
132
        try:
 
133
            ret = PyInt_FromString(self.tail, NULL, 10)
 
134
        finally:
 
135
            self.tail[i] = c'e'
 
136
        self._update_tail(i+1)
 
137
        return ret
 
138
 
 
139
    cdef object _decode_string(self):
 
140
        cdef int n, i
 
141
        i = self._read_digits(c':')
 
142
        n = strtol(self.tail, NULL, 10)
 
143
        self._update_tail(i+1)
 
144
        if n == 0:
 
145
            return ''
 
146
        if n > self.size:
 
147
            raise ValueError('stream underflow')
 
148
        if n < 0:
 
149
            raise ValueError('string size below zero: %d' % n)
 
150
 
 
151
        result = PyString_FromStringAndSize(self.tail, n)
 
152
        self._update_tail(n)
 
153
        return result
 
154
 
 
155
    cdef object _decode_list(self):
 
156
        result = []
 
157
 
 
158
        while self.size > 0:
 
159
            if self.tail[0] == c'e':
 
160
                self._update_tail(1)
 
161
                if self._yield_tuples:
 
162
                    return tuple(result)
 
163
                else:
 
164
                    return result
 
165
            else:
 
166
                result.append(self.decode_object())
 
167
 
 
168
        raise ValueError('malformed list')
 
169
 
 
170
    cdef object _decode_dict(self):
 
171
        cdef char ch
 
172
 
 
173
        result = {}
 
174
        lastkey = None
 
175
 
 
176
        while self.size > 0:
 
177
            ch = self.tail[0]
 
178
            if ch == c'e':
 
179
                self._update_tail(1)
 
180
                return result
 
181
            else:
 
182
                # keys should be strings only
 
183
                key = self._decode_string()
 
184
                if lastkey >= key:
 
185
                    raise ValueError('dict keys disordered')
 
186
                else:
 
187
                    lastkey = key
 
188
                value = self.decode_object()
 
189
                result[key] = value
 
190
 
 
191
        raise ValueError('malformed dict')
 
192
 
 
193
 
 
194
def bdecode(object s):
 
195
    """Decode string x to Python object"""
 
196
    return Decoder(s).decode()
 
197
 
 
198
 
 
199
def bdecode_as_tuple(object s):
 
200
    """Decode string x to Python object, using tuples rather than lists."""
 
201
    return Decoder(s, True).decode()
 
202
 
 
203
 
 
204
class Bencached(object):
 
205
    __slots__ = ['bencoded']
 
206
 
 
207
    def __init__(self, s):
 
208
        self.bencoded = s
 
209
 
 
210
 
 
211
cdef enum:
 
212
    INITSIZE = 1024     # initial size for encoder buffer
 
213
    INT_BUF_SIZE = 32
 
214
 
 
215
 
 
216
cdef class Encoder:
 
217
    """Bencode encoder"""
 
218
 
 
219
    cdef readonly char *buffer
 
220
    cdef readonly int maxsize
 
221
    cdef readonly char *tail
 
222
    cdef readonly int size
 
223
 
 
224
    def __init__(self, int maxsize=INITSIZE):
 
225
        """Initialize encoder engine
 
226
        @param  maxsize:    initial size of internal char buffer
 
227
        """
 
228
        cdef char *p
 
229
 
 
230
        self.maxsize = 0
 
231
        self.size = 0
 
232
        self.tail = NULL
 
233
 
 
234
        p = <char*>malloc(maxsize)
 
235
        if p == NULL:
 
236
            raise MemoryError('Not enough memory to allocate buffer '
 
237
                              'for encoder')
 
238
        self.buffer = p
 
239
        self.maxsize = maxsize
 
240
        self.tail = p
 
241
 
 
242
    def __del__(self):
 
243
        free(self.buffer)
 
244
        self.buffer = NULL
 
245
        self.maxsize = 0
 
246
 
 
247
    def __str__(self):
 
248
        if self.buffer != NULL and self.size != 0:
 
249
            return PyString_FromStringAndSize(self.buffer, self.size)
 
250
        else:
 
251
            return ''
 
252
 
 
253
    cdef int _ensure_buffer(self, int required) except 0:
 
254
        """Ensure that tail of CharTail buffer has enough size.
 
255
        If buffer is not big enough then function try to
 
256
        realloc buffer.
 
257
        """
 
258
        cdef char *new_buffer
 
259
        cdef int   new_size
 
260
 
 
261
        if self.size + required < self.maxsize:
 
262
            return 1
 
263
 
 
264
        new_size = self.maxsize
 
265
        while new_size < self.size + required:
 
266
            new_size = new_size * 2
 
267
        new_buffer = <char*>realloc(self.buffer, <size_t>new_size)
 
268
        if new_buffer == NULL:
 
269
            raise MemoryError('Cannot realloc buffer for encoder')
 
270
 
 
271
        self.buffer = new_buffer
 
272
        self.maxsize = new_size
 
273
        self.tail = &new_buffer[self.size]
 
274
        return 1
 
275
 
 
276
    cdef void _update_tail(self, int n):
 
277
        """Update tail pointer and resulting size by n characters"""
 
278
        self.size = self.size + n
 
279
        self.tail = &self.tail[n]
 
280
 
 
281
    cdef int _encode_int(self, int x) except 0:
 
282
        """Encode int to bencode string iNNNe
 
283
        @param  x:  value to encode
 
284
        """
 
285
        cdef int n
 
286
        self._ensure_buffer(INT_BUF_SIZE)
 
287
        n = snprintf(self.tail, INT_BUF_SIZE, "i%de", x)
 
288
        if n < 0:
 
289
            raise MemoryError('int %d too big to encode' % x)
 
290
        self._update_tail(n)
 
291
        return 1
 
292
 
 
293
    cdef int _encode_long(self, x) except 0:
 
294
        return self._append_string(''.join(('i', str(x), 'e')))
 
295
 
 
296
    cdef int _append_string(self, s) except 0:
 
297
        self._ensure_buffer(PyString_GET_SIZE(s))
 
298
        memcpy(self.tail, PyString_AS_STRING(s), PyString_GET_SIZE(s))
 
299
        self._update_tail(PyString_GET_SIZE(s))
 
300
        return 1
 
301
 
 
302
    cdef int _encode_string(self, x) except 0:
 
303
        cdef int n
 
304
        self._ensure_buffer(PyString_GET_SIZE(x) + 32)
 
305
        n = snprintf(self.tail, 32, '%d:', PyString_GET_SIZE(x))
 
306
        if n < 0:
 
307
            raise MemoryError('string %s too big to encode' % x)
 
308
        memcpy(<void *>(self.tail+n), PyString_AS_STRING(x),
 
309
               PyString_GET_SIZE(x))
 
310
        self._update_tail(n+PyString_GET_SIZE(x))
 
311
        return 1
 
312
 
 
313
    cdef int _encode_list(self, x) except 0:
 
314
        self._ensure_buffer(2)
 
315
        self.tail[0] = c'l'
 
316
        self._update_tail(1)
 
317
 
 
318
        for i in x:
 
319
            self.process(i)
 
320
 
 
321
        self.tail[0] = c'e'
 
322
        self._update_tail(1)
 
323
        return 1
 
324
 
 
325
    cdef int _encode_dict(self, x) except 0:
 
326
        self._ensure_buffer(2)
 
327
        self.tail[0] = c'd'
 
328
        self._update_tail(1)
 
329
 
 
330
        keys = x.keys()
 
331
        keys.sort()
 
332
        for k in keys:
 
333
            if not PyString_CheckExact(k):
 
334
                raise TypeError('key in dict should be string')
 
335
            self._encode_string(k)
 
336
            self.process(x[k])
 
337
 
 
338
        self.tail[0] = c'e'
 
339
        self._update_tail(1)
 
340
        return 1
 
341
 
 
342
    def process(self, object x):
 
343
        if Py_EnterRecursiveCall("encode"):
 
344
            raise RuntimeError("too deeply nested")
 
345
        try:
 
346
            if PyString_CheckExact(x):
 
347
                self._encode_string(x)
 
348
            elif PyInt_CheckExact(x):
 
349
                self._encode_int(x)
 
350
            elif PyLong_CheckExact(x):
 
351
                self._encode_long(x)
 
352
            elif PyList_CheckExact(x) or PyTuple_CheckExact(x):
 
353
                self._encode_list(x)
 
354
            elif PyDict_CheckExact(x):
 
355
                self._encode_dict(x)
 
356
            elif PyBool_Check(x):
 
357
                self._encode_int(int(x))
 
358
            elif isinstance(x, Bencached):
 
359
                self._append_string(x.bencoded)
 
360
            else:
 
361
                raise TypeError('unsupported type %r' % x)
 
362
        finally:
 
363
            Py_LeaveRecursiveCall()
 
364
 
 
365
 
 
366
def bencode(x):
 
367
    """Encode Python object x to string"""
 
368
    encoder = Encoder()
 
369
    encoder.process(x)
 
370
    return str(encoder)