waifuc/model/item.py

import json
import os.path
import pickle
from dataclasses import dataclass
from typing import Optional

from PIL import Image
from hbutils.encoding import base64_decode, base64_encode
from hbutils.reflection import quick_import_object

NoneType = type(None)

_TYPE_META = '__type'
_BASE64_META = 'base64'


def load_meta(data, path=()):
    if isinstance(data, (int, float, str, NoneType)):
        return data
    elif isinstance(data, list):
        return [load_meta(item, (*path, i)) for i, item in enumerate(data)]
    elif isinstance(data, dict):
        if _TYPE_META not in data:
            return {key: load_meta(value, (*path, key)) for key, value in data.items()}
        else:
            cls, _, _ = quick_import_object(data[_TYPE_META])
            binary = base64_decode(data[_BASE64_META])
            obj = pickle.loads(binary)
            if isinstance(obj, cls):
                return obj
            else:
                raise TypeError(f'{cls!r} expected but {obj!r} found at {path!r}.')
    else:
        raise TypeError(f'Unknown type {data!r} at {path!r}.')


def dump_meta(data, path=()):
    if isinstance(data, (int, float, str, NoneType)):
        return data
    elif isinstance(data, list):
        return [dump_meta(item, (*path, i)) for i, item in enumerate(data)]
    elif isinstance(data, dict):
        return {key: dump_meta(value, (*path, key)) for key, value in data.items()}
    else:
        cls = type(data)
        type_str = f'{cls.__module__}.{cls.__name__}' if hasattr(cls, '__module__') else cls.__name__
        base64_str = base64_encode(pickle.dumps(data))
        return {
            _TYPE_META: type_str,
            _BASE64_META: base64_str
        }


@dataclass
class ImageItem:
    image: Image.Image
    meta: dict

    def __init__(self, image: Image.Image, meta: Optional[dict] = None):
        self.image = image
        self.meta = meta or {}

    @classmethod
    def _image_file_to_meta_file(cls, image_file):
        directory, filename = os.path.split(image_file)
        filebody, _ = os.path.splitext(filename)
        meta_file = os.path.join(directory, f'.{filebody}_meta.json')
        return meta_file

    @classmethod
    def load_from_image(cls, image_file):
        image = Image.open(image_file)
        meta_file = cls._image_file_to_meta_file(image_file)

        if os.path.exists(meta_file):
            with open(meta_file, 'r', encoding='utf-8') as f:
                meta = load_meta(json.load(f))
        else:
            meta = {}

        return cls(image, meta)

    def save(self, image_file, no_meta: bool = False, skip_when_image_exist: bool = False):
        if not skip_when_image_exist or not os.path.exists(image_file):
            self.image.save(image_file)
        if not no_meta and self.meta:
            meta_file = self._image_file_to_meta_file(image_file)
            with open(meta_file, 'w', encoding='utf-8') as f:
                json.dump(dump_meta(self.meta), f)

    def __repr__(self):
        values = {'size': self.image.size}
        for key, value in self.meta.items():
            if isinstance(value, (int, float, str)):
                values[key] = value

        content = ', '.join(f'{key}: {values[key]!r}' for key in sorted(values.keys()))
        return f'<{self.__class__.__name__} {content}>'