Returning generator comprehension in function causing NoneType not iterable error

97 Views Asked by At

I have a class with a generator function that works as intended. A caller iterates through the yield'ed values, and never executes a for ... in loop when nothing is yield'ed.

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

I just learned about generator comprehensions, and would like to use that if possible in the generator function above. It's really just in an effort to be more Pythonic, but there are a few other situations where they will shorten code a lot more. I tried this:

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)

Unfortunately, if somethings is None, my original generator function worked fine, and the caller just wouldn't have anything to iterate over -- but, my generator comprehension function returns NoneType which causes a 'NoneType' object is not iterable TypeError.

I tried adding raise StopIteration to the end of the function, wondering if that might be needed, but that didn't help. It just changes the error to that the StopIteration error was raised but nothing 'caught' it.

I avoid the problem by making the generator function:

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 ()

But, I'm not sure if that's correct (even though it works) since that's returning an empty tuple. And, it seems so bizarre to me that in the working scenario, the function trails off without a return or yield so I have no idea why when code that isn't executing is changed from a yield to a return why it would have to do anything different.

I also avoid the problem by making the generator function:

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

And maybe that's what I have to do, but it still feels wrong and like there must be a better way.

How can I make the generator function work with a generator comprehension? Or, should I be leaving it in its original form and am I trying to misuse generator comprehensions somehow? Or, is there a better way to be doing this?

I see the type my generator comprehension is creating is Generator[Something, Any, None] instead of Iterator[Package], but I think inheritance must be involved there because pyCharm and mypy aren't complaining about a mismatched return type.

2

There are 2 best solutions below

0
Barmar On

Rather than returning the generator, use yield from to pass along the iteration.

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

Another alternative, always return (the result of) the generator expression:

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

Note the added or ().

Barmar's with yield from makes every value go through an extra generator layer, which is less efficient.

But I really think your original without a generator expression is the most pythonic.