函数返回生成器推导式引发NoneType不可迭代错误的解决方法

问题背景 #

原有生成器函数运行正常：调用方可迭代yield返回的值，无值可yield时，调用方的for...in循环不会执行。代码如下：

class Repository:
   ...
   def gen_by_id(self, related_id: int) -> Iterator[Something]:
      somethings = self._related_id_to_somethings[related_id]
      if somethings:
         for something in somethings:
            yield something

尝试用生成器推导式改写以更贴合Python风格，改写后代码如下：

class Repository:
   ...
   def gen_by_id(self, related_id: int) -> Iterator[Something]:
      somethings = self._related_id_to_somethings[related_id]
      if somethings:
         return (something for something in somethings)

但改写后出现问题：当somethings为None时，函数默认返回NoneType，引发'NoneType' object is not iterable错误，而原函数可正常运行。

错误原因 #

• 原函数是生成器函数（包含yield关键字），即使不进入if分支，Python也会自动返回一个空的生成器对象，该对象可被迭代（只是不会产生任何值）。
• 改写后的函数是普通函数：当不进入if分支时，没有显式return语句，Python默认返回None，而None不可被迭代，导致调用方报错。

正确实现方式 #

方式1：确保始终返回生成器推导式 #

通过处理somethings为None的情况，让它始终指向可迭代对象（空容器），再返回生成器推导式：

class Repository:
   ...
   def gen_by_id(self, related_id: int) -> Iterator[Something]:
      somethings = self._related_id_to_somethings[related_id] or ()
      return (something for something in somethings)

or ()会在somethings为None或空容器时，将其替换为空元组，生成器推导式遍历空元组时会产生空生成器，符合原函数的行为。

方式2：显式返回空生成器 #

在else分支返回空的生成器推导式，确保函数始终返回生成器类型：

class Repository:
   ...
   def gen_by_id(self, related_id: int) -> Iterator[Something]:
      somethings = self._related_id_to_somethings[related_id]
      if somethings:
         return (something for something in somethings)
      return (x for x in ())

方式3：结合yield from的简洁写法 #

如果不执着于生成器推导式，yield from可以更简洁地实现原逻辑，同时保留生成器函数的特性：

class Repository:
   ...
   def gen_by_id(self, related_id: int) -> Iterator[Something]:
      somethings = self._related_id_to_somethings[related_id]
      if somethings:
         yield from somethings

这种写法和原函数行为完全一致，且代码更短。

方案对比与选择 #

• 原写法：逻辑直观，可读性强，适合团队成员对生成器特性不熟悉的场景，不会出现返回None的风险。缺点是代码稍显冗长。
• 生成器推导式改写：更简洁贴合Python风格，通过上述两种方式可以避免返回None的问题，适合追求代码简洁性的场景。
• yield from写法：兼顾简洁性和生成器函数的安全性，是平衡可读性与简洁性的最优方案之一。

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