1
0
mirror of https://github.com/trezor/trezor-firmware.git synced 2024-12-26 16:18:22 +00:00

feat(core): make protobuf.dump_uvarint more reusable

(cherry picked from commit 6acc1cd6ab)
This commit is contained in:
matejcik 2021-01-14 16:15:03 +01:00 committed by Tomas Susanka
parent 78a2ff16d4
commit 2231fe6842
2 changed files with 17 additions and 15 deletions

View File

@ -4,7 +4,7 @@ bytes, string, embedded message and repeated fields.
"""
if False:
from typing import Any, Dict, Iterable, List, Tuple, Type, TypeVar, Union
from typing import Any, Callable, Dict, Iterable, List, Tuple, Type, TypeVar, Union
from typing_extensions import Protocol
class Reader(Protocol):
@ -19,6 +19,8 @@ if False:
Writes all bytes from `buf`, or raises `EOFError`.
"""
WriteMethod = Callable[[bytes], Any]
_UVARINT_BUFFER = bytearray(1)
@ -36,7 +38,7 @@ def load_uvarint(reader: Reader) -> int:
return result
def dump_uvarint(writer: Writer, n: int) -> None:
def dump_uvarint(write: WriteMethod, n: int) -> None:
if n < 0:
raise ValueError("Cannot dump signed value, convert it to unsigned first.")
buffer = _UVARINT_BUFFER
@ -44,7 +46,7 @@ def dump_uvarint(writer: Writer, n: int) -> None:
while shifted:
shifted = n >> 7
buffer[0] = (n & 0x7F) | (0x80 if shifted else 0x00)
writer.write(buffer)
write(buffer)
n = shifted
@ -318,32 +320,32 @@ def dump_message(
fvalue = repvalue
for svalue in fvalue:
dump_uvarint(writer, fkey)
dump_uvarint(writer.write, fkey)
if ftype is UVarintType:
dump_uvarint(writer, svalue)
dump_uvarint(writer.write, svalue)
elif ftype is SVarintType:
dump_uvarint(writer, sint_to_uint(svalue))
dump_uvarint(writer.write, sint_to_uint(svalue))
elif ftype is BoolType:
dump_uvarint(writer, int(svalue))
dump_uvarint(writer.write, int(svalue))
elif isinstance(ftype, EnumType):
dump_uvarint(writer, svalue)
dump_uvarint(writer.write, svalue)
elif ftype is BytesType:
if isinstance(svalue, list):
dump_uvarint(writer, _count_bytes_list(svalue))
dump_uvarint(writer.write, _count_bytes_list(svalue))
for sub_svalue in svalue:
writer.write(sub_svalue)
else:
dump_uvarint(writer, len(svalue))
dump_uvarint(writer.write, len(svalue))
writer.write(svalue)
elif ftype is UnicodeType:
svalue = svalue.encode()
dump_uvarint(writer, len(svalue))
dump_uvarint(writer.write, len(svalue))
writer.write(svalue)
elif issubclass(ftype, MessageType):
@ -351,7 +353,7 @@ def dump_message(
if ffields is None:
ffields = ftype.get_fields()
field_cache[ftype] = ffields
dump_uvarint(writer, count_message(svalue, field_cache))
dump_uvarint(writer.write, count_message(svalue, field_cache))
dump_message(writer, svalue, field_cache)
else:

View File

@ -36,9 +36,9 @@ def load_uvarint(data: bytes) -> int:
def dump_uvarint(value: int) -> bytearray:
writer = BufferWriter(bytearray(16))
protobuf.dump_uvarint(writer, value)
return memoryview(writer.buffer)[: writer.offset]
w = bytearray()
protobuf.dump_uvarint(w.extend, value)
return w
def dump_message(msg: protobuf.MessageType) -> bytearray: