Is it safe to create a Python class which is simultaneously sync and async iterator?

614 Views Asked by At

Is it safe/bad practice to make a class both iterator and async iterator? Example:

import asyncio


class Iter:
    def __init__(self):
        self.i = 0
        self.elems = list(range(10))
    
    def __iter__(self):
        return self
    
    def __aiter__(self):
        return self
    
    def __next__(self):
        if self.i >= len(self.elems):
            raise StopIteration
        self.i += 1
        return self.elems[self.i - 1]
    
    async def __anext__(self):
        if self.i >= len(self.elems):
            raise StopAsyncIteration
        self.i += 1
        return self.elems[self.i - 1]


async def main():
    print("async usage:")
    async for elem in Iter():
        print(elem)
    print("sync usage:")
    for elem in Iter():
        print(elem)


try:
    asyncio.run(main())
except RuntimeError:
    await main()

I surfed the net and didn't find anybody asking similar question or discussing the problem.

1

There are 1 best solutions below

0
On

The problem there is not the sync iterator: unless used in multi-threaded code it should be ok (but not otherwise).

Your async code, however, keeps the state in a single instance, and if it is ever used in more than a task at once, the states will mix up. (Also, if you have a single task using it, but nests iterations - async or otherwise - in the same instance).

It is easily resolved by returning an auxiliar object that will contain your i counter, for each of __iter__ and __aiter__ - or simply return a generator with the counter in a closure - that is way easier than implementing __next__ and __anext__:

class Iter:
    def __init__(self):
        self.elems = list(range(10))
    
    def __iter__(self):
        return iter(self.elems)
    
    def __aiter__(self):
        async def _iter():
            yield from iter(self.elems)
        return _iter
    

Now if you need some custom logic to Actually go into __next__, other than just sequentially yeld the values of a list, the associated iterator class logic can be of help:

...

class _InnerItter:
    def __init__(self, parent):
        self.parent = parent
        self.i = 0
    
    def __next__(self):
        if self.i >= len(self.parent.elems):
            raise StopIteration
        self.i += 1
        return self.parent.elems[self.i - 1]
    
    async def __anext__(self):
        if self.i >= len(self.parent.elems):
            raise StopAsyncIteration
        self.i += 1
        return self.parent.elems[self.i - 1]


class Iter:
    def __init__(self):
        self.elems = list(range(10))
    
    def __iter__(self):
        return _InnerItter(self)
    
    def __aiter__(self):
        return _InnerItter(self)
...