diff --git a/wechat/avatar.py b/wechat/avatar.py index dde908e..dd9d047 100644 --- a/wechat/avatar.py +++ b/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_from_avtdb(self, avtid): + try: + 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 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}*")) + candidates = sorted(set(candidates), key=_filename_priority, reverse=True) + for cand in candidates: + try: + 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')) - try: - candidates = self.query_index(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: - pass - 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) + if self.avt_db is not None: + ret = self.get_avatar_from_avtdb(avtid) + if ret is not None: + return ret - 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) - except Exception: - logger.exception("") - pass + 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 = []