icon_loader.py

# icon_loader.py — 从 Windows 文件提取图标
#
# 只依赖 ctypes（内置）和 Pillow，不需要 pywin32。
# 使用 SHGetFileInfoW 获取系统图标句柄，
# 通过 CreateDIBSection 渲染到内存位图再转 PIL Image。

import ctypes
from ctypes import wintypes
import os

from PIL import Image, ImageDraw, ImageFont, ImageTk

# ── Windows 结构 ──────────────────────────────────────────

class SHFILEINFOW(ctypes.Structure):
    _fields_ = [
        ('hIcon',        wintypes.HANDLE),
        ('iIcon',        ctypes.c_int),
        ('dwAttributes', wintypes.DWORD),
        ('szDisplayName', ctypes.c_wchar * 260),
        ('szTypeName',    ctypes.c_wchar * 80),
    ]

class BITMAPINFOHEADER(ctypes.Structure):
    _fields_ = [
        ('biSize',          wintypes.DWORD),
        ('biWidth',         wintypes.LONG),
        ('biHeight',        wintypes.LONG),
        ('biPlanes',        wintypes.WORD),
        ('biBitCount',      wintypes.WORD),
        ('biCompression',   wintypes.DWORD),
        ('biSizeImage',     wintypes.DWORD),
        ('biXPelsPerMeter', wintypes.LONG),
        ('biYPelsPerMeter', wintypes.LONG),
        ('biClrUsed',       wintypes.DWORD),
        ('biClrImportant',  wintypes.DWORD),
    ]

class BITMAPINFO(ctypes.Structure):
    _fields_ = [('bmiHeader', BITMAPINFOHEADER)]

# ── 缓存 ──────────────────────────────────────────────────

_icon_cache = {}

# ── 后备图标 ──────────────────────────────────────────────

_FALLBACK_COLORS = [
    (70, 130, 180),    # 钢蓝
    (60, 179, 113),    # 海洋绿
    (218, 165, 32),    # 金黄
    (205, 92, 92),     # 印度红
    (147, 112, 219),   # 中紫
    (0, 139, 139),     # 深青
    (210, 105, 30),    # 巧克力
]

def _fallback_icon(name, size=48):
    if not name:
        name = '?'
    letter = name[0].upper()
    color = _FALLBACK_COLORS[hash(name) % len(_FALLBACK_COLORS)]
    img = Image.new('RGBA', (size, size), (*color, 255))
    draw = ImageDraw.Draw(img)
    try:
        font = ImageFont.truetype('segoeui.ttf', size // 2)
    except (IOError, OSError):
        font = ImageFont.load_default()
    bbox = draw.textbbox((0, 0), letter, font=font)
    tw = bbox[2] - bbox[0]
    th = bbox[3] - bbox[1]
    x = (size - tw) // 2 - bbox[0]
    y = (size - th) // 2 - bbox[1]
    draw.text((x, y), letter, fill=(255, 255, 255, 255), font=font)
    return img

# ── 核心图标提取 ──────────────────────────────────────────

def _extract_pywin32_shgfi(path, size):
    try:
        import win32gui, win32con
        attrs = _get_file_attrs(path)
        hicon, _, _ = win32gui.SHGetFileInfo(path, attrs,
            win32con.SHGFI_ICON | win32con.SHGFI_LARGEICON | win32con.SHGFI_USEFILEATTRIBUTES)
        if hicon == 0:
            return None
        user32 = ctypes.windll.user32; gdi32 = ctypes.windll.gdi32
        hdc_screen = user32.GetDC(0); hdc_mem = gdi32.CreateCompatibleDC(hdc_screen)
        bmi = BITMAPINFO()
        bmi.bmiHeader.biSize = ctypes.sizeof(BITMAPINFOHEADER)
        bmi.bmiHeader.biWidth = size; bmi.bmiHeader.biHeight = -size
        bmi.bmiHeader.biPlanes = 1; bmi.bmiHeader.biBitCount = 32
        bmi.bmiHeader.biCompression = 0
        bits_ptr = ctypes.POINTER(ctypes.c_byte)()
        hbmp = gdi32.CreateDIBSection(hdc_screen, ctypes.byref(bmi), 0, ctypes.byref(bits_ptr), None, 0)
        if not hbmp:
            user32.DestroyIcon(hicon); gdi32.DeleteDC(hdc_mem)
            user32.ReleaseDC(0, hdc_screen); return None
        hbmp_old = gdi32.SelectObject(hdc_mem, hbmp)
        user32.DrawIconEx(hdc_mem, 0, 0, hicon, size, size, 0, None, 0x0003)
        n = size * size * 4
        buf = (ctypes.c_ubyte * n).from_address(ctypes.addressof(bits_ptr.contents))
        img = Image.frombuffer('RGBA', (size, size), bytes(buf), 'raw', 'BGRA', 0, 1)
        gdi32.SelectObject(hdc_mem, hbmp_old); gdi32.DeleteObject(hbmp)
        gdi32.DeleteDC(hdc_mem); user32.ReleaseDC(0, hdc_screen)
        win32gui.DestroyIcon(hicon)
        return img
    except Exception:
        return None

def get_file_icon(path, size=48):
    cache_key = (path.lower(), size)
    if cache_key in _icon_cache:
        return _icon_cache[cache_key]
    name = os.path.splitext(os.path.basename(path))[0]
    # 尝试三种方式提取图标 (含 USEFILEATTRIBUTES, 文件不存在也能按扩展名取)
    img = _extract_dib(path, size)
    if img is None:
        img = _extract_win32(path, size)
    if img is None:
        img = _extract_pywin32_shgfi(path, size)
    if img is None:
        img = _fallback_icon(name, size)
    _icon_cache[cache_key] = img
    return img

def _get_file_attrs(path):
    if os.path.isdir(path):
        return 0x10  # FILE_ATTRIBUTE_DIRECTORY
    return 0x80      # FILE_ATTRIBUTE_NORMAL

def _extract_dib(path, size):
    shell32 = ctypes.windll.shell32
    user32  = ctypes.windll.user32
    gdi32   = ctypes.windll.gdi32
    sfi = SHFILEINFOW()
    flags = 0x000000110  # SHGFI_ICON | SHGFI_USEFILEATTRIBUTES
    attrs = _get_file_attrs(path)
    ret = shell32.SHGetFileInfoW(path, attrs, ctypes.byref(sfi), ctypes.sizeof(sfi), flags)
    if ret == 0 or sfi.hIcon == 0:
        return None
    hicon = sfi.hIcon
    hdc_screen = user32.GetDC(0)
    hdc_mem    = gdi32.CreateCompatibleDC(hdc_screen)
    bmi = BITMAPINFO()
    bmi.bmiHeader.biSize = ctypes.sizeof(BITMAPINFOHEADER)
    bmi.bmiHeader.biWidth = size; bmi.bmiHeader.biHeight = -size
    bmi.bmiHeader.biPlanes = 1; bmi.bmiHeader.biBitCount = 32
    bmi.bmiHeader.biCompression = 0
    bits_ptr = ctypes.POINTER(ctypes.c_byte)()
    hbmp = gdi32.CreateDIBSection(hdc_screen, ctypes.byref(bmi), 0, ctypes.byref(bits_ptr), None, 0)
    if not hbmp:
        user32.DestroyIcon(hicon); gdi32.DeleteDC(hdc_mem)
        user32.ReleaseDC(0, hdc_screen); return None
    hbmp_old = gdi32.SelectObject(hdc_mem, hbmp)
    user32.DrawIconEx(hdc_mem, 0, 0, hicon, size, size, 0, None, 0x0003)
    n = size * size * 4
    buf = (ctypes.c_ubyte * n).from_address(ctypes.addressof(bits_ptr.contents))
    img = Image.frombuffer('RGBA', (size, size), bytes(buf), 'raw', 'BGRA', 0, 1)
    gdi32.SelectObject(hdc_mem, hbmp_old); gdi32.DeleteObject(hbmp)
    gdi32.DeleteDC(hdc_mem); user32.ReleaseDC(0, hdc_screen)
    user32.DestroyIcon(hicon)
    return img

def _extract_win32(path, size):
    try:
        import win32gui
        large, _ = win32gui.ExtractIconEx(path, 0)
        if not large:
            return None
        hicon = large[0]
        user32 = ctypes.windll.user32
        gdi32 = ctypes.windll.gdi32
        hdc_screen = user32.GetDC(0)
        hdc_mem = gdi32.CreateCompatibleDC(hdc_screen)
        bmi = BITMAPINFO()
        bmi.bmiHeader.biSize = ctypes.sizeof(BITMAPINFOHEADER)
        bmi.bmiHeader.biWidth = size; bmi.bmiHeader.biHeight = -size
        bmi.bmiHeader.biPlanes = 1; bmi.bmiHeader.biBitCount = 32
        bmi.bmiHeader.biCompression = 0
        bits_ptr = ctypes.POINTER(ctypes.c_byte)()
        hbmp = gdi32.CreateDIBSection(hdc_screen, ctypes.byref(bmi), 0, ctypes.byref(bits_ptr), None, 0)
        if not hbmp:
            user32.DestroyIcon(hicon); gdi32.DeleteDC(hdc_mem)
            user32.ReleaseDC(0, hdc_screen); return None
        hbmp_old = gdi32.SelectObject(hdc_mem, hbmp)
        user32.DrawIconEx(hdc_mem, 0, 0, hicon, size, size, 0, None, 0x0003)
        n = size * size * 4
        buf = (ctypes.c_ubyte * n).from_address(ctypes.addressof(bits_ptr.contents))
        img = Image.frombuffer('RGBA', (size, size), bytes(buf), 'raw', 'BGRA', 0, 1)
        gdi32.SelectObject(hdc_mem, hbmp_old); gdi32.DeleteObject(hbmp)
        gdi32.DeleteDC(hdc_mem); user32.ReleaseDC(0, hdc_screen)
        win32gui.DestroyIcon(hicon)
        return img
    except Exception:
        return None

def as_photoimage(path, size=48):
    img = get_file_icon(path, size)
    if img is None:
        return None
    if img.mode == 'RGBA':
        bg = Image.new('RGB', img.size, (43, 43, 43))
        try:
            bg.paste(img, mask=img.split()[3])
        except ValueError:
            bg.paste(img)
        img = bg
    return ImageTk.PhotoImage(img)