mirror of
https://github.com/zhenxun-org/zhenxun_bot.git
synced 2026-10-02 10:10:01 +08:00
refactor code
This commit is contained in:
hibiki
+140
-107
@@ -1,127 +1,160 @@
|
||||
from configs.path_config import IMAGE_PATH, TXT_PATH
|
||||
import os
|
||||
from utils.user_agent import get_user_agent
|
||||
from services.log import logger
|
||||
from datetime import datetime
|
||||
from utils.img_utils import rar_imgs, get_img_hash
|
||||
from utils.image_utils import compressed_image, get_img_hash
|
||||
from utils.utils import get_bot, get_local_proxy
|
||||
from asyncio.exceptions import TimeoutError
|
||||
from models.setu import Setu
|
||||
from aiohttp.client_exceptions import ClientConnectorError
|
||||
from asyncpg.exceptions import UniqueViolationError
|
||||
from pathlib import Path
|
||||
from nonebot import Driver
|
||||
import nonebot
|
||||
import aiofiles
|
||||
import aiohttp
|
||||
from aiohttp.client_exceptions import ClientConnectorError
|
||||
import os
|
||||
import ujson as json
|
||||
|
||||
try:
|
||||
import ujson as json
|
||||
except ModuleNotFoundError:
|
||||
import json
|
||||
driver: Driver = nonebot.get_driver()
|
||||
|
||||
_path = Path(IMAGE_PATH)
|
||||
|
||||
|
||||
@driver.on_startup
|
||||
async def update_old_setu_data():
|
||||
path = Path(TXT_PATH)
|
||||
setu_data_file = path / "setu_data.json"
|
||||
r18_data_file = path / "r18_setu_data.json"
|
||||
index = 0
|
||||
r18_index = 0
|
||||
count = 0
|
||||
fail_count = 0
|
||||
for file in [setu_data_file, r18_data_file]:
|
||||
if file.exists():
|
||||
data = json.load(open(file, "r", encoding="utf8"))
|
||||
for x in data:
|
||||
if file == setu_data_file:
|
||||
idx = index
|
||||
if 'R-18' in data[x]["tags"]:
|
||||
data[x]["tags"].remove('R-18')
|
||||
else:
|
||||
idx = r18_index
|
||||
img_url = (
|
||||
data[x]["img_url"].replace("i.pixiv.cat", "i.pximg.net")
|
||||
if "i.pixiv.cat" in data[x]["img_url"]
|
||||
else data[x]["img_url"]
|
||||
)
|
||||
# idx = r18_index if 'R-18' in data[x]["tags"] else index
|
||||
try:
|
||||
await Setu.add_setu_data(
|
||||
idx,
|
||||
data[x]["title"],
|
||||
data[x]["author"],
|
||||
data[x]["pid"],
|
||||
data[x]["img_hash"],
|
||||
img_url,
|
||||
",".join(data[x]["tags"]),
|
||||
)
|
||||
count += 1
|
||||
if 'R-18' in data[x]["tags"]:
|
||||
r18_index += 1
|
||||
else:
|
||||
index += 1
|
||||
logger.info(f'添加旧色图数据成功 PID：{data[x]["pid"]} index：{idx}....')
|
||||
except UniqueViolationError:
|
||||
fail_count += 1
|
||||
logger.info(f'添加旧色图数据失败，色图重复 PID：{data[x]["pid"]} index：{idx}....')
|
||||
file.unlink()
|
||||
setu_url_path = path / "setu_url.json"
|
||||
setu_r18_url_path = path / "setu_r18_url.json"
|
||||
if setu_url_path.exists():
|
||||
setu_url_path.unlink()
|
||||
if setu_r18_url_path.exists():
|
||||
setu_r18_url_path.unlink()
|
||||
logger.info(f"更新旧色图数据完成，成功更新数据：{count} 条，累计失败：{fail_count} 条")
|
||||
|
||||
|
||||
headers = {
|
||||
"User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10.6;"
|
||||
" rv:2.0.1) Gecko/20100101 Firefox/4.0.1",
|
||||
"Referer": "https://www.pixiv.net",
|
||||
}
|
||||
|
||||
|
||||
async def update_setu_img():
|
||||
async with aiohttp.ClientSession(headers=get_user_agent()) as session:
|
||||
for file_name in ['setu_url.json', 'setu_r18_url.json']:
|
||||
if file_name == 'setu_url.json':
|
||||
json_name = 'setu_data.json'
|
||||
path = '_setu/'
|
||||
rar_path = 'setu_rar/'
|
||||
else:
|
||||
json_name = 'r18_setu_data.json'
|
||||
path = '_r18/'
|
||||
rar_path = 'r18_rar/'
|
||||
if not os.path.exists(IMAGE_PATH + path):
|
||||
os.mkdir(IMAGE_PATH + path)
|
||||
if not os.path.exists(IMAGE_PATH + rar_path):
|
||||
os.mkdir(IMAGE_PATH + rar_path)
|
||||
try:
|
||||
data = json.load(open(TXT_PATH + json_name, encoding='utf8'))
|
||||
if not data:
|
||||
data = {}
|
||||
except (FileNotFoundError, TypeError):
|
||||
data = {}
|
||||
_success = 0
|
||||
_similar = 0
|
||||
try:
|
||||
with open(TXT_PATH + file_name, 'r', encoding='utf8') as f:
|
||||
txt_data = json.load(f)
|
||||
if not txt_data:
|
||||
continue
|
||||
except (FileNotFoundError, ValueError):
|
||||
continue
|
||||
total = len(txt_data)
|
||||
urls = [data[x]['img_url'] for x in data.keys()]
|
||||
for pid in txt_data.keys():
|
||||
index = str(len(os.listdir(IMAGE_PATH + path)))
|
||||
url = txt_data[pid]["img_url"].replace('img-master', 'img-original').replace('_master1200', '')
|
||||
if url in urls or txt_data[pid]["img_url"] in urls:
|
||||
continue
|
||||
logger.info(f'开始更新 index:{index} --> {url}')
|
||||
image_list = await Setu.get_all_setu()
|
||||
image_list.reverse()
|
||||
_success = 0
|
||||
error_info = []
|
||||
error_type = []
|
||||
async with aiohttp.ClientSession(headers=headers) as session:
|
||||
for image in image_list:
|
||||
path = _path / "_r18" if image.is_r18 else _path / "_setu"
|
||||
rar_path = "r18_rar" if image.is_r18 else "rar"
|
||||
local_image = path / f"{image.local_id}.jpg"
|
||||
if not local_image.exists():
|
||||
for _ in range(3):
|
||||
try:
|
||||
async with session.get(url, proxy=get_local_proxy(), timeout=15) as response:
|
||||
async with session.get(
|
||||
image.img_url, proxy=get_local_proxy(), timeout=30
|
||||
) as response:
|
||||
if response.status == 200:
|
||||
async with aiofiles.open(IMAGE_PATH + rar_path + index + ".jpg", 'wb') as f:
|
||||
async with aiofiles.open(
|
||||
f"{IMAGE_PATH}/{rar_path}/{image.local_id}.jpg",
|
||||
"wb",
|
||||
) as f:
|
||||
await f.write(await response.read())
|
||||
_success += 1
|
||||
else:
|
||||
logger.info(f'{url} 不存在，使用更新原url')
|
||||
url = txt_data[pid]["img_url"]
|
||||
async with session.get(txt_data[pid]["img_url"], proxy=get_local_proxy(),
|
||||
timeout=15) as response:
|
||||
if response.status == 200:
|
||||
async with aiofiles.open(IMAGE_PATH + rar_path + index + ".jpg", 'wb') as f:
|
||||
await f.write(await response.read())
|
||||
_success += 1
|
||||
try:
|
||||
if os.path.getsize(IMAGE_PATH + rar_path + str(index) + ".jpg") > 1024 * 1024 * 1.5:
|
||||
rar_imgs(
|
||||
rar_path,
|
||||
path,
|
||||
in_file_name=index,
|
||||
out_file_name=index
|
||||
try:
|
||||
if (
|
||||
os.path.getsize(
|
||||
f"{IMAGE_PATH}/{rar_path}/{image.local_id}.jpg"
|
||||
)
|
||||
> 1024 * 1024 * 1.5
|
||||
):
|
||||
compressed_image(
|
||||
os.path.join(
|
||||
rar_path, f"{image.local_id}.jpg"
|
||||
),
|
||||
os.path.join(path, f"{image.local_id}.jpg"),
|
||||
)
|
||||
else:
|
||||
logger.info(
|
||||
f"不需要压缩，移动图片 {IMAGE_PATH}/{rar_path}/{image.local_id}.jpg "
|
||||
f"--> /{path}/{image.local_id}.jpg"
|
||||
)
|
||||
os.rename(
|
||||
f"{IMAGE_PATH}/{rar_path}/{image.local_id}.jpg",
|
||||
f"{path}/{image.local_id}.jpg",
|
||||
)
|
||||
except FileNotFoundError:
|
||||
logger.warning(f"文件 {image.local_id}.jpg 不存在，跳过...")
|
||||
_success -= 1
|
||||
continue
|
||||
img_hash = str(
|
||||
get_img_hash(
|
||||
f"{path}/{image.local_id}.jpg"
|
||||
)
|
||||
)
|
||||
else:
|
||||
logger.info('不需要压缩，移动图片 ' + IMAGE_PATH + rar_path + index + ".jpg --> "
|
||||
+ IMAGE_PATH + path + index + ".jpg")
|
||||
os.rename(IMAGE_PATH + rar_path + index + ".jpg",
|
||||
IMAGE_PATH + path + index + ".jpg")
|
||||
except FileNotFoundError:
|
||||
logger.warning(f'文件 {index}.jpg 不存在，跳过...')
|
||||
_success -= 1
|
||||
continue
|
||||
img_hash = str(get_img_hash(f'{IMAGE_PATH}{path}{index}.jpg'))
|
||||
if img_hash in [data[x]['img_hash'] for x in data.keys()]:
|
||||
logger.info(f'index:{index} 与 '
|
||||
f'{[data[x]["img_hash"] for x in data.keys()].index(img_hash)} 存在重复，删除')
|
||||
os.remove(IMAGE_PATH + path + index + ".jpg")
|
||||
_similar += 1
|
||||
else:
|
||||
data[index] = {
|
||||
'title': txt_data[pid]['title'],
|
||||
'author': txt_data[pid]['author'],
|
||||
'pid': txt_data[pid]['pid'],
|
||||
'img_hash': img_hash,
|
||||
'img_url': url,
|
||||
'tags': txt_data[pid]['tags'],
|
||||
}
|
||||
break
|
||||
await Setu.update_setu_data(
|
||||
image.pid, img_hash=img_hash
|
||||
)
|
||||
break
|
||||
except (TimeoutError, ClientConnectorError) as e:
|
||||
logger.warning(f'{url} 更新失败 ..{type(e)}：{e}')
|
||||
continue
|
||||
logger.warning(f"{image.local_id}.jpg 更新失败 ..{type(e)}：{e}")
|
||||
pass
|
||||
except Exception as e:
|
||||
await get_bot().send_private_msg(
|
||||
user_id=int(list(get_bot().config.superusers)[0]),
|
||||
message=f'更新 {index}.jpg 色图错误 {type(e)}: {e}'
|
||||
)
|
||||
_success -= 1
|
||||
logger.error(f'更新色图 {index}.jpg 错误 {type(e)}: {e}')
|
||||
continue
|
||||
with open(TXT_PATH + json_name, 'w', encoding='utf-8') as f:
|
||||
json.dump(data, f, indent=4, ensure_ascii=False)
|
||||
open(TXT_PATH + file_name, 'w')
|
||||
logger.info(
|
||||
f'{str(datetime.now()).split(".")[0]} 更新 {file_name.split(".")[0]}完成，预计更新 {total} 张，'
|
||||
f'实际更新 {_success} 张，相似 {_similar} 张，实际存入 {_success - _similar} 张')
|
||||
await get_bot().send_private_msg(
|
||||
user_id=int(list(get_bot().config.superusers)[0]),
|
||||
message=f'{str(datetime.now()).split(".")[0]} 更新{file_name.split(".")[0]}完成，预计更新 {total} 张，'
|
||||
f'实际更新 {_success} 张，相似 {_similar} 张，实际存入 {_success - _similar} 张'
|
||||
)
|
||||
logger.error(f"更新色图 {image.local_id}.jpg 错误 {type(e)}: {e}")
|
||||
if type(e) not in error_type:
|
||||
error_type.append(type(e))
|
||||
error_info.append(
|
||||
f"更新色图 {image.local_id}.jpg 错误 {type(e)}: {e}"
|
||||
)
|
||||
await get_bot().send_private_msg(
|
||||
user_id=int(list(get_bot().config.superusers)[0]),
|
||||
message=f'{str(datetime.now()).split(".")[0]} 更新 色图 完成，实际更新 {_success} 张，以下为更新时未知错误：\n'
|
||||
+ "\n".join(error_info),
|
||||
)
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user