常用代码片段

MySql 操作

import pymysql.cursors
from loguru import logger
from dbutils.pooled_db import PooledDB


class DBPool:
    def __init__(self, **kwargs):
        self.NAME = kwargs.get("NAME")
        self.HOST = kwargs.get("HOST")
        self.PORT = kwargs.get("PORT")
        self.USER = kwargs.get("USER")
        self.PASSWORD = kwargs.get("PASSWORD")
        self.return_dict = kwargs.get("return_dict", False)

        # 断言配置信息完整
        check_list = [x for x in [self.HOST, self.PORT, self.NAME, self.USER] if not x]
        assert not check_list, "配置信息不正确！"

        # 使用字典游标，返回字典类型的查询结果
        self.pool = PooledDB(
            mincached=10, creator=pymysql, host=self.HOST, port=int(self.PORT),
            user=self.USER, password=self.PASSWORD, database=self.NAME,
            cursorclass=pymysql.cursors.DictCursor if self.return_dict else pymysql.cursors.Cursor
        )

    def get_connection(self):
        """获取数据库连接"""
        conn = self.pool.connection()
        if conn.open:
            return conn
        else:
            logger.error("数据库连接失效，重新获取连接...")
            return self.pool.connection()

    def close_connection(self):
        """关闭数据库连接池中的所有连接"""
        try:
            self.pool.close()  # PooledDB 已经提供了关闭连接池的方法
            logger.info("数据库连接池已关闭")
        except Exception as e:
            logger.error(f"关闭数据库连接池失败: {e}")

    def fetch_one(self, sql, args=None):
        """查询一条数据"""
        try:
            with self.get_connection() as conn:
                with conn.cursor() as cursor:
                    cursor.execute(sql, args)
                    return cursor.fetchone()
        except Exception as e:
            logger.error(f"执行查询失败: {e}")
            return None

    def fetch_all(self, sql, args=None):
        """查询多条数据"""
        try:
            with self.get_connection() as conn:
                with conn.cursor() as cursor:
                    cursor.execute(sql, args)
                    return cursor.fetchall()
        except Exception as e:
            logger.error(f"执行查询失败: {e}")
            return None

    def exe_sql(self, sql, args=None):
        """执行单条 SQL 语句"""
        try:
            with self.get_connection() as conn:
                with conn.cursor() as cursor:
                    cursor.execute(sql, args)
                    conn.commit()
                    return True
        except Exception as e:
            logger.error(f"执行 SQL 失败: {e}")
            return False

    def exe_many_sql(self, sql, args):
        """执行多条 SQL 语句"""
        try:
            with self.get_connection() as conn:
                with conn.cursor() as cursor:
                    cursor.execute('START TRANSACTION')
                    cursor.executemany(sql, args)
                    conn.commit()
                    return True
        except Exception as e:
            if 'Duplicate' in str(e):
                logger.warning(f"重复数据错误: {e}")
            else:
                logger.error(f"执行多条 SQL 失败: {e}")
            conn.rollback()
            return False


config = {
    'ENGINE': 'pymysql',
    'NAME': 'db_name',
    'USER': 'root',
    'PASSWORD': 'root',
    'HOST': '127.0.0.1',
    'PORT': '3306',
    'return_dict': True
}

工具类操作

def insert_data(mysql_db, table, data):
    sql = "INSERT INTO %s(%s) VALUES(%s)" % (
        table,
        ','.join(data.keys()),
        ','.join(['%s'] * len(data))
    )
    status = mysql_db.exe_sql(sql, list(data.values()))
    logger.debug(f"入库提示: {status}")
    return status

def select_data(mysql_db, table, col="*", condition=None):
    if not condition:
        sql = f"SELECT {col} FROM {table}"
    else:
        sql = f"SELECT {col} FROM {table} {condition}"
    results = mysql_db.fetch_all(sql)
    return results

def update_data(mysql_db, table, data, condition):
    set_clause = ', '.join([f"{key} = %s" for key in data.keys()])
    sql = f"UPDATE {table} SET {set_clause} WHERE {condition}"
    status = mysql_db.exe_sql(sql, list(data.values()))
    logger.debug(f"更新提示: {status}")
    return status

Redis 操作

import json
import redis

redis_client = redis.Redis('127.0.0.1', password='', decode_responses=True, db=1, health_check_interval=30)

def add_redis(key, value):
    redis_client.sadd(key, json.dumps(value, ensure_ascii=False))
    logger.success(f'redis插入提示 {key}')

def get_redis(key):
    return json.loads(redis_client.spop(key))

路径操作

import os

# 项目路径
BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))

# 浏览器用户文件
TMP_PATH = os.path.join(BASE_DIR, 'Tmp')
os.makedirs(TMP_PATH, exist_ok=True)

# 日志配置
log_path = os.path.join(BASE_DIR, 'Logs')
os.makedirs(log_path, exist_ok=True)

# 文件配置
file_path = os.path.join(BASE_DIR, 'Files')
os.makedirs(file_path, exist_ok=True)

生成md5

import hashlib

def create_md5(input_string):
    # 创建md5对象
    md5 = hashlib.md5()
    # 更新md5对象以包含输入字符串的字节
    md5.update(input_string.encode('utf-8'))
    # 获取16进制的md5值
    return md5.hexdigest()

OSSTool工具

class OSSClient:
    def __init__(self, access_key_id, access_key_secret, endpoint, bucket_name):
        self.auth = oss2.Auth(access_key_id, access_key_secret)
        self.bucket = oss2.Bucket(self.auth, endpoint, bucket_name, connect_timeout=600000)

    def check_bucket_exists(self):
        try:
            self.bucket.get_bucket_info()
            return True
        except oss2.exceptions.NoSuchBucket:
            return False
        except Exception as e:
            logger.error(f"检查 Bucket 是否存在时出错：{e}")
            return False

    def check_resource_exist(self, key):
        try:
            return self.bucket.object_exists(key)
        except Exception as e:
            logger.error(f"检查资源 {key} 是否存在时出错：{e}")
            return False

    def upload_resource(self, key, file_path):
        if not os.path.exists(file_path):
            logger.warning(f"文件 {file_path} 不存在。")
            return
        try:
            self.bucket.put_object_from_file(key, file_path)
            logger.success(f"成功上传资源：{key}")
        except Exception as e:
            logger.error(f"上传资源 {key} 失败：{e}")

    def set_public_read_acl(self, key):
        try:
            self.bucket.put_object_acl(key, oss2.OBJECT_ACL_PUBLIC_READ)
            logger.success(f"成功设置资源 {key} 的公共读取权限")
        except Exception as e:
            logger.error(f"设置资源 {key} 的公共读取权限失败：{e}")

    def get_download_url(self, key, expires=60 * 60 * 24 * 365 * 20):  # 默认有效期为20年
        try:
            self.set_public_read_acl(key)
            download_url = self.bucket.sign_url(
                method='GET',
                key=key,
                expires=expires,
                slash_safe=True
            )
            return download_url
        except Exception as e:
            logger.error(f"获取资源 {key} 的下载链接失败：{e}")
            return None
 
# 阿里云 OSS 访问凭证配置
AccessKeyId = ''
AccessKeySecret = ""
BucketName = ""
Endpoint = ""

oss_client = OSSClient(
    access_key_id=AccessKeyId,
    access_key_secret=AccessKeySecret,
    endpoint=Endpoint,
    bucket_name=BucketName,
)

日志配置

from loguru import logger

# 日志配置
logger.add(
    os.path.join(file_config['Logs'], "{time:YYYY-MM-DD}.log"),
    encoding="utf-8",
    format="{time:YYYY-MM-DD HH:mm:ss.SSS} | {level} | {module}:{function}:{line} - {message}",
    rotation="00:00",
    enqueue=True,
    level="DEBUG",
)

企业机器人推送

def web_hook(content):
    """
    企业微信机器人, 推送程序信息
    :param content: 要推送的文本
    :return: None
    """
    key = ""
    webhook_url = f'https://qyapi.weixin.qq.com/cgi-bin/webhook/send?key={key}'
    data = {
        "msgtype": "markdown",
        "markdown": {
            "content": content
        }
    }
    header = {'Content-Type': 'application/json'}
    try:
        response = requests.post(webhook_url, headers=header, json=data)
        logger.info(f"Webhook推送结果通知 {response.json()}")
    except Exception as e:
        logger.error(f"Webhook 推送结果失败通知 {e}")

去除文件名中无效字符

from uuid import uuid4

def sanitize_filename(filename: str, replacement: str = "_") -> str:
    # 不允许的字符: \ / : * ? " < > |
    illegal_chars = r'[\\/:*?"<>|]'
    filename = re.sub(illegal_chars, replacement, filename)

    # 去掉结尾的空格和点
    filename = filename.rstrip(' .')

    # 禁用 Windows 保留文件名（不区分大小写）
    reserved_names = {
        'CON', 'PRN', 'AUX', 'NUL',
        *(f'COM{i}' for i in range(1, 10)),
        *(f'LPT{i}' for i in range(1, 10)),
    }
    name_part = filename.split('.')[0].upper()
    if name_part in reserved_names:
        filename = f"{filename}_safe"

    # 如果结果为空，则返回默认名
    if not filename:
        filename = uuid4().hex

    return filename

去除url中文件名部分

from urllib.parse import urlparse

def extract_base_url(url: str) -> str:
    parsed = urlparse(url)
    path = os.path.dirname(parsed.path)  # 去除最后的文件名部分
    base_url = f"{parsed.scheme}://{parsed.netloc}{path}"
    return base_url