用于列表的 Python 函数 `yield`，返回单个元素

是否有可能有一个函数，当传递一个列表时会产生一个生成器，但在给定单个值时会返回？

以这个为例:

def duty2015(x):

  if type(x) in [list, np.ndarray]:
    for xi in x:
      yield new_duty(xi)

  else:

    sd = 0
    if x <= 120000:
      return sd
    elif x <= 250000:
      return (x-125000) * 0.02
    elif x <= 925000:
      return 2500 + (x-250000)*0.05
    elif  x <= 1500000:
      return 36250 + (x-925000)*0.1
    else:
      return 93750 + (x-1500000)*0.12

显然这不起作用，我收到了 SyntaxError: 'return' with argument inside generator 错误。

我意识到我可以做这样的事情:

def duty2015(x):

  if type(x) in [list, np.ndarray]:
    for xi in x:
      for result in duty2015(xi):
        yield result

  else:

    sd = 0
    if x <= 125000:
      yield sd
    elif x <= 250000:
      yield (x-125000) * 0.02
    elif x <= 925000:
      yield 2500 + (x-250000)*0.05
    elif  x <= 1500000:
      yield 36250 + (x-925000)*0.1
    else:
      yield 93750 + (x-1500000)*0.12

但是当我在单个项目上调用它时，它给了我一个生成器，我宁愿只在较大的项目上调用它时才得到它。

显然我可以把它作为一个列表来做，但这又不是最优的。

对于下面答案中的评论，这样的事情会更好吗:

def duty_new(x, generator=False):

  if type(x) in [list, np.ndarray]:
    if generator:
      return (duty_new(xi) for xi in x)
    else:
      return [duty_new(xi) for xi in x]

  else:

    sd = 0
    if x <= 125000:
      return sd
    elif x <= 250000:
      return (x-125000) * 0.02
    elif x <= 925000:
      return 2500 + (x-250000)*0.05
    elif  x <= 1500000:
      return 36250 + (x-925000)*0.1
    else:
      return 93750 + (x-1500000)*0.12

因此，在正常使用下，它将具有可预测的行为或返回与传递给它的相同类型的参数(至少对于合理的参数，并且可能是这样的，以便它不只是遍历 numpy 数组)，但是如果需要生成器，可以明确要求吗？

最佳答案

返回一个生成器表达式，像这样

def duty2015(x):

  if isinstance(x, list) or isinstance(x, np.ndarray):
      return (result for xi in x for result in duty2015(xi))
  else:
     ...
     ...

现在，每当您使用单个元素调用 duty2015 时，您将获得单独的值，否则您将获得一个生成器表达式，该表达式必须使用 next 协议(protocol)进行迭代。

就我个人而言，我觉得你的第二个版本很好并且是一致的，因为它使 duty2015 成为一个生成器函数，并且正如 Martijn Pieters 在评论中提到的那样，它没有让调用者猜测它得到了什么，最好坚持下去。

注意: 你在第一个版本和第二个版本中的代码是不同的。我选择了第二个版本中的代码来展示这个想法。

关于用于列表的 Python 函数 `yield`，返回单个元素，我们在Stack Overflow上找到一个类似的问题：https://stackoverflow.com/questions/28236695/