Spaces:
Runtime error
Runtime error
File size: 1,379 Bytes
4f8ad24 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 |
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
|