KikkaShortcut/Kikka/Utils/IconHelper.py

import os
import sys
import json
import hashlib
from icoextract import IconExtractor

from PySide6.QtCore import Qt, QFileInfo, QSize
from PySide6.QtGui import QPixmap, QIcon, QImage
from PySide6.QtWidgets import QFileIconProvider

from Kikka.KikkaConst import *
from Kikka.Utils.Singleton import singleton


@singleton
class IconHelper():
    
    def __init__(self):
        self._cache = {}
        self._icons = {}
        self.loadCache()
        
    def loadCache(self):
        os.makedirs(CACHE_ICON_PATH, exist_ok=True)
        if os.path.exists(ICON_CACHE_FILE):
            with open(ICON_CACHE_FILE, 'r') as fp:
                self._cache = json.load(fp)
            
    def saveCache(self):
        with open(ICON_CACHE_FILE, 'w') as fp:
             json.dump(self._cache, fp)
    
    def hash(self, text):
        return hashlib.sha256(text.encode('utf-8')).hexdigest()
    
    def getIconByHash(self, hash):
        if hash in self._cache.values():
            if hash in self._icons:
                return self._icons[hash]
            else:
                icon = QIcon(os.path.join(CACHE_ICON_PATH, hash + ".png"))
                self._icons[hash] = icon
                return icon
        return None
        
    def getIconByFile(self, filepath):
        hash = self.hash(filepath)
        if filepath in self._cache:
            if hash in self._icons:
                return self._icons[hash]
            else:
                icon_path = os.path.join(CACHE_ICON_PATH, hash + ".png")
                if os.path.exists(icon_path):
                    icon = QIcon(icon_path)
                    self._icons[hash] = icon
                    return icon
                
        file = QFileInfo(filepath)
        if not file.exists():
            return None
        
        icon = None
        if file.suffix().lower() in ["exe", "dll", "mun"]:
            temp_icon = os.path.join(CACHE_ICON_PATH, "_temp.ico")
            extractor = IconExtractor(file.filePath())
            extractor.export_icon(temp_icon)
            icon = QIcon(temp_icon)
            available_sizes = icon.availableSizes()
            if not available_sizes:
                return None
        else:
            p = QFileIconProvider()
            icon = p.icon(file)
        
        if icon is None or icon.isNull():
            return None
            
        # for size in available_sizes:
        #     pixmap = icon.pixmap(size)
        #     name = "%dx%s.png" % (size.width(), size.height())
        #     pixmap.save(os.path.join(CacheDir, name))
        
        actualSize = icon.actualSize(QSize(72, 72))
        pixmap = icon.pixmap(actualSize)
        pixmap.save(os.path.join(CACHE_ICON_PATH, hash + ".png"))
        self._cache[filepath] = hash
        return self.getIconByHash(hash)
    

if __name__ == "__main__":
    from PySide6.QtWidgets import QApplication
    app = QApplication(sys.argv)
    
    file = r"C:\example.exe"
    # file = C:\example.svg"
    # file = C:\example.lnk"
    icon = IconHelper().getIconByFile(file)
    pixmap = icon.pixmap(64, 64)
    pixmap.save("icon.png")