mirror of
https://gist.github.com/6ba37e4d4084e858f917e271550ce5f6.git
synced 2024-09-20 00:34:20 +03:00
Add support for aibooru.online
This commit is contained in:
parent
5a970dc2ba
commit
aa40b93401
20
database.py
20
database.py
@ -11,27 +11,29 @@ class Database:
|
||||
conn = sqlite3.connect(self.db_name)
|
||||
c = conn.cursor()
|
||||
c.executescript("""
|
||||
CREATE TABLE IF NOT EXISTS danbooru (
|
||||
id INTEGER PRIMARY KEY NOT NULL UNIQUE,
|
||||
CREATE TABLE IF NOT EXISTS images (
|
||||
id INTEGER PRIMARY KEY NOT NULL,
|
||||
provider TEXT NOT NULL,
|
||||
tags TEXT NOT NULL,
|
||||
created_at TIMESTAMP
|
||||
);
|
||||
created_at TIMESTAMP,
|
||||
UNIQUE(id, provider) ON CONFLICT REPLACE
|
||||
)
|
||||
""")
|
||||
conn.commit()
|
||||
conn.close()
|
||||
|
||||
def is_exists(self, _id) -> bool:
|
||||
def is_exists(self, _id, provider) -> bool:
|
||||
conn = sqlite3.connect(self.db_name)
|
||||
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]
|
||||
conn.close()
|
||||
return bool(result)
|
||||
|
||||
def add(self, _id, tags):
|
||||
def add(self, _id, provider, tags):
|
||||
conn = sqlite3.connect(self.db_name)
|
||||
c = conn.cursor()
|
||||
sql = 'INSERT INTO danbooru(id, tags, created_at) VALUES (?,?,?)'
|
||||
c.execute(sql, (_id, tags, datetime.now()))
|
||||
sql = 'INSERT INTO images(id, provider, tags, created_at) VALUES (?,?,?,?)'
|
||||
c.execute(sql, (_id, provider, tags, datetime.now()))
|
||||
conn.commit()
|
||||
conn.close()
|
||||
|
15
picsorter.py
15
picsorter.py
@ -95,20 +95,23 @@ class PicSorter:
|
||||
return url
|
||||
|
||||
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:
|
||||
return False
|
||||
post_id = int(m.group(1))
|
||||
if self.db.is_exists(post_id):
|
||||
logging.info("Skipping exists post %d", post_id)
|
||||
provider = m.group(2)
|
||||
post_id = int(m.group(3))
|
||||
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
|
||||
|
||||
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:
|
||||
return False
|
||||
image_path, tags = meta_result
|
||||
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)
|
||||
return True
|
||||
|
||||
|
Loading…
Reference in New Issue
Block a user