|
8 | 8 | except ImportError:
|
9 | 9 | from io import StringIO
|
10 | 10 |
|
| 11 | +from io import BytesIO |
11 | 12 | from gzip import GzipFile
|
12 | 13 | import json
|
13 | 14 | import logging
|
@@ -75,14 +76,26 @@ def user_rooms(self, userid, shard='shard0'):
|
75 | 76 | def memory(self, path='', shard='shard0'):
|
76 | 77 | ret = self.get('user/memory', path=path, shard=shard)
|
77 | 78 | if 'data' in ret:
|
78 |
| - ret['data'] = json.load(GzipFile(fileobj=StringIO(b64decode(ret['data'][3:])))) |
| 79 | + try: |
| 80 | + gzip_input = StringIO(b64decode(ret['data'][3:])) |
| 81 | + except: |
| 82 | + gzip_input = BytesIO(b64decode(ret['data'][3:])) |
| 83 | + ret['data'] = json.load(GzipFile(fileobj=gzip_input)) |
79 | 84 | return ret
|
80 | 85 |
|
81 | 86 | def set_memory(self, path, value, shard='shard0'):
|
82 | 87 | return self.post('user/memory', path=path, value=value, shard=shard)
|
83 | 88 |
|
84 | 89 | def get_segment(self, segment, shard='shard0'):
|
85 |
| - return self.get('user/memory-segment', segment=segment, shard=shard) |
| 90 | + ret = self.get('user/memory-segment', segment=segment, shard=shard) |
| 91 | + if 'data' in ret and ret['data'][:3] == 'gz:': |
| 92 | + try: |
| 93 | + gzip_input = StringIO(b64decode(ret['data'][3:])) |
| 94 | + except: |
| 95 | + gzip_input = BytesIO(b64decode(ret['data'][3:])) |
| 96 | + ret['data'] = GzipFile(fileobj=gzip_input) |
| 97 | + return ret |
| 98 | + |
86 | 99 |
|
87 | 100 | def set_segment(self, segment, data, shard='shard0'):
|
88 | 101 | return self.post('user/memory-segment', segment=segment, data=data, shard=shard)
|
|
0 commit comments