Add async_iterator util (#153194)

homeassistant/components/backup/util.py

@@ -4,7 +4,6 @@ from __future__ import annotations
 
 import asyncio
 from collections.abc import AsyncIterator, Callable, Coroutine
-from concurrent.futures import CancelledError, Future
 import copy
 from dataclasses import dataclass, replace
 from io import BytesIO
@@ -14,7 +13,7 @@ from pathlib import Path, PurePath
 from queue import SimpleQueue
 import tarfile
 import threading
-from typing import IO, Any, Self, cast
+from typing import IO, Any, cast
 
 import aiohttp
 from securetar import SecureTarError, SecureTarFile, SecureTarReadError
@@ -23,6 +22,11 @@ from homeassistant.backup_restore import password_to_key
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import HomeAssistantError
 from homeassistant.util import dt as dt_util
+from homeassistant.util.async_iterator import (
+    Abort,
+    AsyncIteratorReader,
+    AsyncIteratorWriter,
+)
 from homeassistant.util.json import JsonObjectType, json_loads_object
 
 from .const import BUF_SIZE, LOGGER
@@ -59,12 +63,6 @@ class BackupEmpty(DecryptError):
     _message = "No tar files found in the backup."
 
 
-class AbortCipher(HomeAssistantError):
-    """Abort the cipher operation."""
-
-    _message = "Abort cipher operation."
-
-
 def make_backup_dir(path: Path) -> None:
     """Create a backup directory if it does not exist."""
     path.mkdir(exist_ok=True)
@@ -166,106 +164,6 @@ def validate_password(path: Path, password: str | None) -> bool:
     return False
 
 
-class AsyncIteratorReader:
-    """Wrap an AsyncIterator."""
-
-    def __init__(self, hass: HomeAssistant, stream: AsyncIterator[bytes]) -> None:
-        """Initialize the wrapper."""
-        self._aborted = False
-        self._hass = hass
-        self._stream = stream
-        self._buffer: bytes | None = None
-        self._next_future: Future[bytes | None] | None = None
-        self._pos: int = 0
-
-    async def _next(self) -> bytes | None:
-        """Get the next chunk from the iterator."""
-        return await anext(self._stream, None)
-
-    def abort(self) -> None:
-        """Abort the reader."""
-        self._aborted = True
-        if self._next_future is not None:
-            self._next_future.cancel()
-
-    def read(self, n: int = -1, /) -> bytes:
-        """Read data from the iterator."""
-        result = bytearray()
-        while n < 0 or len(result) < n:
-            if not self._buffer:
-                self._next_future = asyncio.run_coroutine_threadsafe(
-                    self._next(), self._hass.loop
-                )
-                if self._aborted:
-                    self._next_future.cancel()
-                    raise AbortCipher
-                try:
-                    self._buffer = self._next_future.result()
-                except CancelledError as err:
-                    raise AbortCipher from err
-                self._pos = 0
-            if not self._buffer:
-                # The stream is exhausted
-                break
-            chunk = self._buffer[self._pos : self._pos + n]
-            result.extend(chunk)
-            n -= len(chunk)
-            self._pos += len(chunk)
-            if self._pos == len(self._buffer):
-                self._buffer = None
-        return bytes(result)
-
-    def close(self) -> None:
-        """Close the iterator."""
-
-
-class AsyncIteratorWriter:
-    """Wrap an AsyncIterator."""
-
-    def __init__(self, hass: HomeAssistant) -> None:
-        """Initialize the wrapper."""
-        self._aborted = False
-        self._hass = hass
-        self._pos: int = 0
-        self._queue: asyncio.Queue[bytes | None] = asyncio.Queue(maxsize=1)
-        self._write_future: Future[bytes | None] | None = None
-
-    def __aiter__(self) -> Self:
-        """Return the iterator."""
-        return self
-
-    async def __anext__(self) -> bytes:
-        """Get the next chunk from the iterator."""
-        if data := await self._queue.get():
-            return data
-        raise StopAsyncIteration
-
-    def abort(self) -> None:
-        """Abort the writer."""
-        self._aborted = True
-        if self._write_future is not None:
-            self._write_future.cancel()
-
-    def tell(self) -> int:
-        """Return the current position in the iterator."""
-        return self._pos
-
-    def write(self, s: bytes, /) -> int:
-        """Write data to the iterator."""
-        self._write_future = asyncio.run_coroutine_threadsafe(
-            self._queue.put(s), self._hass.loop
-        )
-        if self._aborted:
-            self._write_future.cancel()
-            raise AbortCipher
-        try:
-            self._write_future.result()
-        except CancelledError as err:
-            raise AbortCipher from err
-        self._pos += len(s)
-        return len(s)
-
-
 def validate_password_stream(
     input_stream: IO[bytes],
     password: str | None,
@@ -342,7 +240,7 @@ def decrypt_backup(
         finally:
             # Write an empty chunk to signal the end of the stream
             output_stream.write(b"")
-    except AbortCipher:
+    except Abort:
         LOGGER.debug("Cipher operation aborted")
     finally:
         on_done(error)
@@ -430,7 +328,7 @@ def encrypt_backup(
         finally:
             # Write an empty chunk to signal the end of the stream
             output_stream.write(b"")
-    except AbortCipher:
+    except Abort:
         LOGGER.debug("Cipher operation aborted")
     finally:
         on_done(error)
@@ -557,8 +455,8 @@ class _CipherBackupStreamer:
             self._hass.loop.call_soon_threadsafe(worker_status.done.set)
 
         stream = await self._open_stream()
-        reader = AsyncIteratorReader(self._hass, stream)
-        writer = AsyncIteratorWriter(self._hass)
+        reader = AsyncIteratorReader(self._hass.loop, stream)
+        writer = AsyncIteratorWriter(self._hass.loop)
         worker = threading.Thread(
             target=self._cipher_func,
             args=[