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

@ -5,7 +5,7 @@ import yaml
class Config: class Config:
@staticmethod @staticmethod
def load(path = 'config.yml'): def load(path='config.yml'):
with open(path, 'rt', encoding='utf8') as f: with open(path, 'rt', encoding='utf8') as f:
config = yaml.load(f.read(), Loader=yaml.FullLoader) config = yaml.load(f.read(), Loader=yaml.FullLoader)
return Config(config) return Config(config)
@ -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:
url = Iqdb.search(filename) self.__process_file(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)
except Exception as ex: except Exception as ex:
raise 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__': if __name__ == '__main__':
PicSorter().process_folder() PicSorter.parse_args()