Skip to content

Commit 10d3a57

Browse files
author
Dominik Andreas
committed
added binary support in dictionaries via base64 encoding
1 parent 3a3adfb commit 10d3a57

File tree

2 files changed

+48
-16
lines changed

2 files changed

+48
-16
lines changed

capnp/lib/capnp.pyx

Lines changed: 27 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@ import array
2727
import asyncio
2828
import collections as _collections
2929
import contextlib
30+
import base64
3031
import enum as _enum
3132
import inspect as _inspect
3233
import os as _os
@@ -969,17 +970,17 @@ cdef _DynamicStructBuilder temp_msg_b
969970
cdef _DynamicStructReader temp_msg_r
970971

971972

972-
cdef _to_dict(msg, bint verbose, bint ordered):
973+
cdef _to_dict(msg, bint verbose, bint ordered, bint encode_bytes_as_base64=False):
973974
msg_type = type(msg)
974975
if msg_type is _DynamicListBuilder:
975976
temp_list_b = msg
976-
return [_to_dict(temp_list_b._get(i), verbose, ordered) for i in range(len(msg))]
977+
return [_to_dict(temp_list_b._get(i), verbose, ordered, encode_bytes_as_base64) for i in range(len(msg))]
977978
elif msg_type is _DynamicListReader:
978979
temp_list_r = msg
979-
return [_to_dict(temp_list_r._get(i), verbose, ordered) for i in range(len(msg))]
980+
return [_to_dict(temp_list_r._get(i), verbose, ordered, encode_bytes_as_base64) for i in range(len(msg))]
980981
elif msg_type is _DynamicResizableListBuilder:
981982
temp_list_rb = msg
982-
return [_to_dict(temp_list_rb._get(i), verbose, ordered) for i in range(len(msg))]
983+
return [_to_dict(temp_list_rb._get(i), verbose, ordered, encode_bytes_as_base64) for i in range(len(msg))]
983984

984985
if msg_type is _DynamicStructBuilder or isinstance(msg, _Request):
985986
temp_msg_b = msg
@@ -989,13 +990,13 @@ cdef _to_dict(msg, bint verbose, bint ordered):
989990
ret = {}
990991
try:
991992
which = temp_msg_b.which()
992-
ret[which] = _to_dict(temp_msg_b._get(which), verbose, ordered)
993+
ret[which] = _to_dict(temp_msg_b._get(which), verbose, ordered, encode_bytes_as_base64)
993994
except KjException:
994995
pass
995996

996997
for field in temp_msg_b.schema.non_union_fields:
997998
if verbose or temp_msg_b._has(field):
998-
ret[field] = _to_dict(temp_msg_b._get(field), verbose, ordered)
999+
ret[field] = _to_dict(temp_msg_b._get(field), verbose, ordered, encode_bytes_as_base64)
9991000

10001001
return ret
10011002
elif msg_type is _DynamicStructReader or isinstance(msg, _Response):
@@ -1006,13 +1007,13 @@ cdef _to_dict(msg, bint verbose, bint ordered):
10061007
ret = {}
10071008
try:
10081009
which = temp_msg_r.which()
1009-
ret[which] = _to_dict(temp_msg_r._get(which), verbose, ordered)
1010+
ret[which] = _to_dict(temp_msg_r._get(which), verbose, ordered, encode_bytes_as_base64)
10101011
except KjException:
10111012
pass
10121013

10131014
for field in temp_msg_r.schema.non_union_fields:
10141015
if verbose or temp_msg_r._has(field):
1015-
ret[field] = _to_dict(temp_msg_r._get(field), verbose, ordered)
1016+
ret[field] = _to_dict(temp_msg_r._get(field), verbose, ordered, encode_bytes_as_base64)
10161017

10171018
return ret
10181019

@@ -1022,6 +1023,10 @@ cdef _to_dict(msg, bint verbose, bint ordered):
10221023
if msg_type is _DynamicEnum:
10231024
return str(msg)
10241025

1026+
if encode_bytes_as_base64 and msg_type is bytes:
1027+
# encode the message as base64 and return utf-8 string
1028+
return base64.b64encode(msg).decode('utf-8')
1029+
10251030
return msg
10261031

10271032

@@ -1234,8 +1239,8 @@ cdef class _DynamicStructReader:
12341239
def __repr__(self):
12351240
return '<%s reader %s>' % (self.schema.node.displayName, <char*>strStructReader(self.thisptr).cStr())
12361241

1237-
def to_dict(self, verbose=False, ordered=False):
1238-
return _to_dict(self, verbose, ordered)
1242+
def to_dict(self, verbose=False, ordered=False, encode_bytes_as_base64=False):
1243+
return _to_dict(self, verbose, ordered, encode_bytes_as_base64)
12391244

12401245
cpdef as_builder(self, num_first_segment_words=None):
12411246
"""A method for casting this Reader to a Builder
@@ -1618,12 +1623,18 @@ cdef class _DynamicStructBuilder:
16181623
def __repr__(self):
16191624
return '<%s builder %s>' % (self.schema.node.displayName, <char*>strStructBuilder(self.thisptr).cStr())
16201625

1621-
def to_dict(self, verbose=False, ordered=False):
1622-
return _to_dict(self, verbose, ordered)
1626+
def to_dict(self, verbose=False, ordered=False, encode_bytes_as_base64=False):
1627+
return _to_dict(self, verbose, ordered, encode_bytes_as_base64)
16231628

16241629
def from_dict(self, dict d):
16251630
for key, val in d.iteritems():
16261631
if key != 'which':
1632+
field = self.schema.fields.get(key)
1633+
if isinstance(val, str):
1634+
dtype = field.proto.slot.type.which()
1635+
if dtype == "data":
1636+
# decode bytes from utf-8 base64 encoding
1637+
val = base64.b64decode(val)
16271638
try:
16281639
self._set(key, val)
16291640
except Exception as e:
@@ -1703,8 +1714,8 @@ cdef class _DynamicStructPipeline:
17031714
# def __repr__(self):
17041715
# return '<%s reader %s>' % (self.schema.node.displayName, strStructReader(self.thisptr).cStr())
17051716

1706-
def to_dict(self, verbose=False, ordered=False):
1707-
return _to_dict(self, verbose, ordered)
1717+
def to_dict(self, verbose=False, ordered=False, encode_bytes_as_base64=False):
1718+
return _to_dict(self, verbose, ordered, encode_bytes_as_base64)
17081719

17091720

17101721
cdef class _DynamicOrphan:
@@ -2086,8 +2097,8 @@ cdef class _RemotePromise:
20862097
def __dir__(self):
20872098
return list(set(self.schema.fieldnames + tuple(dir(self.__class__))))
20882099

2089-
def to_dict(self, verbose=False, ordered=False):
2090-
return _to_dict(self, verbose, ordered)
2100+
def to_dict(self, verbose=False, ordered=False, encode_bytes_as_base64=False):
2101+
return _to_dict(self, verbose, ordered, encode_bytes_as_base64)
20912102

20922103
cpdef cancel(self):
20932104
self.thisptr = Own[RemotePromise]()

test/test_blob_to_dict_base64.py

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,21 @@
1+
import os
2+
import capnp
3+
import base64
4+
import pytest
5+
6+
this_dir = os.path.dirname(__file__)
7+
8+
9+
@pytest.fixture
10+
def blob_schema():
11+
return capnp.load(os.path.join(this_dir, "blob_test.capnp"))
12+
13+
14+
def test_blob_to_dict(blob_schema):
15+
blob_value = b"hello world"
16+
blob = blob_schema.BlobTest(blob=blob_value)
17+
blob_dict = blob.to_dict(encode_bytes_as_base64=True)
18+
assert base64.b64decode(blob_dict["blob"]) == blob_value
19+
msg = blob_schema.BlobTest.new_message()
20+
msg.from_dict(blob_dict)
21+
assert blob.blob == blob_value

0 commit comments

Comments
 (0)