encoder.py 27 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825
  1. # Protocol Buffers - Google's data interchange format
  2. # Copyright 2008 Google Inc. All rights reserved.
  3. # https://developers.google.com/protocol-buffers/
  4. #
  5. # Redistribution and use in source and binary forms, with or without
  6. # modification, are permitted provided that the following conditions are
  7. # met:
  8. #
  9. # * Redistributions of source code must retain the above copyright
  10. # notice, this list of conditions and the following disclaimer.
  11. # * Redistributions in binary form must reproduce the above
  12. # copyright notice, this list of conditions and the following disclaimer
  13. # in the documentation and/or other materials provided with the
  14. # distribution.
  15. # * Neither the name of Google Inc. nor the names of its
  16. # contributors may be used to endorse or promote products derived from
  17. # this software without specific prior written permission.
  18. #
  19. # THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
  20. # "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
  21. # LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
  22. # A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
  23. # OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
  24. # SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
  25. # LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
  26. # DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
  27. # THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
  28. # (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
  29. # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
  30. # Copyright 2009 Google Inc. All Rights Reserved.
  31. """Code for encoding protocol message primitives.
  32. Contains the logic for encoding every logical protocol field type
  33. into one of the 5 physical wire types.
  34. This code is designed to push the Python interpreter's performance to the
  35. limits.
  36. The basic idea is that at startup time, for every field (i.e. every
  37. FieldDescriptor) we construct two functions: a "sizer" and an "encoder". The
  38. sizer takes a value of this field's type and computes its byte size. The
  39. encoder takes a writer function and a value. It encodes the value into byte
  40. strings and invokes the writer function to write those strings. Typically the
  41. writer function is the write() method of a BytesIO.
  42. We try to do as much work as possible when constructing the writer and the
  43. sizer rather than when calling them. In particular:
  44. * We copy any needed global functions to local variables, so that we do not need
  45. to do costly global table lookups at runtime.
  46. * Similarly, we try to do any attribute lookups at startup time if possible.
  47. * Every field's tag is encoded to bytes at startup, since it can't change at
  48. runtime.
  49. * Whatever component of the field size we can compute at startup, we do.
  50. * We *avoid* sharing code if doing so would make the code slower and not sharing
  51. does not burden us too much. For example, encoders for repeated fields do
  52. not just call the encoders for singular fields in a loop because this would
  53. add an extra function call overhead for every loop iteration; instead, we
  54. manually inline the single-value encoder into the loop.
  55. * If a Python function lacks a return statement, Python actually generates
  56. instructions to pop the result of the last statement off the stack, push
  57. None onto the stack, and then return that. If we really don't care what
  58. value is returned, then we can save two instructions by returning the
  59. result of the last statement. It looks funny but it helps.
  60. * We assume that type and bounds checking has happened at a higher level.
  61. """
  62. __author__ = 'kenton@google.com (Kenton Varda)'
  63. import struct
  64. import six
  65. from google.protobuf.internal import wire_format
  66. # This will overflow and thus become IEEE-754 "infinity". We would use
  67. # "float('inf')" but it doesn't work on Windows pre-Python-2.6.
  68. _POS_INF = 1e10000
  69. _NEG_INF = -_POS_INF
  70. def _VarintSize(value):
  71. """Compute the size of a varint value."""
  72. if value <= 0x7f: return 1
  73. if value <= 0x3fff: return 2
  74. if value <= 0x1fffff: return 3
  75. if value <= 0xfffffff: return 4
  76. if value <= 0x7ffffffff: return 5
  77. if value <= 0x3ffffffffff: return 6
  78. if value <= 0x1ffffffffffff: return 7
  79. if value <= 0xffffffffffffff: return 8
  80. if value <= 0x7fffffffffffffff: return 9
  81. return 10
  82. def _SignedVarintSize(value):
  83. """Compute the size of a signed varint value."""
  84. if value < 0: return 10
  85. if value <= 0x7f: return 1
  86. if value <= 0x3fff: return 2
  87. if value <= 0x1fffff: return 3
  88. if value <= 0xfffffff: return 4
  89. if value <= 0x7ffffffff: return 5
  90. if value <= 0x3ffffffffff: return 6
  91. if value <= 0x1ffffffffffff: return 7
  92. if value <= 0xffffffffffffff: return 8
  93. if value <= 0x7fffffffffffffff: return 9
  94. return 10
  95. def _TagSize(field_number):
  96. """Returns the number of bytes required to serialize a tag with this field
  97. number."""
  98. # Just pass in type 0, since the type won't affect the tag+type size.
  99. return _VarintSize(wire_format.PackTag(field_number, 0))
  100. # --------------------------------------------------------------------
  101. # In this section we define some generic sizers. Each of these functions
  102. # takes parameters specific to a particular field type, e.g. int32 or fixed64.
  103. # It returns another function which in turn takes parameters specific to a
  104. # particular field, e.g. the field number and whether it is repeated or packed.
  105. # Look at the next section to see how these are used.
  106. def _SimpleSizer(compute_value_size):
  107. """A sizer which uses the function compute_value_size to compute the size of
  108. each value. Typically compute_value_size is _VarintSize."""
  109. def SpecificSizer(field_number, is_repeated, is_packed):
  110. tag_size = _TagSize(field_number)
  111. if is_packed:
  112. local_VarintSize = _VarintSize
  113. def PackedFieldSize(value):
  114. result = 0
  115. for element in value:
  116. result += compute_value_size(element)
  117. return result + local_VarintSize(result) + tag_size
  118. return PackedFieldSize
  119. elif is_repeated:
  120. def RepeatedFieldSize(value):
  121. result = tag_size * len(value)
  122. for element in value:
  123. result += compute_value_size(element)
  124. return result
  125. return RepeatedFieldSize
  126. else:
  127. def FieldSize(value):
  128. return tag_size + compute_value_size(value)
  129. return FieldSize
  130. return SpecificSizer
  131. def _ModifiedSizer(compute_value_size, modify_value):
  132. """Like SimpleSizer, but modify_value is invoked on each value before it is
  133. passed to compute_value_size. modify_value is typically ZigZagEncode."""
  134. def SpecificSizer(field_number, is_repeated, is_packed):
  135. tag_size = _TagSize(field_number)
  136. if is_packed:
  137. local_VarintSize = _VarintSize
  138. def PackedFieldSize(value):
  139. result = 0
  140. for element in value:
  141. result += compute_value_size(modify_value(element))
  142. return result + local_VarintSize(result) + tag_size
  143. return PackedFieldSize
  144. elif is_repeated:
  145. def RepeatedFieldSize(value):
  146. result = tag_size * len(value)
  147. for element in value:
  148. result += compute_value_size(modify_value(element))
  149. return result
  150. return RepeatedFieldSize
  151. else:
  152. def FieldSize(value):
  153. return tag_size + compute_value_size(modify_value(value))
  154. return FieldSize
  155. return SpecificSizer
  156. def _FixedSizer(value_size):
  157. """Like _SimpleSizer except for a fixed-size field. The input is the size
  158. of one value."""
  159. def SpecificSizer(field_number, is_repeated, is_packed):
  160. tag_size = _TagSize(field_number)
  161. if is_packed:
  162. local_VarintSize = _VarintSize
  163. def PackedFieldSize(value):
  164. result = len(value) * value_size
  165. return result + local_VarintSize(result) + tag_size
  166. return PackedFieldSize
  167. elif is_repeated:
  168. element_size = value_size + tag_size
  169. def RepeatedFieldSize(value):
  170. return len(value) * element_size
  171. return RepeatedFieldSize
  172. else:
  173. field_size = value_size + tag_size
  174. def FieldSize(value):
  175. return field_size
  176. return FieldSize
  177. return SpecificSizer
  178. # ====================================================================
  179. # Here we declare a sizer constructor for each field type. Each "sizer
  180. # constructor" is a function that takes (field_number, is_repeated, is_packed)
  181. # as parameters and returns a sizer, which in turn takes a field value as
  182. # a parameter and returns its encoded size.
  183. Int32Sizer = Int64Sizer = EnumSizer = _SimpleSizer(_SignedVarintSize)
  184. UInt32Sizer = UInt64Sizer = _SimpleSizer(_VarintSize)
  185. SInt32Sizer = SInt64Sizer = _ModifiedSizer(
  186. _SignedVarintSize, wire_format.ZigZagEncode)
  187. Fixed32Sizer = SFixed32Sizer = FloatSizer = _FixedSizer(4)
  188. Fixed64Sizer = SFixed64Sizer = DoubleSizer = _FixedSizer(8)
  189. BoolSizer = _FixedSizer(1)
  190. def StringSizer(field_number, is_repeated, is_packed):
  191. """Returns a sizer for a string field."""
  192. tag_size = _TagSize(field_number)
  193. local_VarintSize = _VarintSize
  194. local_len = len
  195. assert not is_packed
  196. if is_repeated:
  197. def RepeatedFieldSize(value):
  198. result = tag_size * len(value)
  199. for element in value:
  200. l = local_len(element.encode('utf-8'))
  201. result += local_VarintSize(l) + l
  202. return result
  203. return RepeatedFieldSize
  204. else:
  205. def FieldSize(value):
  206. l = local_len(value.encode('utf-8'))
  207. return tag_size + local_VarintSize(l) + l
  208. return FieldSize
  209. def BytesSizer(field_number, is_repeated, is_packed):
  210. """Returns a sizer for a bytes field."""
  211. tag_size = _TagSize(field_number)
  212. local_VarintSize = _VarintSize
  213. local_len = len
  214. assert not is_packed
  215. if is_repeated:
  216. def RepeatedFieldSize(value):
  217. result = tag_size * len(value)
  218. for element in value:
  219. l = local_len(element)
  220. result += local_VarintSize(l) + l
  221. return result
  222. return RepeatedFieldSize
  223. else:
  224. def FieldSize(value):
  225. l = local_len(value)
  226. return tag_size + local_VarintSize(l) + l
  227. return FieldSize
  228. def GroupSizer(field_number, is_repeated, is_packed):
  229. """Returns a sizer for a group field."""
  230. tag_size = _TagSize(field_number) * 2
  231. assert not is_packed
  232. if is_repeated:
  233. def RepeatedFieldSize(value):
  234. result = tag_size * len(value)
  235. for element in value:
  236. result += element.ByteSize()
  237. return result
  238. return RepeatedFieldSize
  239. else:
  240. def FieldSize(value):
  241. return tag_size + value.ByteSize()
  242. return FieldSize
  243. def MessageSizer(field_number, is_repeated, is_packed):
  244. """Returns a sizer for a message field."""
  245. tag_size = _TagSize(field_number)
  246. local_VarintSize = _VarintSize
  247. assert not is_packed
  248. if is_repeated:
  249. def RepeatedFieldSize(value):
  250. result = tag_size * len(value)
  251. for element in value:
  252. l = element.ByteSize()
  253. result += local_VarintSize(l) + l
  254. return result
  255. return RepeatedFieldSize
  256. else:
  257. def FieldSize(value):
  258. l = value.ByteSize()
  259. return tag_size + local_VarintSize(l) + l
  260. return FieldSize
  261. # --------------------------------------------------------------------
  262. # MessageSet is special: it needs custom logic to compute its size properly.
  263. def MessageSetItemSizer(field_number):
  264. """Returns a sizer for extensions of MessageSet.
  265. The message set message looks like this:
  266. message MessageSet {
  267. repeated group Item = 1 {
  268. required int32 type_id = 2;
  269. required string message = 3;
  270. }
  271. }
  272. """
  273. static_size = (_TagSize(1) * 2 + _TagSize(2) + _VarintSize(field_number) +
  274. _TagSize(3))
  275. local_VarintSize = _VarintSize
  276. def FieldSize(value):
  277. l = value.ByteSize()
  278. return static_size + local_VarintSize(l) + l
  279. return FieldSize
  280. # --------------------------------------------------------------------
  281. # Map is special: it needs custom logic to compute its size properly.
  282. def MapSizer(field_descriptor):
  283. """Returns a sizer for a map field."""
  284. # Can't look at field_descriptor.message_type._concrete_class because it may
  285. # not have been initialized yet.
  286. message_type = field_descriptor.message_type
  287. message_sizer = MessageSizer(field_descriptor.number, False, False)
  288. def FieldSize(map_value):
  289. total = 0
  290. for key in map_value:
  291. value = map_value[key]
  292. # It's wasteful to create the messages and throw them away one second
  293. # later since we'll do the same for the actual encode. But there's not an
  294. # obvious way to avoid this within the current design without tons of code
  295. # duplication.
  296. entry_msg = message_type._concrete_class(key=key, value=value)
  297. total += message_sizer(entry_msg)
  298. return total
  299. return FieldSize
  300. # ====================================================================
  301. # Encoders!
  302. def _VarintEncoder():
  303. """Return an encoder for a basic varint value (does not include tag)."""
  304. def EncodeVarint(write, value):
  305. bits = value & 0x7f
  306. value >>= 7
  307. while value:
  308. write(six.int2byte(0x80|bits))
  309. bits = value & 0x7f
  310. value >>= 7
  311. return write(six.int2byte(bits))
  312. return EncodeVarint
  313. def _SignedVarintEncoder():
  314. """Return an encoder for a basic signed varint value (does not include
  315. tag)."""
  316. def EncodeSignedVarint(write, value):
  317. if value < 0:
  318. value += (1 << 64)
  319. bits = value & 0x7f
  320. value >>= 7
  321. while value:
  322. write(six.int2byte(0x80|bits))
  323. bits = value & 0x7f
  324. value >>= 7
  325. return write(six.int2byte(bits))
  326. return EncodeSignedVarint
  327. _EncodeVarint = _VarintEncoder()
  328. _EncodeSignedVarint = _SignedVarintEncoder()
  329. def _VarintBytes(value):
  330. """Encode the given integer as a varint and return the bytes. This is only
  331. called at startup time so it doesn't need to be fast."""
  332. pieces = []
  333. _EncodeVarint(pieces.append, value)
  334. return b"".join(pieces)
  335. def TagBytes(field_number, wire_type):
  336. """Encode the given tag and return the bytes. Only called at startup."""
  337. return _VarintBytes(wire_format.PackTag(field_number, wire_type))
  338. # --------------------------------------------------------------------
  339. # As with sizers (see above), we have a number of common encoder
  340. # implementations.
  341. def _SimpleEncoder(wire_type, encode_value, compute_value_size):
  342. """Return a constructor for an encoder for fields of a particular type.
  343. Args:
  344. wire_type: The field's wire type, for encoding tags.
  345. encode_value: A function which encodes an individual value, e.g.
  346. _EncodeVarint().
  347. compute_value_size: A function which computes the size of an individual
  348. value, e.g. _VarintSize().
  349. """
  350. def SpecificEncoder(field_number, is_repeated, is_packed):
  351. if is_packed:
  352. tag_bytes = TagBytes(field_number, wire_format.WIRETYPE_LENGTH_DELIMITED)
  353. local_EncodeVarint = _EncodeVarint
  354. def EncodePackedField(write, value):
  355. write(tag_bytes)
  356. size = 0
  357. for element in value:
  358. size += compute_value_size(element)
  359. local_EncodeVarint(write, size)
  360. for element in value:
  361. encode_value(write, element)
  362. return EncodePackedField
  363. elif is_repeated:
  364. tag_bytes = TagBytes(field_number, wire_type)
  365. def EncodeRepeatedField(write, value):
  366. for element in value:
  367. write(tag_bytes)
  368. encode_value(write, element)
  369. return EncodeRepeatedField
  370. else:
  371. tag_bytes = TagBytes(field_number, wire_type)
  372. def EncodeField(write, value):
  373. write(tag_bytes)
  374. return encode_value(write, value)
  375. return EncodeField
  376. return SpecificEncoder
  377. def _ModifiedEncoder(wire_type, encode_value, compute_value_size, modify_value):
  378. """Like SimpleEncoder but additionally invokes modify_value on every value
  379. before passing it to encode_value. Usually modify_value is ZigZagEncode."""
  380. def SpecificEncoder(field_number, is_repeated, is_packed):
  381. if is_packed:
  382. tag_bytes = TagBytes(field_number, wire_format.WIRETYPE_LENGTH_DELIMITED)
  383. local_EncodeVarint = _EncodeVarint
  384. def EncodePackedField(write, value):
  385. write(tag_bytes)
  386. size = 0
  387. for element in value:
  388. size += compute_value_size(modify_value(element))
  389. local_EncodeVarint(write, size)
  390. for element in value:
  391. encode_value(write, modify_value(element))
  392. return EncodePackedField
  393. elif is_repeated:
  394. tag_bytes = TagBytes(field_number, wire_type)
  395. def EncodeRepeatedField(write, value):
  396. for element in value:
  397. write(tag_bytes)
  398. encode_value(write, modify_value(element))
  399. return EncodeRepeatedField
  400. else:
  401. tag_bytes = TagBytes(field_number, wire_type)
  402. def EncodeField(write, value):
  403. write(tag_bytes)
  404. return encode_value(write, modify_value(value))
  405. return EncodeField
  406. return SpecificEncoder
  407. def _StructPackEncoder(wire_type, format):
  408. """Return a constructor for an encoder for a fixed-width field.
  409. Args:
  410. wire_type: The field's wire type, for encoding tags.
  411. format: The format string to pass to struct.pack().
  412. """
  413. value_size = struct.calcsize(format)
  414. def SpecificEncoder(field_number, is_repeated, is_packed):
  415. local_struct_pack = struct.pack
  416. if is_packed:
  417. tag_bytes = TagBytes(field_number, wire_format.WIRETYPE_LENGTH_DELIMITED)
  418. local_EncodeVarint = _EncodeVarint
  419. def EncodePackedField(write, value):
  420. write(tag_bytes)
  421. local_EncodeVarint(write, len(value) * value_size)
  422. for element in value:
  423. write(local_struct_pack(format, element))
  424. return EncodePackedField
  425. elif is_repeated:
  426. tag_bytes = TagBytes(field_number, wire_type)
  427. def EncodeRepeatedField(write, value):
  428. for element in value:
  429. write(tag_bytes)
  430. write(local_struct_pack(format, element))
  431. return EncodeRepeatedField
  432. else:
  433. tag_bytes = TagBytes(field_number, wire_type)
  434. def EncodeField(write, value):
  435. write(tag_bytes)
  436. return write(local_struct_pack(format, value))
  437. return EncodeField
  438. return SpecificEncoder
  439. def _FloatingPointEncoder(wire_type, format):
  440. """Return a constructor for an encoder for float fields.
  441. This is like StructPackEncoder, but catches errors that may be due to
  442. passing non-finite floating-point values to struct.pack, and makes a
  443. second attempt to encode those values.
  444. Args:
  445. wire_type: The field's wire type, for encoding tags.
  446. format: The format string to pass to struct.pack().
  447. """
  448. value_size = struct.calcsize(format)
  449. if value_size == 4:
  450. def EncodeNonFiniteOrRaise(write, value):
  451. # Remember that the serialized form uses little-endian byte order.
  452. if value == _POS_INF:
  453. write(b'\x00\x00\x80\x7F')
  454. elif value == _NEG_INF:
  455. write(b'\x00\x00\x80\xFF')
  456. elif value != value: # NaN
  457. write(b'\x00\x00\xC0\x7F')
  458. else:
  459. raise
  460. elif value_size == 8:
  461. def EncodeNonFiniteOrRaise(write, value):
  462. if value == _POS_INF:
  463. write(b'\x00\x00\x00\x00\x00\x00\xF0\x7F')
  464. elif value == _NEG_INF:
  465. write(b'\x00\x00\x00\x00\x00\x00\xF0\xFF')
  466. elif value != value: # NaN
  467. write(b'\x00\x00\x00\x00\x00\x00\xF8\x7F')
  468. else:
  469. raise
  470. else:
  471. raise ValueError('Can\'t encode floating-point values that are '
  472. '%d bytes long (only 4 or 8)' % value_size)
  473. def SpecificEncoder(field_number, is_repeated, is_packed):
  474. local_struct_pack = struct.pack
  475. if is_packed:
  476. tag_bytes = TagBytes(field_number, wire_format.WIRETYPE_LENGTH_DELIMITED)
  477. local_EncodeVarint = _EncodeVarint
  478. def EncodePackedField(write, value):
  479. write(tag_bytes)
  480. local_EncodeVarint(write, len(value) * value_size)
  481. for element in value:
  482. # This try/except block is going to be faster than any code that
  483. # we could write to check whether element is finite.
  484. try:
  485. write(local_struct_pack(format, element))
  486. except SystemError:
  487. EncodeNonFiniteOrRaise(write, element)
  488. return EncodePackedField
  489. elif is_repeated:
  490. tag_bytes = TagBytes(field_number, wire_type)
  491. def EncodeRepeatedField(write, value):
  492. for element in value:
  493. write(tag_bytes)
  494. try:
  495. write(local_struct_pack(format, element))
  496. except SystemError:
  497. EncodeNonFiniteOrRaise(write, element)
  498. return EncodeRepeatedField
  499. else:
  500. tag_bytes = TagBytes(field_number, wire_type)
  501. def EncodeField(write, value):
  502. write(tag_bytes)
  503. try:
  504. write(local_struct_pack(format, value))
  505. except SystemError:
  506. EncodeNonFiniteOrRaise(write, value)
  507. return EncodeField
  508. return SpecificEncoder
  509. # ====================================================================
  510. # Here we declare an encoder constructor for each field type. These work
  511. # very similarly to sizer constructors, described earlier.
  512. Int32Encoder = Int64Encoder = EnumEncoder = _SimpleEncoder(
  513. wire_format.WIRETYPE_VARINT, _EncodeSignedVarint, _SignedVarintSize)
  514. UInt32Encoder = UInt64Encoder = _SimpleEncoder(
  515. wire_format.WIRETYPE_VARINT, _EncodeVarint, _VarintSize)
  516. SInt32Encoder = SInt64Encoder = _ModifiedEncoder(
  517. wire_format.WIRETYPE_VARINT, _EncodeVarint, _VarintSize,
  518. wire_format.ZigZagEncode)
  519. # Note that Python conveniently guarantees that when using the '<' prefix on
  520. # formats, they will also have the same size across all platforms (as opposed
  521. # to without the prefix, where their sizes depend on the C compiler's basic
  522. # type sizes).
  523. Fixed32Encoder = _StructPackEncoder(wire_format.WIRETYPE_FIXED32, '<I')
  524. Fixed64Encoder = _StructPackEncoder(wire_format.WIRETYPE_FIXED64, '<Q')
  525. SFixed32Encoder = _StructPackEncoder(wire_format.WIRETYPE_FIXED32, '<i')
  526. SFixed64Encoder = _StructPackEncoder(wire_format.WIRETYPE_FIXED64, '<q')
  527. FloatEncoder = _FloatingPointEncoder(wire_format.WIRETYPE_FIXED32, '<f')
  528. DoubleEncoder = _FloatingPointEncoder(wire_format.WIRETYPE_FIXED64, '<d')
  529. def BoolEncoder(field_number, is_repeated, is_packed):
  530. """Returns an encoder for a boolean field."""
  531. false_byte = b'\x00'
  532. true_byte = b'\x01'
  533. if is_packed:
  534. tag_bytes = TagBytes(field_number, wire_format.WIRETYPE_LENGTH_DELIMITED)
  535. local_EncodeVarint = _EncodeVarint
  536. def EncodePackedField(write, value):
  537. write(tag_bytes)
  538. local_EncodeVarint(write, len(value))
  539. for element in value:
  540. if element:
  541. write(true_byte)
  542. else:
  543. write(false_byte)
  544. return EncodePackedField
  545. elif is_repeated:
  546. tag_bytes = TagBytes(field_number, wire_format.WIRETYPE_VARINT)
  547. def EncodeRepeatedField(write, value):
  548. for element in value:
  549. write(tag_bytes)
  550. if element:
  551. write(true_byte)
  552. else:
  553. write(false_byte)
  554. return EncodeRepeatedField
  555. else:
  556. tag_bytes = TagBytes(field_number, wire_format.WIRETYPE_VARINT)
  557. def EncodeField(write, value):
  558. write(tag_bytes)
  559. if value:
  560. return write(true_byte)
  561. return write(false_byte)
  562. return EncodeField
  563. def StringEncoder(field_number, is_repeated, is_packed):
  564. """Returns an encoder for a string field."""
  565. tag = TagBytes(field_number, wire_format.WIRETYPE_LENGTH_DELIMITED)
  566. local_EncodeVarint = _EncodeVarint
  567. local_len = len
  568. assert not is_packed
  569. if is_repeated:
  570. def EncodeRepeatedField(write, value):
  571. for element in value:
  572. encoded = element.encode('utf-8')
  573. write(tag)
  574. local_EncodeVarint(write, local_len(encoded))
  575. write(encoded)
  576. return EncodeRepeatedField
  577. else:
  578. def EncodeField(write, value):
  579. encoded = value.encode('utf-8')
  580. write(tag)
  581. local_EncodeVarint(write, local_len(encoded))
  582. return write(encoded)
  583. return EncodeField
  584. def BytesEncoder(field_number, is_repeated, is_packed):
  585. """Returns an encoder for a bytes field."""
  586. tag = TagBytes(field_number, wire_format.WIRETYPE_LENGTH_DELIMITED)
  587. local_EncodeVarint = _EncodeVarint
  588. local_len = len
  589. assert not is_packed
  590. if is_repeated:
  591. def EncodeRepeatedField(write, value):
  592. for element in value:
  593. write(tag)
  594. local_EncodeVarint(write, local_len(element))
  595. write(element)
  596. return EncodeRepeatedField
  597. else:
  598. def EncodeField(write, value):
  599. write(tag)
  600. local_EncodeVarint(write, local_len(value))
  601. return write(value)
  602. return EncodeField
  603. def GroupEncoder(field_number, is_repeated, is_packed):
  604. """Returns an encoder for a group field."""
  605. start_tag = TagBytes(field_number, wire_format.WIRETYPE_START_GROUP)
  606. end_tag = TagBytes(field_number, wire_format.WIRETYPE_END_GROUP)
  607. assert not is_packed
  608. if is_repeated:
  609. def EncodeRepeatedField(write, value):
  610. for element in value:
  611. write(start_tag)
  612. element._InternalSerialize(write)
  613. write(end_tag)
  614. return EncodeRepeatedField
  615. else:
  616. def EncodeField(write, value):
  617. write(start_tag)
  618. value._InternalSerialize(write)
  619. return write(end_tag)
  620. return EncodeField
  621. def MessageEncoder(field_number, is_repeated, is_packed):
  622. """Returns an encoder for a message field."""
  623. tag = TagBytes(field_number, wire_format.WIRETYPE_LENGTH_DELIMITED)
  624. local_EncodeVarint = _EncodeVarint
  625. assert not is_packed
  626. if is_repeated:
  627. def EncodeRepeatedField(write, value):
  628. for element in value:
  629. write(tag)
  630. local_EncodeVarint(write, element.ByteSize())
  631. element._InternalSerialize(write)
  632. return EncodeRepeatedField
  633. else:
  634. def EncodeField(write, value):
  635. write(tag)
  636. local_EncodeVarint(write, value.ByteSize())
  637. return value._InternalSerialize(write)
  638. return EncodeField
  639. # --------------------------------------------------------------------
  640. # As before, MessageSet is special.
  641. def MessageSetItemEncoder(field_number):
  642. """Encoder for extensions of MessageSet.
  643. The message set message looks like this:
  644. message MessageSet {
  645. repeated group Item = 1 {
  646. required int32 type_id = 2;
  647. required string message = 3;
  648. }
  649. }
  650. """
  651. start_bytes = b"".join([
  652. TagBytes(1, wire_format.WIRETYPE_START_GROUP),
  653. TagBytes(2, wire_format.WIRETYPE_VARINT),
  654. _VarintBytes(field_number),
  655. TagBytes(3, wire_format.WIRETYPE_LENGTH_DELIMITED)])
  656. end_bytes = TagBytes(1, wire_format.WIRETYPE_END_GROUP)
  657. local_EncodeVarint = _EncodeVarint
  658. def EncodeField(write, value):
  659. write(start_bytes)
  660. local_EncodeVarint(write, value.ByteSize())
  661. value._InternalSerialize(write)
  662. return write(end_bytes)
  663. return EncodeField
  664. # --------------------------------------------------------------------
  665. # As before, Map is special.
  666. def MapEncoder(field_descriptor):
  667. """Encoder for extensions of MessageSet.
  668. Maps always have a wire format like this:
  669. message MapEntry {
  670. key_type key = 1;
  671. value_type value = 2;
  672. }
  673. repeated MapEntry map = N;
  674. """
  675. # Can't look at field_descriptor.message_type._concrete_class because it may
  676. # not have been initialized yet.
  677. message_type = field_descriptor.message_type
  678. encode_message = MessageEncoder(field_descriptor.number, False, False)
  679. def EncodeField(write, value):
  680. for key in value:
  681. entry_msg = message_type._concrete_class(key=key, value=value[key])
  682. encode_message(write, entry_msg)
  683. return EncodeField