memory loading bug fix for python3 (bytes/string change)

screepers · Aug 30, 2017 · bdf32f7 · bdf32f7
1 parent a33a1df
commit bdf32f7
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 3 deletions.
diff --git a/screepsapi/screepsapi.py b/screepsapi/screepsapi.py
@@ -8,6 +8,7 @@
 except ImportError:
     from io import StringIO
 
+from io import BytesIO
 from gzip import GzipFile
 import json
 import logging
@@ -75,14 +76,26 @@ def user_rooms(self, userid, shard='shard0'):
     def memory(self, path='', shard='shard0'):
         ret = self.get('user/memory', path=path, shard=shard)
         if 'data' in ret:
-            ret['data'] = json.load(GzipFile(fileobj=StringIO(b64decode(ret['data'][3:]))))
+            try:
+                gzip_input = StringIO(b64decode(ret['data'][3:]))
+            except:
+                gzip_input = BytesIO(b64decode(ret['data'][3:]))
+            ret['data'] = json.load(GzipFile(fileobj=gzip_input))
         return ret
 
     def set_memory(self, path, value, shard='shard0'):
         return self.post('user/memory', path=path, value=value, shard=shard)
 
     def get_segment(self, segment, shard='shard0'):
-        return self.get('user/memory-segment', segment=segment, shard=shard)
+        ret = self.get('user/memory-segment', segment=segment, shard=shard)
+        if 'data' in ret and ret['data'][:3] == 'gz:':
+            try:
+                gzip_input = StringIO(b64decode(ret['data'][3:]))
+            except:
+                gzip_input = BytesIO(b64decode(ret['data'][3:]))
+            ret['data'] = GzipFile(fileobj=gzip_input)
+        return ret
+
 
     def set_segment(self, segment, data, shard='shard0'):
         return self.post('user/memory-segment', segment=segment, data=data, shard=shard)

diff --git a/setup.py b/setup.py
@@ -8,7 +8,7 @@
     long_description = open('README.md').read()
 
 
-version = '0.3.0'
+version = '0.4.1'
 setup(
   name = 'screepsapi',
   version = version,