diff --git a/.gitignore b/.gitignore index 5d62c19..8516963 100644 --- a/.gitignore +++ b/.gitignore @@ -1,4 +1,5 @@ __pycache__ +.idea input library logs diff --git a/config.py b/config.py index 82ee122..450d103 100644 --- a/config.py +++ b/config.py @@ -5,7 +5,7 @@ import yaml class Config: @staticmethod - def load(path = 'config.yml'): + def load(path='config.yml'): with open(path, 'rt', encoding='utf8') as f: config = yaml.load(f.read(), Loader=yaml.FullLoader) return Config(config) @@ -13,15 +13,13 @@ class Config: def __init__(self, config): dirs = config.get('dirs', {}) self.dir_tmp = Path(dirs.get('tmp', '/tmp/')) - self.dir_input = Path(dirs.get('input', './input')) - self.dir_processed = Path(dirs.get('processed', './input/processed')) + self.dir_processed = Path(dirs.get('processed', './processed')) self.dir_logs = Path(dirs.get('logs', './logs')) self.dir_library = Path(dirs.get('library', './library')) self.__setup_folders() def __setup_folders(self): self.dir_tmp.mkdir(exist_ok=True) - self.dir_processed.mkdir(exist_ok=True) self.dir_logs.mkdir(exist_ok=True) self.dir_library.mkdir(exist_ok=True) diff --git a/config.yml b/config.yml index c0a7458..ff0b394 100644 --- a/config.yml +++ b/config.yml @@ -1,7 +1,6 @@ dirs: tmp: /tmp/ - input: ./input - processed: ./input/processed + processed: ./processed logs: ./logs library: ./library diff --git a/picsorter.py b/picsorter.py index eeb1055..7d40847 100644 --- a/picsorter.py +++ b/picsorter.py @@ -1,3 +1,4 @@ +import argparse import logging import os import re @@ -5,6 +6,7 @@ import shutil import time from datetime import datetime from pathlib import Path +from typing import Optional from config import Config from database import Database @@ -14,13 +16,32 @@ from metadata import Metadata class PicSorter: - def __init__(self): - self.config = Config.load('config.yml') - self.setup_logging() + @staticmethod + def parse_args(): + parser = argparse.ArgumentParser( + description='Finds an image on danbooru, writes tags as IPTC keywords, than places the image in the library' + ) + parser.add_argument('-c', '--config', + type=Path, + default='config.yml', + help='config.yml file path') + parser.add_argument('input', nargs=argparse.REMAINDER) + args = parser.parse_args() + if len(args.input) >= 1: + PicSorter(args.config).process(args.input) - def setup_logging(self): + def __init__(self, config_file='config.yml'): + config = Config.load(config_file) + self.config = config + self.__setup_logging(config.dir_logs) + self.library = Library(config.dir_library) + self.metadata = Metadata(config.dir_tmp) + self.db = Database() + + @staticmethod + def __setup_logging(dir_logs: Path): filename = datetime.now().strftime('%Y-%m-%d.log') - logfile = Path(self.config.dir_logs, filename) + logfile = Path(dir_logs, filename) logging.basicConfig( filename=os.fspath(logfile), level=logging.INFO, @@ -28,41 +49,66 @@ class PicSorter: datefmt='%H:%M:%S', ) - def process_folder(self): - config = self.config - library = Library(config.dir_library) - metadata = Metadata(config.dir_tmp) - db = Database() - files = {p for p in config.dir_input.iterdir() + def process(self, inputs: list[str]) -> None: + for input in inputs: + p = Path(input) + if p.is_dir(): + self.__process_folder(p) + elif p.is_file(): + print("Processing file", input) + self.__process_file(input) + else: + print("Processing url", input) + self.__process_url(input) + + def __process_folder(self, dir_input: Path) -> None: + files = {p for p in dir_input.iterdir() if p.suffix in [".jpg", ".png"]} for filename in files: - print("Process", filename) + print("Processing", filename) try: - url = Iqdb.search(filename) - if url is None: - logging.warning("%s not found", filename) - library.move_to_orphan(Path(filename)) - continue - - m = re.search(r".*posts/(\d{3,})", url) - if not m: - continue - post_id = int(m.group(1)) - if db.is_exists(post_id): - logging.info("Skipping exists post %d", post_id) - continue - - meta_result = metadata.process(url) - if meta_result is None: - continue - image_path, tags = meta_result - library.move(image_path, tags) - db.add(post_id, tags.tags_string) - shutil.move(os.fspath(filename), os.fspath(config.dir_processed)) - time.sleep(5) + self.__process_file(filename) except Exception as ex: raise ex + time.sleep(5) + + def __process_file(self, filename: str) -> bool: + url = self.__search_iqdb(filename) + if url is None: + return False + if self.__process_url(url): + self.config.dir_processed.mkdir(exist_ok=True, parents=True) + from_path = os.fspath(filename) + to_path = os.fspath(self.config.dir_processed) + shutil.move(from_path, to_path) + return True + return False + + def __search_iqdb(self, filename: str) -> Optional[str]: + url = Iqdb.search(filename) + if url is None: + logging.warning("%s not found", filename) + self.library.move_to_orphan(Path(filename)) + return None + return url + + def __process_url(self, url: str) -> bool: + m = re.search(r".*posts/(\d{3,})", url) + if not m: + return False + post_id = int(m.group(1)) + if self.db.is_exists(post_id): + logging.info("Skipping exists post %d", post_id) + return False + + meta_result = self.metadata.process(url) + if meta_result is None: + return False + image_path, tags = meta_result + self.library.move(image_path, tags) + self.db.add(post_id, tags.tags_string) + return True if __name__ == '__main__': - PicSorter().process_folder() + PicSorter.parse_args()