如何借助functools扩展生成器函数，不修改原函数即可支持返回列表？

问题背景 #

我正在开发一个包含大量返回生成器的模块，希望在不修改底层函数代码的前提下，给这些函数增加一个参数，让用户可以选择返回生成器还是列表。尝试用multimethod实现时，遇到了TypeError: type 'bool' is not an acceptable base type的错误，需要找到可行的实现方式。

错误原因 #

用Literal[True]/Literal[False]作为参数类型注册multimethod分支时，multimethod默认基于参数的类型调度，但Literal是类型注解里的字面量约束，并非实际的类型对象，无法被multimethod识别，因此报错。

可行实现方案 #

方案一：装饰器统一扩展（推荐） #

无需修改原函数，用装饰器包装即可批量给生成器函数添加return_as_list参数，简洁高效：

from typing import Union, Generator, List
from functools import wraps

def with_list_option(func):
    """给生成器函数添加返回列表的选项"""
    @wraps(func)
    def wrapper(*args, return_as_list: bool = False, **kwargs):
        result = func(*args, **kwargs)
        return list(result) if return_as_list else result
    return wrapper

# 原生成器函数（无需修改）
def base_func(a: int, b: int) -> Generator[int, None, None]:
    assert b > a
    for i in range(a, b):
        yield i

# 扩展后的函数
extended_func = with_list_option(base_func)

# 使用示例
print(extended_func(1, 3))  # 输出生成器对象: <generator object base_func at 0x...>
print(extended_func(1, 3, return_as_list=True))  # 输出列表: [1, 2]

方案二：调整multimethod使用方式 #

如果坚持用multimethod，可通过调整注册逻辑实现按参数值调度：

from typing import Union, Generator, List, Literal
from functools import wraps
from multimethod import multimethod

def as_list(func):
    """将生成器转换为列表的装饰器"""
    @wraps(func)
    def wrapper(*args, **kwargs):
        return list(func(*args, **kwargs))
    return wrapper

# 原生成器函数（无需修改）
def base_func(a: int, b: int) -> Generator[int, None, None]:
    assert b > a
    for i in range(a, b):
        yield i

# 定义多方法调度入口
@multimethod
def test_func(a: int, b: int, return_as_list: bool):
    pass

# 注册返回生成器的分支
@test_func.register
def _(a: int, b: int, return_as_list: Literal[False]):
    return base_func(a, b)

# 注册返回列表的分支，用as_list装饰转换结果
@test_func.register
@as_list
def _(a: int, b: int, return_as_list: Literal[True]):
    return base_func(a, b)

# 使用示例
print(test_func(1, 3, return_as_list=False))  # 输出生成器对象
print(test_func(1, 3, return_as_list=True))  # 输出列表: [1, 2]

注意：需确保multimethod版本支持识别Literal，若版本不兼容，可直接通过参数值判断实现分支匹配。

总结 #

• 装饰器方案更简洁，适合批量扩展现有生成器函数，完全无需修改原代码。
• 多方法调度方案适合需要复杂分支逻辑的场景，但实现相对繁琐。

内容的提问来源于stack exchange，提问作者Le Chase