|
|
@@ -13,6 +13,13 @@ logger = logging.getLogger(__name__)
|
|
|
from .common.textutil import ensure_unicode, md5
|
|
|
|
|
|
|
|
|
+def _filename_priority(s):
|
|
|
+ if "_hd" in s and s.endswith(".png"):
|
|
|
+ return 10
|
|
|
+ else:
|
|
|
+ return 1
|
|
|
+
|
|
|
+
|
|
|
class AvatarReader(object):
|
|
|
def __init__(self, res_dir, avt_db="avatar.index"):
|
|
|
self.sfs_dir = os.path.join(res_dir, 'sfs')
|
|
|
@@ -22,14 +29,15 @@ class AvatarReader(object):
|
|
|
self.avt_db = avt_db
|
|
|
self._use_avt = True
|
|
|
if os.path.isdir(self.avt_dir) and len(os.listdir(self.avt_dir)):
|
|
|
- self.avt_use_db = False
|
|
|
+ # 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*')):
|
|
|
- self.avt_use_db = True
|
|
|
+ # has sfs/avatar*
|
|
|
+ pass
|
|
|
else:
|
|
|
- logger.warn(
|
|
|
- "Cannot find avatar files. Will not use avatar!")
|
|
|
+ logger.warn("Cannot find avatar storage. Will not use avatar!")
|
|
|
self._use_avt = False
|
|
|
|
|
|
def get_avatar(self, username):
|
|
|
@@ -38,41 +46,39 @@ class AvatarReader(object):
|
|
|
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)
|
|
|
|
|
|
- def priority(s):
|
|
|
- if "_hd" in s and s.endswith(".png"):
|
|
|
- return 10
|
|
|
- else:
|
|
|
- return 1
|
|
|
-
|
|
|
- candidates = sorted(set(candidates), key=priority, reverse=True)
|
|
|
-
|
|
|
+ candidates = sorted(set(candidates), key=_filename_priority, reverse=True)
|
|
|
for cand in candidates:
|
|
|
try:
|
|
|
- if self.avt_use_db:
|
|
|
- pos, size = self.query_index(cand)
|
|
|
- return self.read_img(pos, size)
|
|
|
- else:
|
|
|
- if os.path.exists(cand):
|
|
|
- if cand.endswith(".bm"):
|
|
|
- return self.read_bm_file(cand)
|
|
|
- else:
|
|
|
- return Image.open(cand)
|
|
|
+ if os.path.exists(cand):
|
|
|
+ if cand.endswith(".bm"):
|
|
|
+ return self.read_bm_file(cand)
|
|
|
+ else:
|
|
|
+ return Image.open(cand)
|
|
|
except Exception:
|
|
|
- logger.exception("HHH")
|
|
|
+ logger.exception("")
|
|
|
pass
|
|
|
- logger.warning("Avatar for {} not found in avatar database.".format(username))
|
|
|
+ logger.warning("Avatar for {} not found anywhere.".format(username))
|
|
|
|
|
|
- def read_img(self, pos, size):
|
|
|
+ def read_img_from_block(self, filename, pos, size):
|
|
|
file_idx = pos >> 32
|
|
|
fname = os.path.join(self.sfs_dir,
|
|
|
'avatar.block.' + '{:05d}'.format(file_idx))
|
|
|
- # a 64-byte offset of each block file
|
|
|
- start_pos = pos - file_idx * (2**32) + 64
|
|
|
+ # offset of each block file: 17 + len(path)
|
|
|
+ start_pos = pos - file_idx * (2**32) + 16 + len(filename) + 1
|
|
|
try:
|
|
|
with open(fname, 'rb') as f:
|
|
|
f.seek(start_pos)
|
|
|
@@ -95,11 +101,14 @@ class AvatarReader(object):
|
|
|
img[i,j] = (r, g, b)
|
|
|
return Image.fromarray(img, mode="RGB")
|
|
|
|
|
|
- def query_index(self, filename):
|
|
|
+ def query_index(self, avtid):
|
|
|
conn = sqlite3.connect(self.avt_db)
|
|
|
- cursor = conn.execute("select Offset,Size from Index_avatar where FileName='{}'".format(filename))
|
|
|
- pos, size = cursor.fetchone()
|
|
|
- return pos, size
|
|
|
+ cursor = conn.execute("select FileName,Offset,Size from Index_avatar")
|
|
|
+ candidates = []
|
|
|
+ for path, offset, size in cursor:
|
|
|
+ if avtid in path:
|
|
|
+ candidates.append((path, offset, size))
|
|
|
+ return candidates
|
|
|
|
|
|
if __name__ == '__main__':
|
|
|
import sys
|
Yuxin Wu