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