Python function `yield` for lists, return for individual elements

Question

Is it possible to have a function which will produce a generator when passed a list, but will then return when given individual values?

Take this for example:

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

obviously this doesn't work, and i get the SyntaxError: 'return' with argument inside generator error.

I realise i could do something like this instead:

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

but then when i call it on individual items it gives me a generator, which i'd rather only get when it's called on larger items.

Obviously i could do it as a list, but again this is not optimal.

For the comments in answers below, would something like this be better then:

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

So that under normal use, it would have a predictable behaviour or returning the same type of argument as is passed to it (for sensible arguments at least, and probably something so that it doesn't just iterate through numpy arrays), but if a generator were needed it could be explicitly asked for?

Answer 1

Return a generator expression, like this

def duty2015(x):

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

Now, whenever you call duty2015 with a single element you will get individual value, otherwise you will get a generator expression, which has to be iterated with next protocol.

Personally, I feel that your second version is good and it is consistent, because it makes duty2015 a generator function and as Martijn Pieters mentioned in the comments it doesn't have its caller guess what it got back, better stick to that.

Note: Your code in the first version and second version are different. I choose the code in the second version to show the idea.