python IO with progress

io.py

import io
import sys
import time
from typing import Any

import humanize


class IO(io.BytesIO):
    ANSI_CLEAR_LINE = "\033[K"

    def __init__(self, *args: Any, **kwargs: Any) -> None:
        self.n = 0
        self.i = 0
        self.started_at = time.monotonic()
        super().__init__(*args, **kwargs)

    def read(self, size: int | None = -1) -> bytes:
        data = super().read(size)
        self.n += len(data)
        print(
            f"{self.n}: {len(data)=} {self.i}\r",
            end="",
            flush=True,
            file=sys.stderr,
        )
        if len(data) == 0:
            elapsed = time.monotonic() - self.started_at
            throughput = humanize.naturalsize(self.n / elapsed) + "/s"
            print(
                f"{self.n} bytes read in {elapsed:.2f} seconds "
                f"({throughput=}){IO.ANSI_CLEAR_LINE}",
                IO.ANSI_CLEAR_LINE,
                file=sys.stderr,
            )
        self.i += 1
        return data