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

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 ''