Skip to content

Commit

Permalink
Fix unpickling in Python 3
Browse files Browse the repository at this point in the history
  • Loading branch information
ob-stripe committed Oct 20, 2017
1 parent d1962df commit d19ccf1
Show file tree
Hide file tree
Showing 2 changed files with 30 additions and 1 deletion.
22 changes: 22 additions & 0 deletions stripe/stripe_object.py
Original file line number Diff line number Diff line change
Expand Up @@ -115,6 +115,28 @@ def __delitem__(self, k):
if hasattr(self, '_unsaved_values'):
self._unsaved_values.remove(k)

# Custom unpickling method that uses `update` to update the dictionary
# without calling __setitem__, which would fail if any value is an empty
# string
def __setstate__(self, state):
self.update(state)

# Custom pickling method to ensure the instance is pickled as a custom
# class and not as a dict, otherwise __setstate__ would not be called when
# unpickling.
def __reduce__(self):
reduce_value = (
type(self), # callable
( # args
self.get('id', None),
self.api_key,
self.stripe_version,
self.stripe_account
),
dict(self), # state
)
return reduce_value

@classmethod
def construct_from(cls, values, key, stripe_version=None,
stripe_account=None):
Expand Down
9 changes: 8 additions & 1 deletion stripe/test/resources/test_stripe_object.py
Original file line number Diff line number Diff line change
Expand Up @@ -142,7 +142,13 @@ def test_pickling(self):
'foo', 'bar', myparam=5)

obj['object'] = 'boo'
obj.refresh_from({'fala': 'lalala'}, api_key='bar', partial=True)
obj.refresh_from(
{
'fala': 'lalala',
'emptystring': '',
},
api_key='bar', partial=True
)

self.assertEqual('lalala', obj.fala)

Expand All @@ -153,6 +159,7 @@ def test_pickling(self):
self.assertEqual('bar', newobj.api_key)
self.assertEqual('boo', newobj['object'])
self.assertEqual('lalala', newobj.fala)
self.assertEqual('', newobj.emptystring)

def test_deletion(self):
obj = stripe.stripe_object.StripeObject('id', 'key')
Expand Down

0 comments on commit d19ccf1

Please sign in to comment.