waifuc/action/filename.py

import os
from typing import Iterator, Optional

from .base import BaseAction
from ..model import ImageItem


class FileExtAction(BaseAction):
    def __init__(self, ext: str):
        self.ext = ext
        self.untitles = 0

    def iter(self, item: ImageItem) -> Iterator[ImageItem]:
        if 'filename' in item.meta:
            filebody, _ = os.path.splitext(item.meta['filename'])
            filename = f'{filebody}{self.ext}'
        else:
            self.untitles += 1
            filename = f'untitled_{self.untitles}{self.ext}'

        yield ImageItem(item.image, {**item.meta, 'filename': filename})

    def reset(self):
        self.untitles = 0


class FileOrderAction(BaseAction):
    def __init__(self, ext: Optional[str] = '.png'):
        self.ext = ext
        self._current = 0

    def iter(self, item: ImageItem) -> Iterator[ImageItem]:
        self._current += 1
        if 'filename' in item.meta:
            _, ext = os.path.splitext(item.meta['filename'])
            new_filename = f'{self._current}{self.ext or ext}'
        else:
            if not self.ext:
                raise ValueError('No extension name provided for unnamed file.')
            else:
                new_filename = f'{self._current}{self.ext}'

        yield ImageItem(item.image, {**item.meta, 'filename': new_filename})

    def reset(self):
        self._current = 0