fix bug, add avt_db support

dump-html.py

@@ -5,8 +5,8 @@
 # Author: Yuxin Wu <[email protected]>
 
 import sys
-if len(sys.argv) != 5:
-    sys.exit("Usage: {0} <path to decoded_database.db> <path to resource> <name> <output html>".format(sys.argv[0]))
+if len(sys.argv) != 6:
+    sys.exit("Usage: {0} <path to decoded_database.db> <avt_db> <path to resource>  <name> <output html>".format(sys.argv[0]))
 
 from common.textutil import ensure_unicode
 from wechat.parser import WeChatDBParser
@@ -14,12 +14,13 @@ from wechat.res import Resource
 from wechat.render import HTMLRender
 
 db_file = sys.argv[1]
-resource_dir = sys.argv[2]
-name = ensure_unicode(sys.argv[3])
-output_file = sys.argv[4]
+avt_db = sys.argv[2]
+resource_dir = sys.argv[3]
+name = ensure_unicode(sys.argv[4])
+output_file = sys.argv[5]
 
 parser = WeChatDBParser(db_file)
-res = Resource(resource_dir)
+res = Resource(resource_dir, avt_db)
 
 try:
     msgs = parser.msgs_by_talker[name]

wechat/res.py

@@ -34,7 +34,7 @@ JPEG_QUALITY = 50
 
 class Resource(object):
     """ multimedia resources in chat"""
-    def __init__(self, res_dir):
+    def __init__(self, res_dir, avt_db):
         def check(subdir):
             assert os.path.isdir(os.path.join(res_dir, subdir)), \
                     "No such directory: {}".format(subdir)
@@ -45,7 +45,7 @@ class Resource(object):
         self.img_dir = os.path.join(res_dir, IMG_DIRNAME)
         self.voice_dir = os.path.join(res_dir, VOICE_DIRNAME)
         self.emoji_dir = os.path.join(res_dir, EMOJI_DIRNAME)
-        self.avt_reader = AvatarReader(os.path.join(res_dir, AVATAR_DIRNAME))
+        self.avt_reader = AvatarReader(os.path.join(res_dir, AVATAR_DIRNAME), avt_db)
 
     def get_voice_filename(self, imgpath):
         fname = md5(imgpath)
@@ -140,7 +140,7 @@ class Resource(object):
                imghdr.what(img_file) != 'jpeg':
                 im = Image.open(open(img_file, 'rb'))
                 buf = cStringIO.StringIO()
-                im.save(buf, 'JPEG', quality=JPEG_QUALITY)
+                im.convert('RGB').save(buf, 'JPEG', quality=JPEG_QUALITY)
                 return base64.b64encode(buf.getvalue())
             return get_file_b64(img_file)
         big_file = get_jpg_b64(big_file)