3434from riak .transports .http import RiakHttpPool
3535from riak .transports .pbc import RiakPbcPool
3636from riak .security import SecurityCreds
37- from riak .util import lazy_property
37+ from riak .util import lazy_property , bytes_to_str , str_to_bytes
38+ from six import string_types , PY2
3839
3940
4041def default_encoder (obj ):
4142 """
4243 Default encoder for JSON datatypes, which returns UTF-8 encoded
43- json instead of the default bloated \uXXXX escaped ASCII strings.
44+ json instead of the default bloated backslash u XXXX escaped ASCII strings.
4445 """
45- return json .dumps (obj , ensure_ascii = False ).encode ("utf-8" )
46+ if type (obj ) == bytes :
47+ return json .dumps (bytes_to_str (obj ),
48+ ensure_ascii = False ).encode ("utf-8" )
49+ else :
50+ return json .dumps (obj , ensure_ascii = False ).encode ("utf-8" )
51+
52+
53+ def binary_json_encoder (obj ):
54+ """
55+ Default encoder for JSON datatypes, which returns UTF-8 encoded
56+ json instead of the default bloated backslash u XXXX escaped ASCII strings.
57+ """
58+ if type (obj ) == bytes :
59+ return json .dumps (bytes_to_str (obj ),
60+ ensure_ascii = False ).encode ("utf-8" )
61+ else :
62+ return json .dumps (obj , ensure_ascii = False ).encode ("utf-8" )
63+
64+
65+ def binary_json_decoder (obj ):
66+ """
67+ Default decoder from JSON datatypes.
68+ """
69+ return json .loads (bytes_to_str (obj ))
70+
71+
72+ def binary_encoder_decoder (obj ):
73+ """
74+ Assumes value is already in binary format, so passes unchanged.
75+ """
76+ return obj
4677
4778
4879class RiakClient (RiakMapReduceChain , RiakClientOperations ):
@@ -90,12 +121,22 @@ def __init__(self, protocol='pbc', transport_options={}, nodes=None,
90121 self ._http_pool = RiakHttpPool (self , ** transport_options )
91122 self ._pb_pool = RiakPbcPool (self , ** transport_options )
92123
93- self ._encoders = {'application/json' : default_encoder ,
94- 'text/json' : default_encoder ,
95- 'text/plain' : str }
96- self ._decoders = {'application/json' : json .loads ,
97- 'text/json' : json .loads ,
98- 'text/plain' : str }
124+ if PY2 :
125+ self ._encoders = {'application/json' : default_encoder ,
126+ 'text/json' : default_encoder ,
127+ 'text/plain' : str }
128+ self ._decoders = {'application/json' : json .loads ,
129+ 'text/json' : json .loads ,
130+ 'text/plain' : str }
131+ else :
132+ self ._encoders = {'application/json' : binary_json_encoder ,
133+ 'text/json' : binary_json_encoder ,
134+ 'text/plain' : str_to_bytes ,
135+ 'binary/octet-stream' : binary_encoder_decoder }
136+ self ._decoders = {'application/json' : binary_json_decoder ,
137+ 'text/json' : binary_json_decoder ,
138+ 'text/plain' : bytes_to_str ,
139+ 'binary/octet-stream' : binary_encoder_decoder }
99140 self ._buckets = WeakValueDictionary ()
100141 self ._bucket_types = WeakValueDictionary ()
101142
@@ -167,7 +208,7 @@ def set_encoder(self, content_type, encoder):
167208 :param content_type: the requested media type
168209 :type content_type: str
169210 :param encoder: an encoding function, takes a single object
170- argument and returns a string
211+ argument and returns encoded data
171212 :type encoder: function
172213 """
173214 self ._encoders [content_type ] = encoder
@@ -188,7 +229,7 @@ def set_decoder(self, content_type, decoder):
188229
189230 :param content_type: the requested media type
190231 :type content_type: str
191- :param decoder: a decoding function, takes a string and
232+ :param decoder: a decoding function, takes encoded data and
192233 returns a Python type
193234 :type decoder: function
194235 """
@@ -217,10 +258,10 @@ def bucket(self, name, bucket_type='default'):
217258 :rtype: :class:`RiakBucket <riak.bucket.RiakBucket>`
218259
219260 """
220- if not isinstance (name , basestring ):
261+ if not isinstance (name , string_types ):
221262 raise TypeError ('Bucket name must be a string' )
222263
223- if isinstance (bucket_type , basestring ):
264+ if isinstance (bucket_type , string_types ):
224265 bucket_type = self .bucket_type (bucket_type )
225266 elif not isinstance (bucket_type , BucketType ):
226267 raise TypeError ('bucket_type must be a string '
@@ -243,7 +284,7 @@ def bucket_type(self, name):
243284 :type name: str
244285 :rtype: :class:`BucketType <riak.bucket.BucketType>`
245286 """
246- if not isinstance (name , basestring ):
287+ if not isinstance (name , string_types ):
247288 raise TypeError ('Bucket name must be a string' )
248289
249290 if name in self ._bucket_types :
0 commit comments