浏览代码

Fix Python 3.13 compatibility (#106)

* compat: support Python 3.13 (imghdr removal)

* compat: drop unnecessary future annotations import
CK Zhang 2 月之前
父节点
当前提交
11d60a7487
共有 3 个文件被更改,包括 59 次插入8 次删除
  1. 51 0
      wechat/common/imgutil.py
  2. 5 5
      wechat/emoji.py
  3. 3 3
      wechat/res.py

+ 51 - 0
wechat/common/imgutil.py

@@ -0,0 +1,51 @@
+import io
+from pathlib import Path
+
+from PIL import Image
+
+try:
+    import imghdr as _imghdr
+except ModuleNotFoundError:  # Python 3.13+
+    _imghdr = None
+
+
+def what(file: str | Path | None = None, h: bytes | None = None) -> str | None:
+    """
+    Compatibility wrapper for the removed stdlib `imghdr` module (Python 3.13+).
+
+    Behaves like `imghdr.what(file, h)` and returns a lowercase format string
+    like "jpeg", "png", "gif", or None if unknown.
+    """
+    if _imghdr is not None:
+        return _imghdr.what(file, h)
+    if h is not None:
+        return _what_from_bytes(h)
+    if file is None:
+        return None
+    try:
+        with Image.open(str(file)) as im:
+            fmt = im.format
+    except Exception:
+        return None
+    return fmt.lower() if fmt else None
+
+
+def _what_from_bytes(data: bytes) -> str | None:
+    if data.startswith(b"\xff\xd8\xff"):
+        return "jpeg"
+    if data.startswith(b"\x89PNG\r\n\x1a\n"):
+        return "png"
+    if data[:6] in (b"GIF87a", b"GIF89a"):
+        return "gif"
+    if data.startswith(b"BM"):
+        return "bmp"
+    if len(data) >= 12 and data.startswith(b"RIFF") and data[8:12] == b"WEBP":
+        return "webp"
+    if data.startswith(b"II*\x00") or data.startswith(b"MM\x00*"):
+        return "tiff"
+    try:
+        with Image.open(io.BytesIO(data)) as im:
+            fmt = im.format
+    except Exception:
+        return None
+    return fmt.lower() if fmt else None

+ 5 - 5
wechat/emoji.py

@@ -4,13 +4,13 @@ import logging
 import io
 import requests
 import base64
-import imghdr
 from PIL import Image
 import pickle
 from Crypto.Cipher import AES
 
 from .wxgf import WxgfAndroidDecoder, is_wxgf_buffer
 from .parser import WeChatDBParser
+from .common.imgutil import what as img_what
 from .common.textutil import md5 as get_md5_hex, get_file_b64, get_file_md5
 
 
@@ -124,10 +124,10 @@ class EmojiReader:
                 candidates = []
 
         def get_data_no_fallback(fname):
-            if imghdr.what(fname):
+            if img_what(fname):
                 data_md5 = get_file_md5(fname)
                 if data_md5 == md5:
-                    return get_file_b64(fname), imghdr.what(fname)
+                    return get_file_b64(fname), img_what(fname)
 
             try:
                 content = self._decode_emoji(fname)
@@ -147,9 +147,9 @@ class EmojiReader:
                 logger.error(f"Error decoding emoji {fname} : {str(e)}")
 
         def get_data_fallback(fname):
-            if not imghdr.what(fname):
+            if not img_what(fname):
                 return  # fallback files are not encrypted
-            return get_file_b64(fname), imghdr.what(fname)
+            return get_file_b64(fname), img_what(fname)
 
         get_data_func = get_data_fallback if allow_fallback else get_data_no_fallback
         results = [(x, get_data_func(x)) for x in candidates]

+ 3 - 3
wechat/res.py

@@ -8,12 +8,12 @@ import io
 import base64
 import logging
 logger = logging.getLogger(__name__)
-import imghdr
 from multiprocessing import Pool
 import atexit
 
 from .emoji import EmojiReader
 from .avatar import AvatarReader
+from .common.imgutil import what as img_what
 from .common.textutil import md5 as get_md5_hex, get_file_b64
 from .msg import TYPE_SPEAK
 from .audio import parse_wechat_audio_file
@@ -167,7 +167,7 @@ class Resource(object):
 
             # True jpeg. Simplest case.
             if img_file.endswith('jpg') and \
-                   imghdr.what(img_file) == 'jpeg':
+                   img_what(img_file) == 'jpeg':
                 return get_file_b64(img_file)
 
             if is_wxgf_file(img_file):
@@ -188,7 +188,7 @@ class Resource(object):
                     buf = f.read()
 
             # File is not actually jpeg. Convert.
-            if imghdr.what(file=None, h=buf) != 'jpeg':
+            if img_what(file=None, h=buf) != 'jpeg':
                 try:
                     im = Image.open(io.BytesIO(buf))
                 except: