← Back to branch summary

~bzr-pqm/bzr/bzr.dev

« back to all changes in this revision

Viewing changes to bzrlib/_bencode_pyx.pyx

  • Committer: Robert Collins
  • Date: 2007-07-15 15:40:37 UTC
  • mto: (2592.3.33 repository)
  • mto: This revision was merged to the branch mainline in revision 2624.
  • Revision ID: robertc@robertcollins.net-20070715154037-3ar8g89decddc9su
Make GraphIndex accept nodes as key, value, references, so that the method
signature is closer to what a simple key->value index delivers. Also
change the behaviour when the reference list count is zero to accept
key, value as nodes, and emit key, value to make it identical in that case
to a simple key->value index. This may not be a good idea, but for now it
seems ok.

Show diffs side-by-side

added added

removed removed

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