Skip to content
This repository was archived by the owner on Dec 10, 2018. It is now read-only.

Commit 7b7e083

Browse files
lxyulxyu
committed
Merge pull request #16 from lodevil/master, close #16
new cython binary protocol
2 parents 003d223 + 892217d commit 7b7e083

File tree

11 files changed

+528
-547
lines changed

11 files changed

+528
-547
lines changed

‎setup.py

+4-4
Original file line numberDiff line numberDiff line change
@@ -27,12 +27,12 @@
2727
ext_modules = []
2828

2929
if cython:
30-
ext_modules.append(Extension("thriftpy.protocol.cybinary",
31-
["thriftpy/protocol/cybinary.pyx"]))
30+
ext_modules.append(Extension("thriftpy.protocol.cybin",
31+
["thriftpy/protocol/cybin/cybin.pyx"]))
3232
cmdclass["build_ext"] = build_ext
3333
else:
34-
ext_modules.append(Extension("thriftpy.protocol.cybinary",
35-
["thriftpy/protocol/cybinary.c"]))
34+
ext_modules.append(Extension("thriftpy.protocol.cybin",
35+
["thriftpy/protocol/cybin/cybin.c"]))
3636

3737

3838
setup(name="thriftpy",

‎tests/test_protocol_cybinary.py

+61-36
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,14 @@
11
# -*- coding: utf-8 -*-
22

3-
from io import BytesIO
3+
# from io import BytesIO
44

55
from nose.tools import assert_equal
66

77
from thriftpy._compat import u
88
from thriftpy.thrift import TType, TPayload
99
from thriftpy.utils import hexlify
10-
from thriftpy.protocol import cybinary as proto
10+
from thriftpy.protocol import cybin as proto
11+
from thriftpy.transport import TMemoryBuffer
1112

1213

1314
class TItem(TPayload):
@@ -19,103 +20,127 @@ class TItem(TPayload):
1920

2021

2122
def test_pack_i8():
22-
b = BytesIO()
23-
proto.write_val(b, TType.I08, 123)
23+
b = TMemoryBuffer()
24+
p = proto.TCyBinaryProtocol(b)
25+
p.write_val(TType.I08, 123)
26+
p.write_message_end()
2427
assert_equal("7b", hexlify(b.getvalue()))
2528

2629

2730
def test_unpack_i8():
28-
b = BytesIO(b"{")
29-
assert_equal(123, proto.read_val(b, TType.I08))
31+
b = TMemoryBuffer(b"{")
32+
p = proto.TCyBinaryProtocol(b)
33+
assert_equal(123, p.read_val(TType.I08))
3034

3135

3236
def test_pack_i16():
33-
b = BytesIO()
34-
proto.write_val(b, TType.I16, 12345)
37+
b = TMemoryBuffer()
38+
p = proto.TCyBinaryProtocol(b)
39+
p.write_val(TType.I16, 12345)
40+
p.write_message_end()
3541
assert_equal("30 39", hexlify(b.getvalue()))
3642

3743

3844
def test_unpack_i16():
39-
b = BytesIO(b"09")
40-
assert_equal(12345, proto.read_val(b, TType.I16))
45+
b = TMemoryBuffer(b"09")
46+
p = proto.TCyBinaryProtocol(b)
47+
assert_equal(12345, p.read_val(TType.I16))
4148

4249

4350
def test_pack_i32():
44-
b = BytesIO()
45-
proto.write_val(b, TType.I32, 1234567890)
51+
b = TMemoryBuffer()
52+
p = proto.TCyBinaryProtocol(b)
53+
p.write_val(TType.I32, 1234567890)
54+
p.write_message_end()
4655
assert_equal("49 96 02 d2", hexlify(b.getvalue()))
4756

4857

4958
def test_unpack_i32():
50-
b = BytesIO(b'I\x96\x02\xd2')
51-
assert_equal(1234567890, proto.read_val(b, TType.I32))
59+
b = TMemoryBuffer(b'I\x96\x02\xd2')
60+
p = proto.TCyBinaryProtocol(b)
61+
assert_equal(1234567890, p.read_val(TType.I32))
5262

5363

5464
def test_pack_i64():
55-
b = BytesIO()
56-
proto.write_val(b, TType.I64, 1234567890123456789)
65+
b = TMemoryBuffer()
66+
p = proto.TCyBinaryProtocol(b)
67+
p.write_val(TType.I64, 1234567890123456789)
68+
p.write_message_end()
5769
assert_equal("11 22 10 f4 7d e9 81 15", hexlify(b.getvalue()))
5870

5971

6072
def test_unpack_i64():
61-
b = BytesIO(b'\x11"\x10\xf4}\xe9\x81\x15')
62-
assert_equal(1234567890123456789, proto.read_val(b, TType.I64))
73+
b = TMemoryBuffer(b'\x11"\x10\xf4}\xe9\x81\x15')
74+
p = proto.TCyBinaryProtocol(b)
75+
assert_equal(1234567890123456789, p.read_val(TType.I64))
6376

6477

6578
def test_pack_double():
66-
b = BytesIO()
67-
proto.write_val(b, TType.DOUBLE, 1234567890.1234567890)
79+
b = TMemoryBuffer()
80+
p = proto.TCyBinaryProtocol(b)
81+
p.write_val(TType.DOUBLE, 1234567890.1234567890)
82+
p.write_message_end()
6883
assert_equal("41 d2 65 80 b4 87 e6 b7", hexlify(b.getvalue()))
6984

7085

7186
def test_unpack_double():
72-
b = BytesIO(b'A\xd2e\x80\xb4\x87\xe6\xb7')
73-
assert_equal(1234567890.1234567890, proto.read_val(b, TType.DOUBLE))
87+
b = TMemoryBuffer(b'A\xd2e\x80\xb4\x87\xe6\xb7')
88+
p = proto.TCyBinaryProtocol(b)
89+
assert_equal(1234567890.1234567890, p.read_val(TType.DOUBLE))
7490

7591

7692
def test_pack_string():
77-
b = BytesIO()
78-
proto.write_val(b, TType.STRING, "hello world!")
93+
b = TMemoryBuffer()
94+
p = proto.TCyBinaryProtocol(b)
95+
p.write_val(TType.STRING, "hello world!")
96+
p.write_message_end()
7997
assert_equal("00 00 00 0c 68 65 6c 6c 6f 20 77 6f 72 6c 64 21",
8098
hexlify(b.getvalue()))
8199

82-
b = BytesIO()
83-
proto.write_val(b, TType.STRING, u("你好世界"))
100+
b = TMemoryBuffer()
101+
p = proto.TCyBinaryProtocol(b)
102+
p.write_val(TType.STRING, u("你好世界"))
103+
p.write_message_end()
84104
assert_equal("00 00 00 0c e4 bd a0 e5 a5 bd e4 b8 96 e7 95 8c",
85105
hexlify(b.getvalue()))
86106

87107

88108
def test_unpack_string():
89-
b = BytesIO(b'\x00\x00\x00\x0c'
90-
b'\xe4\xbd\xa0\xe5\xa5\xbd\xe4\xb8\x96\xe7\x95\x8c')
91-
assert_equal(u("你好世界"), proto.read_val(b, TType.STRING))
109+
b = TMemoryBuffer(b'\x00\x00\x00\x0c'
110+
b'\xe4\xbd\xa0\xe5\xa5\xbd\xe4\xb8\x96\xe7\x95\x8c')
111+
p = proto.TCyBinaryProtocol(b)
112+
assert_equal(u("你好世界"), p.read_val(TType.STRING))
92113

93114

94115
def test_write_message_begin():
95-
b = BytesIO()
96-
proto.TCyBinaryProtocol(b).write_message_begin('test', TType.STRING, 1)
116+
b = TMemoryBuffer()
117+
p = proto.TCyBinaryProtocol(b)
118+
p.write_message_begin('test', TType.STRING, 1)
119+
p.write_message_end()
97120
assert_equal("80 01 00 0b 00 00 00 04 74 65 73 74 00 00 00 01",
98121
hexlify(b.getvalue()))
99122

100123

101124
def test_read_message_begin():
102-
b = BytesIO(b'\x80\x01\x00\x0b\x00\x00\x00\x04test\x00\x00\x00\x01')
125+
b = TMemoryBuffer(b'\x80\x01\x00\x0b\x00\x00\x00\x04test\x00\x00\x00\x01')
103126
res = proto.TCyBinaryProtocol(b).read_message_begin()
104127
assert_equal(res, ("test", TType.STRING, 1))
105128

106129

107130
def test_write_struct():
108-
b = BytesIO()
131+
b = TMemoryBuffer()
109132
item = TItem(id=123, phones=['123456', 'abcdef'])
110-
proto.TCyBinaryProtocol(b).write_struct(item)
133+
p = proto.TCyBinaryProtocol(b)
134+
p.write_struct(item)
135+
p.write_message_end()
111136
assert_equal("08 00 01 00 00 00 7b 0f 00 02 0b 00 00 00 02 00 00 00 "
112137
"06 31 32 33 34 35 36 00 00 00 06 61 62 63 64 65 66 00",
113138
hexlify(b.getvalue()))
114139

115140

116141
def test_read_struct():
117-
b = BytesIO(b'\x08\x00\x01\x00\x00\x00{\x0f\x00\x02\x0b\x00\x00\x00'
118-
b'\x02\x00\x00\x00\x06123456\x00\x00\x00\x06abcdef\x00')
142+
b = TMemoryBuffer(b'\x08\x00\x01\x00\x00\x00{\x0f\x00\x02\x0b\x00\x00\x00'
143+
b'\x02\x00\x00\x00\x06123456\x00\x00\x00\x06abcdef\x00')
119144
_item = TItem(id=123, phones=['123456', 'abcdef'])
120145
_item2 = TItem()
121146
proto.TCyBinaryProtocol(b).read_struct(_item2)

‎thriftpy/protocol/__init__.py

+1-1
Original file line numberDiff line numberDiff line change
@@ -5,4 +5,4 @@
55

66

77
from .binary import TBinaryProtocol, TBinaryProtocolFactory
8-
from .cybinary import TCyBinaryProtocol, TCyBinaryProtocolFactory
8+
from .cybin import TCyBinaryProtocol, TCyBinaryProtocolFactory

‎thriftpy/protocol/cybin/binbuf.pxi

+182
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,182 @@
1+
from libc.stdlib cimport malloc, free
2+
from libc.string cimport memcpy
3+
from libc.stdint cimport *
4+
5+
6+
cdef class Buffer(object):
7+
cdef byte *buf
8+
cdef int cur, buf_size, data_size
9+
10+
def __init__(self, buf_size):
11+
self.buf = <byte*>malloc(buf_size)
12+
self.buf_size = buf_size
13+
self.cur = 0
14+
self.data_size = 0
15+
16+
def __dealloc__(self):
17+
if self.buf != NULL:
18+
free(self.buf)
19+
20+
cdef void move_to_start(self):
21+
if self.cur != 0 and self.data_size > 0:
22+
memcpy(self.buf, self.buf + self.cur, self.data_size)
23+
self.cur = 0
24+
25+
cdef void clean(self):
26+
self.cur = 0
27+
self.data_size = 0
28+
29+
30+
class BufferError(Exception):
31+
pass
32+
33+
34+
cdef class BinaryRW(object):
35+
'''binary reader/writer'''
36+
37+
DEF DEFAULT_BUFFER = 4096
38+
cdef object trans
39+
cdef Buffer rbuf, wbuf
40+
41+
def __init__(self, trans, int buf_size=DEFAULT_BUFFER):
42+
self.trans = trans
43+
self.rbuf = Buffer(buf_size)
44+
self.wbuf = Buffer(buf_size)
45+
46+
cdef ensure_rbuf(self, int size):
47+
if size > self.rbuf.buf_size:
48+
raise BufferError('reader buffer out of cap')
49+
50+
cdef int cap
51+
cdef bytes new_data
52+
if self.rbuf.data_size == 0:
53+
self.rbuf.cur = 0
54+
self.rbuf.data_size = 0
55+
if self.rbuf.data_size < size:
56+
cap = self.rbuf.buf_size - self.rbuf.data_size
57+
if cap < 256:
58+
self.rbuf.move_to_start()
59+
new_data = self.trans._read(cap)
60+
memcpy(self.rbuf.buf + self.rbuf.cur + self.rbuf.data_size,
61+
<byte*>new_data, len(new_data))
62+
self.rbuf.data_size += len(new_data)
63+
64+
cdef read_byte(self, byte *ret):
65+
self.ensure_rbuf(1)
66+
ret[0] = (self.rbuf.buf + self.rbuf.cur)[0]
67+
self.rbuf.cur += 1
68+
self.rbuf.data_size -= 1
69+
70+
cdef read_int16(self, int16_t *ret):
71+
self.ensure_rbuf(2)
72+
ret[0] = be16toh((<int16_t*>(self.rbuf.buf + self.rbuf.cur))[0])
73+
self.rbuf.cur += 2
74+
self.rbuf.data_size -= 2
75+
76+
cdef read_int32(self, int32_t *ret):
77+
self.ensure_rbuf(4)
78+
ret[0] = be32toh((<int32_t*>(self.rbuf.buf + self.rbuf.cur))[0])
79+
self.rbuf.cur += 4
80+
self.rbuf.data_size -= 4
81+
82+
cdef read_int64(self, int64_t *ret):
83+
self.ensure_rbuf(8)
84+
ret[0] = be64toh((<int64_t*>(self.rbuf.buf + self.rbuf.cur))[0])
85+
self.rbuf.cur += 8
86+
self.rbuf.data_size -= 8
87+
88+
cdef read_double(self, double *ret):
89+
self.ensure_rbuf(sizeof(double))
90+
cdef int64_t n = be64toh((<int64_t*>(self.rbuf.buf + self.rbuf.cur))[0])
91+
ret[0] = (<double*>(&n))[0]
92+
self.rbuf.cur += sizeof(double)
93+
self.rbuf.data_size -= sizeof(double)
94+
95+
cdef read_string(self):
96+
self.ensure_rbuf(4)
97+
cdef int32_t str_size = be32toh((<int32_t*>(self.rbuf.buf + self.rbuf.cur))[0])
98+
self.rbuf.cur += 4
99+
self.rbuf.data_size -= 4
100+
101+
if str_size == 0:
102+
return ''
103+
self.ensure_rbuf(str_size)
104+
ret_str = (self.rbuf.buf + self.rbuf.cur)[:str_size].decode('utf8')
105+
self.rbuf.cur += str_size
106+
self.rbuf.data_size -= str_size
107+
return ret_str
108+
109+
cdef read_bytes(self, int size):
110+
self.ensure_rbuf(size)
111+
cdef bytes ret_bs = (self.rbuf.buf + self.rbuf.cur)[:size]
112+
self.rbuf.cur += size
113+
self.rbuf.data_size -= size
114+
return ret_bs
115+
116+
cdef read_skip(self, int size):
117+
self.ensure_rbuf(size)
118+
self.rbuf.cur += size
119+
self.rbuf.data_size -= size
120+
121+
cdef ensure_wbuf(self, int size):
122+
cdef int cap = self.rbuf.buf_size - self.rbuf.data_size
123+
124+
if cap < size:
125+
if size > self.wbuf.buf_size:
126+
raise BufferError('writer buffer out of cap')
127+
self.write_flush()
128+
129+
cdef write_byte(self, byte n):
130+
self.ensure_wbuf(1)
131+
(self.wbuf.buf + self.wbuf.data_size)[0] = n
132+
self.wbuf.data_size += 1
133+
134+
cdef write_int16(self, int16_t n):
135+
self.ensure_wbuf(2)
136+
(<int16_t*>(self.wbuf.buf + self.wbuf.data_size))[0] = htobe16(n)
137+
self.wbuf.data_size += 2
138+
139+
cdef write_int32(self, int32_t n):
140+
self.ensure_wbuf(4)
141+
(<int32_t*>(self.wbuf.buf + self.wbuf.data_size))[0] = htobe32(n)
142+
self.wbuf.data_size += 4
143+
144+
cdef write_int64(self, int64_t n):
145+
self.ensure_wbuf(8)
146+
(<int64_t*>(self.wbuf.buf + self.wbuf.data_size))[0] = htobe64(n)
147+
self.wbuf.data_size += 8
148+
149+
cdef write_double(self, double n):
150+
self.ensure_wbuf(8)
151+
cdef int64_t *n64 = <int64_t*>(&n)
152+
(<int64_t*>(self.wbuf.buf + self.wbuf.data_size))[0] = htobe64(n64[0])
153+
self.wbuf.data_size += 8
154+
155+
cdef write_string(self, s):
156+
cdef bytes bs
157+
if isinstance(s, bytes):
158+
bs = s
159+
else:
160+
bs = s.encode('utf8')
161+
162+
cdef int size = len(bs)
163+
self.ensure_wbuf(4)
164+
(<int32_t*>(self.wbuf.buf + self.wbuf.data_size))[0] = htobe32(size)
165+
self.wbuf.data_size += 4
166+
167+
self.ensure_wbuf(size)
168+
memcpy(self.wbuf.buf + self.wbuf.data_size, <byte*>bs, size)
169+
self.wbuf.data_size += size
170+
171+
cdef write_bytes(self, bytes bs):
172+
cdef int size = len(bs)
173+
self.ensure_wbuf(size)
174+
memcpy(self.wbuf.buf + self.wbuf.data_size, <byte*>bs, size)
175+
self.wbuf.data_size += size
176+
177+
cdef write_flush(self):
178+
cdef bytes data
179+
if self.wbuf.data_size > 0:
180+
data = self.wbuf.buf[:self.wbuf.data_size]
181+
self.trans.write(data)
182+
self.wbuf.clean()

0 commit comments

Comments
 (0)