check for two avatar storage

wechat/avatar.py

@@ -26,51 +26,59 @@ class AvatarReader(object):
 
         # new location of avatar, see #50
         self.avt_dir = os.path.join(res_dir, 'avatar')
+        if not os.path.isdir(self.avt_dir) or len(os.listdir(self.avt_dir)) == 0:
+            self.avt_dir = None
+
         self.avt_db = avt_db
         self._use_avt = True
-        if os.path.isdir(self.avt_dir) and len(os.listdir(self.avt_dir)):
-            # has resource/avatar
-            pass
-        elif self.avt_db is not None \
-                and os.path.isfile(self.avt_db) \
-                and glob.glob(os.path.join(self.sfs_dir, 'avatar*')):
-            # has sfs/avatar*
-            pass
-        else:
+        if self.avt_db is not None:
+            if len(glob.glob(os.path.join(self.sfs_dir, 'avatar*'))) == 0:
+                # has sfs/avatar*
+                self.avt_db = None
+        if self.avt_dir is None and self.avt_db is None:
             logger.warn("Cannot find avatar storage. Will not use avatar!")
             self._use_avt = False
 
-    def get_avatar(self, username):
-        """ username: `username` field in db.rcontact"""
-        if not self._use_avt:
-            return None
-        username = ensure_unicode(username)
-        avtid = md5(username.encode('utf-8'))
+    def get_avatar_from_avtdb(self, avtid):
         try:
-            candidates = self.query_index(avtid)
+            candidates = self._search_avt_db(avtid)
             candidates = sorted(candidates, key=lambda x: _filename_priority(x[0]), reverse=True)
             for c in candidates:
                 path, offset, size = c
                 return self.read_img_from_block(path, offset, size)
-        except:
+        except Exception:
             pass
 
+    def get_avatar_from_avtdir(self, avtid):
         dir1, dir2 = avtid[:2], avtid[2:4]
         candidates = glob.glob(os.path.join(self.avt_dir, dir1, dir2, f"*{avtid}*"))
-        default_candidate = os.path.join(self.avt_dir, dir1, dir2, f"user_{avtid}.png")
-        candidates.append(default_candidate)
-
         candidates = sorted(set(candidates), key=_filename_priority, reverse=True)
         for cand in candidates:
             try:
-                if os.path.exists(cand):
-                    if cand.endswith(".bm"):
-                        return self.read_bm_file(cand)
-                    else:
-                        return Image.open(cand)
+                if cand.endswith(".bm"):
+                    return self.read_bm_file(cand)
+                else:
+                    return Image.open(cand)
             except Exception:
                 logger.exception("")
                 pass
+
+    def get_avatar(self, username):
+        """ username: `username` field in db.rcontact"""
+        if not self._use_avt:
+            return None
+        username = ensure_unicode(username)
+        avtid = md5(username.encode('utf-8'))
+
+        if self.avt_db is not None:
+            ret = self.get_avatar_from_avtdb(avtid)
+            if ret is not None:
+                return ret
+
+        if self.avt_dir is not None:
+            ret = self.get_avatar_from_avtdir(avtid)
+            if ret is not None:
+                return ret
         logger.warning("Avatar for {} not found anywhere.".format(username))
 
     def read_img_from_block(self, filename, pos, size):
@@ -101,7 +109,7 @@ class AvatarReader(object):
                     img[i,j] = (r, g, b)
             return Image.fromarray(img, mode="RGB")
 
-    def query_index(self, avtid):
+    def _search_avt_db(self, avtid):
         conn = sqlite3.connect(self.avt_db)
         cursor = conn.execute("select FileName,Offset,Size from Index_avatar")
         candidates = []