1
0
mirror of https://gist.github.com/6ba37e4d4084e858f917e271550ce5f6.git synced 2024-09-20 00:34:20 +03:00

Add cli arguments

This commit is contained in:
Victor 2021-04-15 14:21:36 +03:00
parent 00834d0e23
commit b528698ff7
4 changed files with 85 additions and 41 deletions

1
.gitignore vendored
View File

@ -1,4 +1,5 @@
__pycache__ __pycache__
.idea
input input
library library
logs logs

View File

@ -13,15 +13,13 @@ class Config:
def __init__(self, config): def __init__(self, config):
dirs = config.get('dirs', {}) dirs = config.get('dirs', {})
self.dir_tmp = Path(dirs.get('tmp', '/tmp/')) self.dir_tmp = Path(dirs.get('tmp', '/tmp/'))
self.dir_input = Path(dirs.get('input', './input')) self.dir_processed = Path(dirs.get('processed', './processed'))
self.dir_processed = Path(dirs.get('processed', './input/processed'))
self.dir_logs = Path(dirs.get('logs', './logs')) self.dir_logs = Path(dirs.get('logs', './logs'))
self.dir_library = Path(dirs.get('library', './library')) self.dir_library = Path(dirs.get('library', './library'))
self.__setup_folders() self.__setup_folders()
def __setup_folders(self): def __setup_folders(self):
self.dir_tmp.mkdir(exist_ok=True) self.dir_tmp.mkdir(exist_ok=True)
self.dir_processed.mkdir(exist_ok=True)
self.dir_logs.mkdir(exist_ok=True) self.dir_logs.mkdir(exist_ok=True)
self.dir_library.mkdir(exist_ok=True) self.dir_library.mkdir(exist_ok=True)

View File

@ -1,7 +1,6 @@
dirs: dirs:
tmp: /tmp/ tmp: /tmp/
input: ./input processed: ./processed
processed: ./input/processed
logs: ./logs logs: ./logs
library: ./library library: ./library

View File

@ -1,3 +1,4 @@
import argparse
import logging import logging
import os import os
import re import re
@ -5,6 +6,7 @@ import shutil
import time import time
from datetime import datetime from datetime import datetime
from pathlib import Path from pathlib import Path
from typing import Optional
from config import Config from config import Config
from database import Database from database import Database
@ -14,13 +16,32 @@ from metadata import Metadata
class PicSorter: class PicSorter:
def __init__(self): @staticmethod
self.config = Config.load('config.yml') def parse_args():
self.setup_logging() 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') filename = datetime.now().strftime('%Y-%m-%d.log')
logfile = Path(self.config.dir_logs, filename) logfile = Path(dir_logs, filename)
logging.basicConfig( logging.basicConfig(
filename=os.fspath(logfile), filename=os.fspath(logfile),
level=logging.INFO, level=logging.INFO,
@ -28,41 +49,66 @@ class PicSorter:
datefmt='%H:%M:%S', datefmt='%H:%M:%S',
) )
def process_folder(self): def process(self, inputs: list[str]) -> None:
config = self.config for input in inputs:
library = Library(config.dir_library) p = Path(input)
metadata = Metadata(config.dir_tmp) if p.is_dir():
db = Database() self.__process_folder(p)
files = {p for p in config.dir_input.iterdir() 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"]} if p.suffix in [".jpg", ".png"]}
for filename in files: for filename in files:
print("Process", filename) print("Processing", filename)
try: try:
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) url = Iqdb.search(filename)
if url is None: if url is None:
logging.warning("%s not found", filename) logging.warning("%s not found", filename)
library.move_to_orphan(Path(filename)) self.library.move_to_orphan(Path(filename))
continue return None
return url
def __process_url(self, url: str) -> bool:
m = re.search(r".*posts/(\d{3,})", url) m = re.search(r".*posts/(\d{3,})", url)
if not m: if not m:
continue return False
post_id = int(m.group(1)) post_id = int(m.group(1))
if db.is_exists(post_id): if self.db.is_exists(post_id):
logging.info("Skipping exists post %d", post_id) logging.info("Skipping exists post %d", post_id)
continue return False
meta_result = metadata.process(url) meta_result = self.metadata.process(url)
if meta_result is None: if meta_result is None:
continue return False
image_path, tags = meta_result image_path, tags = meta_result
library.move(image_path, tags) self.library.move(image_path, tags)
db.add(post_id, tags.tags_string) self.db.add(post_id, tags.tags_string)
shutil.move(os.fspath(filename), os.fspath(config.dir_processed)) return True
time.sleep(5)
except Exception as ex:
raise ex
if __name__ == '__main__': if __name__ == '__main__':
PicSorter().process_folder() PicSorter.parse_args()