async filter function with generic typing

Question

I'm trying to implement the python built in filter but async, seem like an easy task right?

async def simulated_data() -> AsyncIterator[int|None]:
    for i in [1, None,3,5]:
        yield i

async def afilter[T](predicate, iterable):
    async for item in iterable:
        if predicate is not none and predicate(item):
            yield item

b = afilter(None, simulated_data())
# or just this!
b = (it async for it in iter if iter is not None) 

Even a comprehension does the trick :D

But what about typing? The type of b still shows "AsyncGenerator[int | None, None]" but it can´t be None.

I tried with TypeGuard, but no luck, then I went to the original filter function, because this problem is solved already there.

class filter(Generic[_T]):
    @overload
    def __new__(cls, function: None, iterable: Iterable[_T | None], /) -> Self: ...
    @overload
    def __new__(cls, function: Callable[[_S], TypeGuard[_T]], iterable: Iterable[_S], /) -> Self: ...
    @overload
    def __new__(cls, function: Callable[[_S], TypeIs[_T]], iterable: Iterable[_S], /) -> Self: ...
    @overload
    def __new__(cls, function: Callable[[_T], Any], iterable: Iterable[_T], /) -> Self: ...
    def __iter__(self) -> Self: ...
    def __next__(self) -> _T: ...

Well it seems filter is not even a function is a generic class, at this point the task doesn't look so easy, anyone has the solution (with generic types) by any chance?

Answer 1

A minimal example of an async filter can be defined with two overloads:

from typing import AsyncIterator, AsyncIterable, Callable, overload, AsyncGenerator

async def simulated_data() -> AsyncIterator[int|None]:
    for i in [1, None,3,5]:
        yield i

@overload
async def afilter[T](predicate: None, iterable: AsyncIterable[T | None]) -> AsyncGenerator[T]: ...

@overload
async def afilter[T](predicate: Callable[[T], bool], iterable: AsyncIterable[T]) -> AsyncGenerator[T]: ...

async def afilter[T](predicate: Callable[[T], bool] | None, iterable: AsyncIterable[T])  -> AsyncGenerator[T]:
    async for item in iterable:
        if predicate is None:
            if item:
                yield item
        elif predicate(item):
            yield item

# No predicate
only_int = afilter(None, simulated_data())
reveal_type(only_int)  # AsyncGenerator[int, None]

# Some predicate
both = afilter(lambda data: False, simulated_data())
reveal_type(both)  # AsyncGenerator[int | None, None] 

# Comprehension
aiter = simulated_data()
comprehension = (it async for it in aiter if it is not None)
reveal_type(comprehension)  # AsyncGenerator[int, None] 

You will realize that when using a predicate there is no further narrowing, it will be just type T. If you want to narrow down types further you need more overloads for predicate similar to the filter function:

@overload
async def afilter[T, S](predicate: Callable[[S], TypeIs[T], iterable: AsyncIterable[S]) -> AsyncGenerator[T]: ...

@overload
async def afilter[T, S](predicate: Callable[[S], TypeGuard[T], iterable: AsyncIterable[S]) -> AsyncGenerator[T]: ...

With these two you can, for example, define a custom predicate:

def remove_none[T](value: T | None) -> TypeGuard[T]:
   return value is not None

without_none = afilter(remove_none, simulated_data())
# AsyncGenerator[int, None]

Answer 2

A few general comments concerning type hints:

1. Many types in the typing module have been deprecated since Python 3.9, for example typing.Callable (see typing.Callable), in favor of similarly named types in module collections.abc.

2. An AsyncGenerator type takes two arguments: the type of values it yields and the type of values you can send to it. When the latter is not used, None is specfied. For example:

async def simulated_data() -> AsyncGenerator[int | None, None]:

3. A Callable type takes two arguments: The first argument is a list of the argument types it accepts and the second argument is the type of return value. For example, Callable[[T], bool] would be a callable with one argument of generic type T that returns a bool type.

mypy reports: Success: no issues found in 1 source file against the following:

import asyncio
from typing import TypeVar
from collections.abc import Callable, AsyncGenerator, AsyncIterable

T = TypeVar('T')

async def afilter(
    predicate: Callable[[T], bool] | None,
    iterable: AsyncIterable[T]
) -> AsyncGenerator[T, None]:
    if predicate is not None and not callable(predicate):
        raise ValueError('predicate must be either None or a callable')

    async for item in iterable:
        if predicate and predicate(item):
            yield item

async def simulated_data() -> AsyncGenerator[int | None, None]:
    for i in [1, None, 3, 5]:
        yield i

async def main():
    async for item in afilter(lambda x: x, simulated_data()):
        print(item)

    # Another way using a generator expression:
    b = (item async for item in afilter(lambda x: x, simulated_data()))
    # Now iterate the generator expression:
    async for item in b:
        print(item)

    # Another way using a comprehension:
    b = [item async for item in afilter(lambda x: x, simulated_data())]
    # Now iterate the list:
    for item in b:
        print(item)

asyncio.run(main())

Prints:

1
3
5
1
3
5
1
3
5

Answer 3

from typing import AsyncIterator, Callable, Optional, TypeVar

# Define a type variable for generic typing
_T = TypeVar("_T")

# Asynchronous filter function
async def afilter(
    predicate: Optional[Callable[[_T], bool]], 
    iterable: AsyncIterator[_T]
) -> AsyncIterator[_T]:
async for item in iterable:
    if predicate is None:
        if item is not None:
            yield item
    elif predicate(item):
        yield item

 #Exampl of Simulated data generator
async def simulated_data() -> AsyncIterator[int | None]:
    for i in [1, None, 3, 5]:
        yield i

# Example usage
async def main():
    # Filter out None values
    filtered = afilter(None, simulated_data())

    # Print the filtered results
    async for item in filtered:
        print(item)
#Example:
 if __name__ == "__main__":
    import asyncio
    asyncio.run(main())