Python Telegram 机器人回调函数代码复用技巧

2025-02-26 05:26:38

不同回调函数间代码复用：告别“复制粘贴”

遇到了一个很常见的问题：有个机器人，一个命令，一个回调。需要在回调函数里复用命令处理函数中的一段代码，但又不想简单地复制粘贴。这篇文章咱们来好好聊聊怎么解决这个问题。

问题：重复代码的烦恼

直接看代码更清楚：

@dp.message(CommandStart())
async def bot_start_handler(message: Message) -> None:
    # 一些代码...
    await message.answer(message.from_user.id)

@dp.callback_query(F.data == "continue")
async def bot_get_sign(callback: types.CallbackQuery):
    # 一些代码...
    functionA()

@dp.callback_query(F.data.startswith('q_'))
async def sign_step(callback: types.CallbackQuery):
    await callback.message.delete()

    await bot_start_handler() # 这里需要调用 bot_start_handler(), 但怎么调？

sign_step 函数里想执行 bot_start_handler 函数里的逻辑，但直接调用有问题，因为 bot_start_handler 需要一个 Message 对象，而 sign_step 接收的是 CallbackQuery 对象。

为啥不能直接“复制粘贴”？

如果只是简单的代码复制，短期内可能没问题，但长期来看，会带来不少麻烦：

维护困难： 同样的逻辑散落在多个地方，一旦需要修改，得改好几处，容易出错，也费时间。
代码冗余： 代码库里充斥着重复的代码块，不够简洁，看着也难受。

几种方法：让代码更优雅

好在，我们有多种方法可以解决这个问题，告别代码复制。

1. 提取核心逻辑：打造独立函数

这是最常用的方法，也是最推荐的方法。基本思路就是：把 bot_start_handler 里需要复用的代码提取出来，放到一个单独的函数里。

async def send_user_id(user_id: int) -> None:
    # 把这部分代码提取出来
     # await message.answer(message.from_user.id) 改成下面的：
    await bot.send_message(user_id, str(user_id))

@dp.message(CommandStart())
async def bot_start_handler(message: Message) -> None:
    # 一些代码...
    await send_user_id(message.from_user.id)

@dp.callback_query(F.data.startswith('q_'))
async def sign_step(callback: types.CallbackQuery):
    await callback.message.delete()
    await send_user_id(callback.from_user.id)

原理：
我们将重复使用的部分 await message.answer(message.from_user.id)提取到 send_user_id 函数。现在，bot_start_handler 和 sign_step 都可以通过调用 send_user_id 来复用这段逻辑了。而且通过用户id调用.
解释:
我们从message或者callback中取出的id参数传入函数

这样，bot_start_handler 和 sign_step 都能通过调用 send_user_id 函数来复用发送用户 ID 的逻辑了。

2. 修改函数签名：接受更通用的参数

如果核心逻辑依赖于 Message 对象中的特定信息（比如用户 ID、用户名等），但这些信息在 CallbackQuery 对象里也有，那我们可以修改函数的签名，让它接受更通用的参数。

@dp.message(CommandStart())
async def bot_start_handler(user_id: int) -> None:  # 修改参数
    # 一些代码...
    await bot.send_message(user_id, str(user_id)) #bot 是一个Bot对象

@dp.callback_query(F.data.startswith('q_'))
async def sign_step(callback: types.CallbackQuery):
    await callback.message.delete()
    await bot_start_handler(callback.from_user.id) # 传递 user_id

原理：

我们将bot_start_handler的参数从Message对象改为user_id(整数类型). 现在, bot_start_handler 不再依赖特定的消息类型, 任何提供user_id的地方都可以调用它.

解释：

在sign_step里，我们从callback.from_user.id取出user_id，然后传给bot_start_handler。

3. 传递必要的数据：构建“桥梁”

如果 bot_start_handler 里需要用到 Message 对象中的多个属性，而且这些属性在 CallbackQuery 对象里都有对应的值，那么我们可以创建一个新的 Message 对象，或者一个包含必要数据的自定义对象，然后传给 bot_start_handler。

from aiogram.types import User, Chat

@dp.message(CommandStart())
async def bot_start_handler(message: Message) -> None:
    # 一些代码...
    await message.answer(str(message.from_user.id))

@dp.callback_query(F.data.startswith('q_'))
async def sign_step(callback: types.CallbackQuery):
    await callback.message.delete()

    # 构建一个“假的” Message 对象
    fake_message = callback.message
    fake_message.from_user = callback.from_user
    await bot_start_handler(fake_message)

原理：

直接将callback.message传入bot_start_handler。在一些情况下, callback.message 已包含了bot所需的大部分message信息, 除了 from_user外.
我们将 callback.from_user 赋给 fake_message.from_user。现在 fake_message 拥有足够信息给 bot_start_handler

解释：
因为是构建message对象, 因此需要注意对象属性类型的匹配, 这是个相对取巧的方法。

4. 进阶技巧：利用上下文管理器

如果你的代码逻辑涉及到一些资源的初始化和清理（比如数据库连接、文件操作等），可以考虑使用上下文管理器（with 语句）来管理这些资源，然后在不同的函数中共享这些资源。

import contextlib

@contextlib.asynccontextmanager
async def db_session():
    # 初始化数据库连接
    session = ...  # 假设这是你的数据库连接对象
    try:
        yield session
    finally:
        # 清理数据库连接
        await session.close()

@dp.message(CommandStart())
async def bot_start_handler(message: Message) -> None:
    async with db_session() as session:
        # 使用 session 进行数据库操作...
        await message.answer(message.from_user.id)
@dp.callback_query(F.data.startswith('q_'))
async def sign_step(callback: types.CallbackQuery):
   async with db_session() as session:
        # 使用同一个 session 进行数据库操作...
         await callback.message.answer(callback.from_user.id)

原理:

这里使用 contextlib.asynccontextmanager 创建了一个异步上下文管理器 db_session。这个管理器负责数据库连接的创建和关闭。bot_start_handler 和 sign_step 都使用 async with db_session() as session: 来获取数据库连接。由于它们处于同一个事件循环中，它们将共享同一个数据库连接。

解释:

使用上下文管理器的好处是确保资源在使用完毕后能够被正确释放, 避免资源泄漏. 同时也能更好地组织和复用涉及资源管理的代码.