Fix image storage enhancement [ER#160: 6bigfire]

Author: littlecodersh

itchat/components/contact.py

@@ -413,10 +413,13 @@ def get_head_img(self, userName=None, chatroomUserName=None, picDir=None):
         tempStorage.write(block)
     if picDir is None:
         return tempStorage.getvalue()
-    with open(picDir, 'wb') as f: f.write(tempStorage.getvalue())
+    with open(picDir, 'wb') as f:
+        f.write(tempStorage.getvalue())
+    tempStorage.seek(0)
     return ReturnValue({'BaseResponse': {
         'ErrMsg': 'Successfully downloaded',
-        'Ret': 0, }})
+        'Ret': 0, },
+        'PostFix': utils.get_image_postfix(tempStorage.read(20)), })
 
 def create_chatroom(self, memberList, topic=''):
     url = '%s/webwxcreatechatroom?pass_ticket=%s&r=%s' % (

itchat/components/messages.py

@@ -31,11 +31,15 @@ def download_fn(downloadDir=None):
         tempStorage = io.BytesIO()
         for block in r.iter_content(1024):
             tempStorage.write(block)
-        if downloadDir is None: return tempStorage.getvalue()
-        with open(downloadDir, 'wb') as f: f.write(tempStorage.getvalue())
+        if downloadDir is None:
+            return tempStorage.getvalue()
+        with open(downloadDir, 'wb') as f:
+            f.write(tempStorage.getvalue())
+        tempStorage.seek(0)
         return ReturnValue({'BaseResponse': {
             'ErrMsg': 'Successfully downloaded',
-            'Ret': 0, }})
+            'Ret': 0, },
+            'PostFix': utils.get_image_postfix(tempStorage.read(20)), })
     return download_fn
 
 def produce_msg(core, msgList):
@@ -101,8 +105,10 @@ def download_video(videoDir=None):
                 tempStorage = io.BytesIO()
                 for block in r.iter_content(1024):
                     tempStorage.write(block)
-                if videoDir is None: return tempStorage.getvalue()
-                with open(videoDir, 'wb') as f: f.write(tempStorage.getvalue())
+                if videoDir is None:
+                    return tempStorage.getvalue()
+                with open(videoDir, 'wb') as f:
+                    f.write(tempStorage.getvalue())
                 return ReturnValue({'BaseResponse': {
                     'ErrMsg': 'Successfully downloaded',
                     'Ret': 0, }})
@@ -128,8 +134,10 @@ def download_atta(attaDir=None):
                     tempStorage = io.BytesIO()
                     for block in r.iter_content(1024):
                         tempStorage.write(block)
-                    if attaDir is None: return tempStorage.getvalue()
-                    with open(attaDir, 'wb') as f: f.write(tempStorage.getvalue())
+                    if attaDir is None:
+                        return tempStorage.getvalue()
+                    with open(attaDir, 'wb') as f:
+                        f.write(tempStorage.getvalue())
                     return ReturnValue({'BaseResponse': {
                         'ErrMsg': 'Successfully downloaded',
                         'Ret': 0, }})
@@ -255,7 +263,8 @@ def upload_file(self, fileDir, isPicture=False, isVideo=False,
             'Ret': -1002, }})
     fileSize = os.path.getsize(fileDir)
     fileSymbol = 'pic' if isPicture else 'video' if isVideo else'doc'
-    with open(fileDir, 'rb') as f: fileMd5 = hashlib.md5(f.read()).hexdigest()
+    with open(fileDir, 'rb') as f:
+        fileMd5 = hashlib.md5(f.read()).hexdigest()
     file_ = open(fileDir, 'rb')
     chunks = int((fileSize - 1) / 524288) + 1
     clientMediaId = int(time.time() * 1e4)

itchat/utils.py

@@ -70,7 +70,8 @@ def msg_formatter(d, k):
 
 def check_file(fileDir):
     try:
-        with open(fileDir): pass
+        with open(fileDir):
+            pass
         return True
     except:
         return False
@@ -130,3 +131,13 @@ def test_connect(retryTime=5):
 def contact_deep_copy(core, contact):
     with core.storageClass.updateLock:
         return copy.deepcopy(contact)
+
+def get_image_postfix(data):
+    data = data[:20]
+    if 'GIF' in data:
+        return 'gif'
+    elif 'PNG' in data:
+        return 'png'
+    elif 'JFIF' in data:
+        return 'jpg'
+    return ''