Building an Async Data Stream Iterator

Hello! Let's continue our exploration of asynchronous programming.

In our previous lesson, we established the formal rules of the asynchronous iteration protocol: the roles of __aiter__, __anext__, and the StopAsyncIteration exception. You now understand the "contract" an object must fulfill to be used in an async for loop.

Today, we'll put that theory into practice. Our goal is to implement a custom asynchronous iterator class that simulates fetching data chunks over a network. This is a common and practical use case, directly analogous to consuming a paginated API, streaming data from a financial exchange, or handling a large file download. We'll focus on how to manage the state of the iteration—such as which chunk comes next—within a class structure.

1. Why a Class-Based Iterator?

Before we dive into coding, let's consider why we'd use a class for this task. As we'll see in a future lesson, asynchronous generators often provide a more concise way to create asynchronous iterators. However, a class-based approach is powerful when the iteration logic is complex or requires explicit state management.

For instance, our "network stream" needs to keep track of:

• The source of the data.
• How many chunks have been sent.
• Whether the connection is still "open".

A class provides a natural way to encapsulate this state and the methods that operate on it.

2. A Blueprint for Class-Based Async Iterators

The Real Python article "Asynchronous Iterators and Iterables in Python" provides an excellent example that is structurally very similar to what we aim to build. It demonstrates creating an async iterator to process a large file in chunks, which is conceptually identical to receiving data chunks over a network.

Please read the following section to see a well-documented implementation of this pattern.

3. Implementing a Simulated Network Stream

Now, let's apply the pattern from the article to our specific goal: simulating a network data stream. Imagine we're connecting to an endpoint that will send us a series of data packets. We'll create a class, NetworkDataStream, that mimics this behavior.

Our iterator will have a predefined list of "data chunks" and will yield them one by one, with a simulated network delay between each chunk.

Here is the full implementation:

import asyncio
import random
import time

class NetworkDataStream:
    """
    An asynchronous iterator that simulates fetching data chunks
    over a network connection.
    """
    def __init__(self, total_chunks: int):
        self._total_chunks = total_chunks
        self._chunks_sent = 0
        # This simulates the data source we are fetching from.
        # In a real scenario, this wouldn't be pre-generated.
        self._data_source = [f"Data chunk {i+1}/{total_chunks}" for i in range(total_chunks)]
        print(f"NetworkDataStream initialized for {total_chunks} chunks.")

    def __aiter__(self):
        """Returns the iterator object itself."""
        print("-> Stream connection established (__aiter__ called).")
        return self

    async def __anext__(self):
        """
        Fetches the next data chunk, simulating network latency.
        Raises StopAsyncIteration when all chunks have been sent.
        """
        if self._chunks_sent >= self._total_chunks:
            print("-> All chunks received, closing stream (raising StopAsyncIteration).")
            raise StopAsyncIteration

        # Simulate variable network latency
        delay = random.uniform(0.2, 0.8)
        await asyncio.sleep(delay)

        # Get the next chunk
        chunk = self._data_source[self._chunks_sent]
        self._chunks_sent += 1
        
        print(f"   (Received after {delay:.2f}s)")
        return chunk

async def main():
    print("Starting data processing pipeline...")
    start_time = time.monotonic()

    # The async for loop seamlessly consumes our custom async iterator
    async for data_chunk in NetworkDataStream(total_chunks=5):
        print(f"Processing: '{data_chunk}'")

    duration = time.monotonic() - start_time
    print(f"\nPipeline finished. Total time: {duration:.2f}s")

if __name__ == "__main__":
    asyncio.run(main())

Deconstructing the Code:

• __init__(self, total_chunks): We initialize the state. _total_chunks is the total number of items we expect, and _chunks_sent tracks our progress. _data_source is our simulated remote data.
• __aiter__(self): As per the protocol, this is a standard def method. It simply returns self, making the object both an async iterable and an async iterator. We've added a print statement to show it's the first thing called by the async for loop, representing the moment a "connection" is established.
• async def __anext__(self): This is the core of our implementation.
  1. Termination Check: The first step is always to check if the iteration is complete. If _chunks_sent has reached the total, we raise StopAsyncIteration to terminate the loop gracefully.
  2. Simulate I/O Wait: await asyncio.sleep(delay) simulates the non-blocking wait for the next packet to arrive from the network. Control is yielded to the event loop here.
  3. Produce Value: Once the "wait" is over, we retrieve the next chunk from our data source.
  4. Update State: We increment _chunks_sent to ensure we fetch the next chunk in the subsequent call.
  5. Return Value: The fetched chunk is returned to the async for loop.

This class-based structure is a versatile pattern. The mCoding video you've seen previously uses the exact same structure to implement a rate limiter, where the state being managed is the time of the last iteration.

Conclusion

In this lesson, we've successfully translated the asynchronous iteration protocol into a practical, class-based implementation.

• Key Takeaway: A custom asynchronous iterator class is an excellent pattern for managing stateful, non-blocking iteration. The structure involves initializing state in __init__, returning self from __aiter__, and implementing the core asynchronous logic in __anext__.
• Practical Application: You can adapt this NetworkDataStream template for many real-world scenarios, such as handling paginated API responses, processing rows from a database cursor asynchronously, or reading from a live data feed.

We've now seen the "verbose" but powerful way to create an async iterator. What if your iteration logic is simpler and doesn't require the full structure of a class?

In our next lesson, we'll explore asynchronous generators. You'll learn how to achieve the exact same result as our NetworkDataStream class, but with significantly less code, by using async def in combination with the yield keyword.