feat: 第三方依赖

This commit is contained in:
hp0912 2026-04-04 22:54:14 +08:00
parent 76d9d9e447
commit 6438efee40

View File

@ -5,23 +5,37 @@ from __future__ import annotations
import json import json
import os import os
import re import re
import subprocess
import sys import sys
import time import time
import urllib.error
import urllib.request import urllib.request
FALLBACK_TEXT = "AI 绘图暂时不可用,请稍后再试。" # ---------------------------------------------------------------------------
# Dependency bootstrap
# ---------------------------------------------------------------------------
def _ensure_package(pip_name: str, import_name: str | None = None) -> None:
"""Install *pip_name* if it cannot be imported."""
try:
__import__(import_name or pip_name)
except ImportError:
subprocess.check_call(
[sys.executable, "-m", "pip", "install", pip_name, "-q"],
stdout=subprocess.DEVNULL,
)
_ensure_package("pymysql")
import pymysql # type: ignore # noqa: E402
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
# Database helpers (pure stdlib, no third-party MySQL driver) # Database helpers
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
def _mysql_connect(): def _mysql_connect():
"""Return a simple MySQL connection via PyMySQL (stdlib-compatible pure-Python driver)
or fall back to mysql.connector. We import lazily so the script still loads
even if neither is installed (it will just raise at call-time)."""
host = os.environ.get("MYSQL_HOST", "127.0.0.1") host = os.environ.get("MYSQL_HOST", "127.0.0.1")
port = int(os.environ.get("MYSQL_PORT", "3306")) port = int(os.environ.get("MYSQL_PORT", "3306"))
user = os.environ.get("MYSQL_USER", "root") user = os.environ.get("MYSQL_USER", "root")
@ -30,27 +44,11 @@ def _mysql_connect():
if not database: if not database:
raise RuntimeError("环境变量 ROBOT_CODE 未配置") raise RuntimeError("环境变量 ROBOT_CODE 未配置")
try: return pymysql.connect(
import pymysql # type: ignore host=host, port=port, user=user, password=password,
return pymysql.connect( database=database, charset="utf8mb4",
host=host, port=port, user=user, password=password, connect_timeout=10, read_timeout=30,
database=database, charset="utf8mb4", )
connect_timeout=10, read_timeout=30,
)
except ImportError:
pass
try:
import mysql.connector # type: ignore
return mysql.connector.connect(
host=host, port=port, user=user, password=password,
database=database, charset="utf8mb4",
connection_timeout=10,
)
except ImportError:
pass
raise RuntimeError("需要安装 pymysql 或 mysql-connector-python: pip install pymysql")
def _query_one(conn, sql: str, params: tuple = ()) -> dict | None: def _query_one(conn, sql: str, params: tuple = ()) -> dict | None:
@ -287,13 +285,13 @@ ZIMAGE_MODELS = {"Z-Image", "Z-Image-Turbo", "Qwen-Image-Edit-2511"}
def main() -> int: def main() -> int:
# Parse input params from first CLI argument # Parse input params from first CLI argument
if len(sys.argv) < 2: if len(sys.argv) < 2:
sys.stdout.write(FALLBACK_TEXT + "\n") sys.stdout.write("缺少输入参数\n")
return 1 return 1
try: try:
params = json.loads(sys.argv[1]) params = json.loads(sys.argv[1])
except json.JSONDecodeError: except json.JSONDecodeError as exc:
sys.stdout.write(FALLBACK_TEXT + "\n") sys.stdout.write(f"参数格式错误: {exc}\n")
return 1 return 1
prompt = params.get("prompt", "").strip() prompt = params.get("prompt", "").strip()
@ -315,16 +313,14 @@ def main() -> int:
try: try:
conn = _mysql_connect() conn = _mysql_connect()
except Exception as exc: except Exception as exc:
sys.stderr.write(f"数据库连接失败: {exc}\n") sys.stdout.write(f"数据库连接失败: {exc}\n")
sys.stdout.write(FALLBACK_TEXT + "\n")
return 1 return 1
try: try:
enabled, settings_json = load_drawing_settings(conn, from_wx_id) enabled, settings_json = load_drawing_settings(conn, from_wx_id)
except Exception as exc: except Exception as exc:
conn.close() conn.close()
sys.stderr.write(f"加载绘图配置失败: {exc}\n") sys.stdout.write(f"加载绘图配置失败: {exc}\n")
sys.stdout.write(FALLBACK_TEXT + "\n")
return 1 return 1
finally: finally:
try: try:
@ -370,8 +366,7 @@ def main() -> int:
return 1 return 1
except Exception as exc: except Exception as exc:
sys.stderr.write(f"调用绘图接口失败: {exc}\n") sys.stdout.write(f"调用绘图接口失败: {exc}\n")
sys.stdout.write(FALLBACK_TEXT + "\n")
return 1 return 1
if not image_urls: if not image_urls: