Make ChunkAsyncStreamIterator an aiohttp helper (#134843)

make ChunkAsyncStreamIterator a generic aiohttp helper
This commit is contained in:
Michael 2025-01-06 04:37:07 +01:00 committed by GitHub
parent bc22e34fc3
commit acd95975e4
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 29 additions and 28 deletions

View File

@ -6,9 +6,9 @@ import base64
from collections.abc import AsyncIterator, Callable, Coroutine, Mapping from collections.abc import AsyncIterator, Callable, Coroutine, Mapping
import hashlib import hashlib
import logging import logging
from typing import Any, Self from typing import Any
from aiohttp import ClientError, ClientTimeout, StreamReader from aiohttp import ClientError, ClientTimeout
from hass_nabucasa import Cloud, CloudError from hass_nabucasa import Cloud, CloudError
from hass_nabucasa.cloud_api import ( from hass_nabucasa.cloud_api import (
async_files_delete_file, async_files_delete_file,
@ -19,6 +19,7 @@ from hass_nabucasa.cloud_api import (
from homeassistant.components.backup import AgentBackup, BackupAgent, BackupAgentError from homeassistant.components.backup import AgentBackup, BackupAgent, BackupAgentError
from homeassistant.core import HomeAssistant, callback from homeassistant.core import HomeAssistant, callback
from homeassistant.helpers.aiohttp_client import ChunkAsyncStreamIterator
from homeassistant.helpers.dispatcher import async_dispatcher_connect from homeassistant.helpers.dispatcher import async_dispatcher_connect
from .client import CloudClient from .client import CloudClient
@ -73,31 +74,6 @@ def async_register_backup_agents_listener(
return unsub return unsub
class ChunkAsyncStreamIterator:
"""Async iterator for chunked streams.
Based on aiohttp.streams.ChunkTupleAsyncStreamIterator, but yields
bytes instead of tuple[bytes, bool].
"""
__slots__ = ("_stream",)
def __init__(self, stream: StreamReader) -> None:
"""Initialize."""
self._stream = stream
def __aiter__(self) -> Self:
"""Iterate."""
return self
async def __anext__(self) -> bytes:
"""Yield next chunk."""
rv = await self._stream.readchunk()
if rv == (b"", False):
raise StopAsyncIteration
return rv[0]
class CloudBackupAgent(BackupAgent): class CloudBackupAgent(BackupAgent):
"""Cloud backup agent.""" """Cloud backup agent."""

View File

@ -9,7 +9,7 @@ import socket
from ssl import SSLContext from ssl import SSLContext
import sys import sys
from types import MappingProxyType from types import MappingProxyType
from typing import TYPE_CHECKING, Any from typing import TYPE_CHECKING, Any, Self
import aiohttp import aiohttp
from aiohttp import web from aiohttp import web
@ -82,6 +82,31 @@ class HassClientResponse(aiohttp.ClientResponse):
return await super().json(*args, loads=loads, **kwargs) return await super().json(*args, loads=loads, **kwargs)
class ChunkAsyncStreamIterator:
"""Async iterator for chunked streams.
Based on aiohttp.streams.ChunkTupleAsyncStreamIterator, but yields
bytes instead of tuple[bytes, bool].
"""
__slots__ = ("_stream",)
def __init__(self, stream: aiohttp.StreamReader) -> None:
"""Initialize."""
self._stream = stream
def __aiter__(self) -> Self:
"""Iterate."""
return self
async def __anext__(self) -> bytes:
"""Yield next chunk."""
rv = await self._stream.readchunk()
if rv == (b"", False):
raise StopAsyncIteration
return rv[0]
@callback @callback
@bind_hass @bind_hass
def async_get_clientsession( def async_get_clientsession(