如何在保持代码结构的前提下自动将普通函数转为Python生成器？

生成器优雅且内存高效，因此我希望将返回序列的普通函数转换为生成器。但未找到简便的转换方法——这里的“简便”指最好能通过工具自动完成，比如简单的字符串或正则替换（如将return替换为yield）。若转换无法始终简便，重构可能引入Bug。当然单元测试能捕获部分Bug，但如果原函数已充分测试，我怀疑是否需要为生成器额外编写测试。

举个例子，普通函数normalize接收src参数：若src是list则返回该列表，否则返回包含src作为单个元素的列表。此处为简化，假设调用者传入的src要么是list要么是简单标量。

测试代码 #

f = normalize # 替换为其他转换后的函数
assert list(f(2)) == [2]
assert list(f([1, 2])) == [1, 2]

普通函数实现 #

def normalize(src):
    if isinstance(src, list):
        return src
    return [src]

我添加else块转换为生成器后，测试通过：

带else块的生成器（测试通过） #

def normalize_generator(src):
    if isinstance(src, list):
        for x in src:
            yield x
    else:
        yield src

但我希望保留normalize的代码结构，仅替换关键词，如下所示，但测试会失败：

保留原结构但测试失败的生成器 #

def normalize_generator_no_else(src):
    if isinstance(src, list):
        for x in src:
            yield x
    yield src

测试结果失败：左侧为[1, 2, [1, 2]]，右侧为[1, 2]

我大致理解该行为，yield不会终止函数执行，后续代码会继续运行。我已查找过类似问题，但未找到合适的解决方案。

内容的提问来源于stack exchange，提问作者Lei Yang