blob: dc2b095de55f68d8864b9f65a0152cb55499ecd4 [file] [log] [blame]
Mark Slee89e2bb82007-03-01 00:20:36 +00001#
David Reissea2cba82009-03-30 21:35:00 +00002# Licensed to the Apache Software Foundation (ASF) under one
3# or more contributor license agreements. See the NOTICE file
4# distributed with this work for additional information
5# regarding copyright ownership. The ASF licenses this file
6# to you under the Apache License, Version 2.0 (the
7# "License"); you may not use this file except in compliance
8# with the License. You may obtain a copy of the License at
9#
10# http://www.apache.org/licenses/LICENSE-2.0
11#
12# Unless required by applicable law or agreed to in writing,
13# software distributed under the License is distributed on an
14# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15# KIND, either express or implied. See the License for the
16# specific language governing permissions and limitations
17# under the License.
18#
Mark Slee89e2bb82007-03-01 00:20:36 +000019
Mark Sleee74306a2007-02-21 05:38:12 +000020from thrift.Thrift import *
Mark Sleecde2b612006-09-03 21:13:07 +000021
Mark Slee92195ae2007-02-21 05:16:30 +000022
Bryan Duxbury69720412012-01-03 17:32:30 +000023class TProtocolException(TException):
Mark Slee92195ae2007-02-21 05:16:30 +000024 """Custom Protocol Exception class"""
25
26 UNKNOWN = 0
27 INVALID_DATA = 1
28 NEGATIVE_SIZE = 2
29 SIZE_LIMIT = 3
Mark Slee808454e2007-06-20 21:51:57 +000030 BAD_VERSION = 4
Mark Slee92195ae2007-02-21 05:16:30 +000031
32 def __init__(self, type=UNKNOWN, message=None):
33 TException.__init__(self, message)
34 self.type = type
35
Mark Sleecde2b612006-09-03 21:13:07 +000036
Bryan Duxbury69720412012-01-03 17:32:30 +000037class TProtocolBase:
Mark Sleecde2b612006-09-03 21:13:07 +000038 """Base class for Thrift protocol driver."""
39
Aditya Agarwal5c468192007-02-06 01:14:33 +000040 def __init__(self, trans):
41 self.trans = trans
Mark Slee4ac459f2006-10-25 21:39:01 +000042
Roger Meierad8154a2012-12-18 21:02:16 +010043 def writeMessageBegin(self, name, ttype, seqid):
Mark Sleecde2b612006-09-03 21:13:07 +000044 pass
45
Mark Slee4ac459f2006-10-25 21:39:01 +000046 def writeMessageEnd(self):
Mark Sleecde2b612006-09-03 21:13:07 +000047 pass
48
Mark Slee4ac459f2006-10-25 21:39:01 +000049 def writeStructBegin(self, name):
Mark Sleecde2b612006-09-03 21:13:07 +000050 pass
51
Mark Slee4ac459f2006-10-25 21:39:01 +000052 def writeStructEnd(self):
Mark Sleecde2b612006-09-03 21:13:07 +000053 pass
54
Roger Meierad8154a2012-12-18 21:02:16 +010055 def writeFieldBegin(self, name, ttype, fid):
Mark Sleecde2b612006-09-03 21:13:07 +000056 pass
57
Mark Slee4ac459f2006-10-25 21:39:01 +000058 def writeFieldEnd(self):
Mark Sleecde2b612006-09-03 21:13:07 +000059 pass
60
Mark Slee4ac459f2006-10-25 21:39:01 +000061 def writeFieldStop(self):
Mark Sleecde2b612006-09-03 21:13:07 +000062 pass
63
Mark Slee4ac459f2006-10-25 21:39:01 +000064 def writeMapBegin(self, ktype, vtype, size):
Mark Sleecde2b612006-09-03 21:13:07 +000065 pass
66
Mark Slee4ac459f2006-10-25 21:39:01 +000067 def writeMapEnd(self):
Mark Sleecde2b612006-09-03 21:13:07 +000068 pass
69
Mark Slee4ac459f2006-10-25 21:39:01 +000070 def writeListBegin(self, etype, size):
Mark Sleecde2b612006-09-03 21:13:07 +000071 pass
72
Mark Slee4ac459f2006-10-25 21:39:01 +000073 def writeListEnd(self):
Mark Sleecde2b612006-09-03 21:13:07 +000074 pass
75
Mark Slee4ac459f2006-10-25 21:39:01 +000076 def writeSetBegin(self, etype, size):
Mark Sleecde2b612006-09-03 21:13:07 +000077 pass
78
Mark Slee4ac459f2006-10-25 21:39:01 +000079 def writeSetEnd(self):
Mark Sleecde2b612006-09-03 21:13:07 +000080 pass
81
Roger Meierad8154a2012-12-18 21:02:16 +010082 def writeBool(self, bool_val):
Mark Sleecde2b612006-09-03 21:13:07 +000083 pass
84
Mark Slee4ac459f2006-10-25 21:39:01 +000085 def writeByte(self, byte):
Mark Sleecde2b612006-09-03 21:13:07 +000086 pass
87
Mark Slee4ac459f2006-10-25 21:39:01 +000088 def writeI16(self, i16):
Mark Sleecde2b612006-09-03 21:13:07 +000089 pass
90
Mark Slee4ac459f2006-10-25 21:39:01 +000091 def writeI32(self, i32):
Mark Sleecde2b612006-09-03 21:13:07 +000092 pass
93
Mark Slee4ac459f2006-10-25 21:39:01 +000094 def writeI64(self, i64):
Mark Sleecde2b612006-09-03 21:13:07 +000095 pass
96
Mark Slee4ac459f2006-10-25 21:39:01 +000097 def writeDouble(self, dub):
Mark Sleec98d0502006-09-06 02:42:25 +000098 pass
99
Roger Meierad8154a2012-12-18 21:02:16 +0100100 def writeString(self, str_val):
Mark Sleecde2b612006-09-03 21:13:07 +0000101 pass
102
Mark Slee4ac459f2006-10-25 21:39:01 +0000103 def readMessageBegin(self):
Mark Sleecde2b612006-09-03 21:13:07 +0000104 pass
105
Mark Slee4ac459f2006-10-25 21:39:01 +0000106 def readMessageEnd(self):
Mark Sleecde2b612006-09-03 21:13:07 +0000107 pass
108
Mark Slee4ac459f2006-10-25 21:39:01 +0000109 def readStructBegin(self):
Mark Sleecde2b612006-09-03 21:13:07 +0000110 pass
111
Mark Slee4ac459f2006-10-25 21:39:01 +0000112 def readStructEnd(self):
Mark Sleecde2b612006-09-03 21:13:07 +0000113 pass
114
Mark Slee4ac459f2006-10-25 21:39:01 +0000115 def readFieldBegin(self):
Mark Sleecde2b612006-09-03 21:13:07 +0000116 pass
117
Mark Slee4ac459f2006-10-25 21:39:01 +0000118 def readFieldEnd(self):
Mark Sleecde2b612006-09-03 21:13:07 +0000119 pass
120
Mark Slee4ac459f2006-10-25 21:39:01 +0000121 def readMapBegin(self):
Mark Sleecde2b612006-09-03 21:13:07 +0000122 pass
123
Mark Slee4ac459f2006-10-25 21:39:01 +0000124 def readMapEnd(self):
Mark Sleecde2b612006-09-03 21:13:07 +0000125 pass
126
Mark Slee4ac459f2006-10-25 21:39:01 +0000127 def readListBegin(self):
Mark Sleecde2b612006-09-03 21:13:07 +0000128 pass
129
Mark Slee4ac459f2006-10-25 21:39:01 +0000130 def readListEnd(self):
Mark Sleecde2b612006-09-03 21:13:07 +0000131 pass
132
Mark Slee4ac459f2006-10-25 21:39:01 +0000133 def readSetBegin(self):
Mark Sleecde2b612006-09-03 21:13:07 +0000134 pass
135
Mark Slee4ac459f2006-10-25 21:39:01 +0000136 def readSetEnd(self):
Mark Sleecde2b612006-09-03 21:13:07 +0000137 pass
138
Mark Slee4ac459f2006-10-25 21:39:01 +0000139 def readBool(self):
Mark Sleecde2b612006-09-03 21:13:07 +0000140 pass
141
Mark Slee4ac459f2006-10-25 21:39:01 +0000142 def readByte(self):
Mark Sleecde2b612006-09-03 21:13:07 +0000143 pass
144
Mark Slee4ac459f2006-10-25 21:39:01 +0000145 def readI16(self):
Mark Sleecde2b612006-09-03 21:13:07 +0000146 pass
147
Mark Slee4ac459f2006-10-25 21:39:01 +0000148 def readI32(self):
Mark Sleecde2b612006-09-03 21:13:07 +0000149 pass
150
Mark Slee4ac459f2006-10-25 21:39:01 +0000151 def readI64(self):
Mark Sleecde2b612006-09-03 21:13:07 +0000152 pass
153
Mark Slee4ac459f2006-10-25 21:39:01 +0000154 def readDouble(self):
Mark Sleec98d0502006-09-06 02:42:25 +0000155 pass
156
Mark Slee4ac459f2006-10-25 21:39:01 +0000157 def readString(self):
Mark Sleecde2b612006-09-03 21:13:07 +0000158 pass
159
Roger Meierad8154a2012-12-18 21:02:16 +0100160 def skip(self, ttype):
161 if ttype == TType.STOP:
Mark Sleecde2b612006-09-03 21:13:07 +0000162 return
Roger Meierad8154a2012-12-18 21:02:16 +0100163 elif ttype == TType.BOOL:
Mark Slee4ac459f2006-10-25 21:39:01 +0000164 self.readBool()
Roger Meierad8154a2012-12-18 21:02:16 +0100165 elif ttype == TType.BYTE:
Mark Slee4ac459f2006-10-25 21:39:01 +0000166 self.readByte()
Roger Meierad8154a2012-12-18 21:02:16 +0100167 elif ttype == TType.I16:
Mark Slee4ac459f2006-10-25 21:39:01 +0000168 self.readI16()
Roger Meierad8154a2012-12-18 21:02:16 +0100169 elif ttype == TType.I32:
Mark Slee4ac459f2006-10-25 21:39:01 +0000170 self.readI32()
Roger Meierad8154a2012-12-18 21:02:16 +0100171 elif ttype == TType.I64:
Mark Slee4ac459f2006-10-25 21:39:01 +0000172 self.readI64()
Roger Meierad8154a2012-12-18 21:02:16 +0100173 elif ttype == TType.DOUBLE:
Mark Slee4ac459f2006-10-25 21:39:01 +0000174 self.readDouble()
Roger Meierad8154a2012-12-18 21:02:16 +0100175 elif ttype == TType.STRING:
Mark Slee4ac459f2006-10-25 21:39:01 +0000176 self.readString()
Roger Meierad8154a2012-12-18 21:02:16 +0100177 elif ttype == TType.STRUCT:
Mark Slee4ac459f2006-10-25 21:39:01 +0000178 name = self.readStructBegin()
Mark Sleecde2b612006-09-03 21:13:07 +0000179 while True:
Roger Meierad8154a2012-12-18 21:02:16 +0100180 (name, ttype, id) = self.readFieldBegin()
181 if ttype == TType.STOP:
Mark Sleecde2b612006-09-03 21:13:07 +0000182 break
Roger Meierad8154a2012-12-18 21:02:16 +0100183 self.skip(ttype)
Mark Slee4ac459f2006-10-25 21:39:01 +0000184 self.readFieldEnd()
185 self.readStructEnd()
Roger Meierad8154a2012-12-18 21:02:16 +0100186 elif ttype == TType.MAP:
Mark Slee4ac459f2006-10-25 21:39:01 +0000187 (ktype, vtype, size) = self.readMapBegin()
Roger Meierad8154a2012-12-18 21:02:16 +0100188 for i in xrange(size):
Mark Slee4ac459f2006-10-25 21:39:01 +0000189 self.skip(ktype)
190 self.skip(vtype)
191 self.readMapEnd()
Roger Meierad8154a2012-12-18 21:02:16 +0100192 elif ttype == TType.SET:
Mark Slee4ac459f2006-10-25 21:39:01 +0000193 (etype, size) = self.readSetBegin()
Roger Meierad8154a2012-12-18 21:02:16 +0100194 for i in xrange(size):
Mark Slee4ac459f2006-10-25 21:39:01 +0000195 self.skip(etype)
196 self.readSetEnd()
Roger Meierad8154a2012-12-18 21:02:16 +0100197 elif ttype == TType.LIST:
Mark Slee4ac459f2006-10-25 21:39:01 +0000198 (etype, size) = self.readListBegin()
Roger Meierad8154a2012-12-18 21:02:16 +0100199 for i in xrange(size):
Mark Slee4ac459f2006-10-25 21:39:01 +0000200 self.skip(etype)
201 self.readListEnd()
Mark Sleecde2b612006-09-03 21:13:07 +0000202
Bryan Duxbury69720412012-01-03 17:32:30 +0000203 # tuple of: ( 'reader method' name, is_container bool, 'writer_method' name )
Roger Meierf4eec7a2011-09-11 18:16:21 +0000204 _TTYPE_HANDLERS = (
Bryan Duxbury69720412012-01-03 17:32:30 +0000205 (None, None, False), # 0 TType.STOP
206 (None, None, False), # 1 TType.VOID # TODO: handle void?
207 ('readBool', 'writeBool', False), # 2 TType.BOOL
208 ('readByte', 'writeByte', False), # 3 TType.BYTE and I08
209 ('readDouble', 'writeDouble', False), # 4 TType.DOUBLE
210 (None, None, False), # 5 undefined
211 ('readI16', 'writeI16', False), # 6 TType.I16
212 (None, None, False), # 7 undefined
213 ('readI32', 'writeI32', False), # 8 TType.I32
214 (None, None, False), # 9 undefined
215 ('readI64', 'writeI64', False), # 10 TType.I64
216 ('readString', 'writeString', False), # 11 TType.STRING and UTF7
217 ('readContainerStruct', 'writeContainerStruct', True), # 12 *.STRUCT
218 ('readContainerMap', 'writeContainerMap', True), # 13 TType.MAP
219 ('readContainerSet', 'writeContainerSet', True), # 14 TType.SET
220 ('readContainerList', 'writeContainerList', True), # 15 TType.LIST
221 (None, None, False), # 16 TType.UTF8 # TODO: handle utf8 types?
222 (None, None, False) # 17 TType.UTF16 # TODO: handle utf16 types?
Roger Meierf4eec7a2011-09-11 18:16:21 +0000223 )
224
225 def readFieldByTType(self, ttype, spec):
226 try:
227 (r_handler, w_handler, is_container) = self._TTYPE_HANDLERS[ttype]
228 except IndexError:
229 raise TProtocolException(type=TProtocolException.INVALID_DATA,
230 message='Invalid field type %d' % (ttype))
231 if r_handler is None:
232 raise TProtocolException(type=TProtocolException.INVALID_DATA,
233 message='Invalid field type %d' % (ttype))
234 reader = getattr(self, r_handler)
235 if not is_container:
236 return reader()
237 return reader(spec)
238
239 def readContainerList(self, spec):
240 results = []
241 ttype, tspec = spec[0], spec[1]
242 r_handler = self._TTYPE_HANDLERS[ttype][0]
243 reader = getattr(self, r_handler)
244 (list_type, list_len) = self.readListBegin()
245 if tspec is None:
246 # list values are simple types
247 for idx in xrange(list_len):
248 results.append(reader())
249 else:
250 # this is like an inlined readFieldByTType
251 container_reader = self._TTYPE_HANDLERS[list_type][0]
252 val_reader = getattr(self, container_reader)
253 for idx in xrange(list_len):
254 val = val_reader(tspec)
255 results.append(val)
256 self.readListEnd()
257 return results
258
259 def readContainerSet(self, spec):
260 results = set()
261 ttype, tspec = spec[0], spec[1]
262 r_handler = self._TTYPE_HANDLERS[ttype][0]
263 reader = getattr(self, r_handler)
264 (set_type, set_len) = self.readSetBegin()
265 if tspec is None:
266 # set members are simple types
267 for idx in xrange(set_len):
268 results.add(reader())
269 else:
270 container_reader = self._TTYPE_HANDLERS[set_type][0]
271 val_reader = getattr(self, container_reader)
272 for idx in xrange(set_len):
Bryan Duxbury69720412012-01-03 17:32:30 +0000273 results.add(val_reader(tspec))
Roger Meierf4eec7a2011-09-11 18:16:21 +0000274 self.readSetEnd()
275 return results
276
277 def readContainerStruct(self, spec):
278 (obj_class, obj_spec) = spec
279 obj = obj_class()
280 obj.read(self)
281 return obj
Bryan Duxbury69720412012-01-03 17:32:30 +0000282
Roger Meierf4eec7a2011-09-11 18:16:21 +0000283 def readContainerMap(self, spec):
284 results = dict()
285 key_ttype, key_spec = spec[0], spec[1]
286 val_ttype, val_spec = spec[2], spec[3]
287 (map_ktype, map_vtype, map_len) = self.readMapBegin()
Bryan Duxbury69720412012-01-03 17:32:30 +0000288 # TODO: compare types we just decoded with thrift_spec and
289 # abort/skip if types disagree
Roger Meierf4eec7a2011-09-11 18:16:21 +0000290 key_reader = getattr(self, self._TTYPE_HANDLERS[key_ttype][0])
291 val_reader = getattr(self, self._TTYPE_HANDLERS[val_ttype][0])
292 # list values are simple types
293 for idx in xrange(map_len):
294 if key_spec is None:
295 k_val = key_reader()
296 else:
297 k_val = self.readFieldByTType(key_ttype, key_spec)
298 if val_spec is None:
299 v_val = val_reader()
300 else:
301 v_val = self.readFieldByTType(val_ttype, val_spec)
Bryan Duxbury69720412012-01-03 17:32:30 +0000302 # this raises a TypeError with unhashable keys types
303 # i.e. this fails: d=dict(); d[[0,1]] = 2
Roger Meierf4eec7a2011-09-11 18:16:21 +0000304 results[k_val] = v_val
305 self.readMapEnd()
306 return results
307
308 def readStruct(self, obj, thrift_spec):
309 self.readStructBegin()
310 while True:
311 (fname, ftype, fid) = self.readFieldBegin()
312 if ftype == TType.STOP:
313 break
314 try:
315 field = thrift_spec[fid]
316 except IndexError:
317 self.skip(ftype)
318 else:
319 if field is not None and ftype == field[1]:
320 fname = field[2]
321 fspec = field[3]
322 val = self.readFieldByTType(ftype, fspec)
323 setattr(obj, fname, val)
324 else:
325 self.skip(ftype)
326 self.readFieldEnd()
327 self.readStructEnd()
328
329 def writeContainerStruct(self, val, spec):
330 val.write(self)
331
332 def writeContainerList(self, val, spec):
333 self.writeListBegin(spec[0], len(val))
Bryan Duxbury69720412012-01-03 17:32:30 +0000334 r_handler, w_handler, is_container = self._TTYPE_HANDLERS[spec[0]]
Roger Meierf4eec7a2011-09-11 18:16:21 +0000335 e_writer = getattr(self, w_handler)
336 if not is_container:
337 for elem in val:
338 e_writer(elem)
339 else:
340 for elem in val:
341 e_writer(elem, spec[1])
342 self.writeListEnd()
343
344 def writeContainerSet(self, val, spec):
345 self.writeSetBegin(spec[0], len(val))
346 r_handler, w_handler, is_container = self._TTYPE_HANDLERS[spec[0]]
347 e_writer = getattr(self, w_handler)
348 if not is_container:
349 for elem in val:
350 e_writer(elem)
351 else:
352 for elem in val:
353 e_writer(elem, spec[1])
354 self.writeSetEnd()
355
356 def writeContainerMap(self, val, spec):
357 k_type = spec[0]
358 v_type = spec[2]
359 ignore, ktype_name, k_is_container = self._TTYPE_HANDLERS[k_type]
360 ignore, vtype_name, v_is_container = self._TTYPE_HANDLERS[v_type]
361 k_writer = getattr(self, ktype_name)
362 v_writer = getattr(self, vtype_name)
363 self.writeMapBegin(k_type, v_type, len(val))
364 for m_key, m_val in val.iteritems():
365 if not k_is_container:
366 k_writer(m_key)
367 else:
368 k_writer(m_key, spec[1])
369 if not v_is_container:
370 v_writer(m_val)
371 else:
372 v_writer(m_val, spec[3])
373 self.writeMapEnd()
374
375 def writeStruct(self, obj, thrift_spec):
376 self.writeStructBegin(obj.__class__.__name__)
377 for field in thrift_spec:
378 if field is None:
379 continue
380 fname = field[2]
381 val = getattr(obj, fname)
382 if val is None:
383 # skip writing out unset fields
384 continue
385 fid = field[0]
386 ftype = field[1]
387 fspec = field[3]
388 # get the writer method for this value
389 self.writeFieldBegin(fname, ftype, fid)
390 self.writeFieldByTType(ftype, val, fspec)
391 self.writeFieldEnd()
392 self.writeFieldStop()
393 self.writeStructEnd()
394
395 def writeFieldByTType(self, ttype, val, spec):
396 r_handler, w_handler, is_container = self._TTYPE_HANDLERS[ttype]
397 writer = getattr(self, w_handler)
398 if is_container:
399 writer(val, spec)
400 else:
401 writer(val)
402
Bryan Duxbury69720412012-01-03 17:32:30 +0000403
Mark Slee4ac459f2006-10-25 21:39:01 +0000404class TProtocolFactory:
Aditya Agarwal5c468192007-02-06 01:14:33 +0000405 def getProtocol(self, trans):
Mark Slee4ac459f2006-10-25 21:39:01 +0000406 pass