other-projects/picsorter
1
0
Fork 0
mirror of https://gist.github.com/6ba37e4d4084e858f917e271550ce5f6.git synced 2024-09-20 00:34:20 +03:00
Code Activity

Add support for aibooru.online

Browse Source
This commit is contained in:
aNNiMON 2023-03-20 19:09:25 +02:00
parent 5a970dc2ba
commit aa40b93401
2 changed files with 20 additions and 15 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

20
database.py
View File

@ -11,27 +11,29 @@ class Database:
conn = sqlite3.connect(self.db_name) conn = sqlite3.connect(self.db_name)
c = conn.cursor() c = conn.cursor()
c.executescript(""" c.executescript("""
CREATE TABLE IF NOT EXISTS danbooru ( CREATE TABLE IF NOT EXISTS images (
id INTEGER PRIMARY KEY NOT NULL UNIQUE, id INTEGER PRIMARY KEY NOT NULL,
provider TEXT NOT NULL,
tags TEXT NOT NULL, tags TEXT NOT NULL,
created_at TIMESTAMP created_at TIMESTAMP,
); UNIQUE(id, provider) ON CONFLICT REPLACE
)
""") """)
conn.commit() conn.commit()
conn.close() conn.close()
def is_exists(self, _id) -> bool: def is_exists(self, _id, provider) -> bool:
conn = sqlite3.connect(self.db_name) conn = sqlite3.connect(self.db_name)
c = conn.cursor() c = conn.cursor()
c.execute("SELECT EXISTS(SELECT 1 FROM danbooru WHERE id=?)", (_id, )) c.execute("SELECT EXISTS(SELECT 1 FROM images WHERE id=? AND provider=?)", (_id, provider))
result = c.fetchone()[0] result = c.fetchone()[0]
conn.close() conn.close()
return bool(result) return bool(result)
def add(self, _id, tags): def add(self, _id, provider, tags):
conn = sqlite3.connect(self.db_name) conn = sqlite3.connect(self.db_name)
c = conn.cursor() c = conn.cursor()
sql = 'INSERT INTO danbooru(id, tags, created_at) VALUES (?,?,?)' sql = 'INSERT INTO images(id, provider, tags, created_at) VALUES (?,?,?,?)'
c.execute(sql, (_id, tags, datetime.now())) c.execute(sql, (_id, provider, tags, datetime.now()))
conn.commit() conn.commit()
conn.close() conn.close()

15
picsorter.py
View File

@ -95,20 +95,23 @@ class PicSorter:
return url return url
def __process_url(self, url: str) -> bool: def __process_url(self, url: str) -> bool:
m = re.search(r"(?:posts/)?(\d{3,})", url) m = re.search(r"(https://((?:dan|ai)booru|yande).*?(?:post(?:s|/show)/)?(\d{3,}))", url)
if not m: if not m:
return False return False
post_id = int(m.group(1)) provider = m.group(2)
if self.db.is_exists(post_id): post_id = int(m.group(3))
logging.info("Skipping exists post %d", post_id) if provider not in ['danbooru', 'aibooru']:
return False
if self.db.is_exists(provider, post_id):
logging.info("Skipping exists post %s %d", provider, post_id)
return False return False
meta_result = self.metadata.process("https://danbooru.donmai.us/posts/" + str(post_id)) meta_result = self.metadata.process(m.group(1))
if meta_result is None: if meta_result is None:
return False return False
image_path, tags = meta_result image_path, tags = meta_result
to_path = self.library.move(image_path, tags) to_path = self.library.move(image_path, tags)
self.db.add(post_id, tags.tags_string) self.db.add(post_id, provider, tags.tags_string)
self.__show_path(to_path) self.__show_path(to_path)
return True return True