Coverage for /pythoncovmergedfiles/medio/medio/src/aiohttp/aiohttp/streams.py: 35%

1import asyncio

2import collections

3import sys

4import warnings

5from collections.abc import Awaitable, Callable, Coroutine

6from typing import Final, Generic, TypeVar

8from .base_protocol import BaseProtocol

9from .helpers import (

10 _EXC_SENTINEL,

11 DEFAULT_CHUNK_SIZE,

12 BaseTimerContext,

13 TimerNoop,

14 set_exception,

15 set_result,

16)

17from .http_exceptions import LineTooLong

18from .log import internal_logger

20__all__ = (

21 "EMPTY_PAYLOAD",

22 "EofStream",

23 "StreamReader",

24 "DataQueue",

25)

27_T = TypeVar("_T")

30class EofStream(Exception):

31 """eof stream indication."""

34class AsyncStreamIterator(Generic[_T]):

36 __slots__ = ("read_func",)

38 def __init__(self, read_func: Callable[[], Awaitable[_T]]) -> None:

39 self.read_func = read_func

41 def __aiter__(self) -> "AsyncStreamIterator[_T]":

42 return self

44 async def __anext__(self) -> _T:

45 try:

46 rv = await self.read_func()

47 except EofStream:

48 raise StopAsyncIteration

49 if rv == b"":

50 raise StopAsyncIteration

51 return rv

54class ChunkTupleAsyncStreamIterator:

56 __slots__ = ("_stream",)

58 def __init__(self, stream: "StreamReader") -> None:

59 self._stream = stream

61 def __aiter__(self) -> "ChunkTupleAsyncStreamIterator":

62 return self

64 async def __anext__(self) -> tuple[bytes, bool]:

65 rv = await self._stream.readchunk()

66 if rv == (b"", False):

67 raise StopAsyncIteration

68 return rv

71class StreamReader:

72 """An enhancement of asyncio.StreamReader.

74 Supports asynchronous iteration by line, chunk or as available::

76 async for line in reader:

77 ...

78 async for chunk in reader.iter_chunked(1024):

79 ...

80 async for slice in reader.iter_any():

81 ...

83 """

85 __slots__ = (

86 "_protocol",

87 "_low_water",

88 "_high_water",

89 "_low_water_chunks",

90 "_high_water_chunks",

91 "_loop",

92 "_size",

93 "_cursor",

94 "_http_chunk_splits",

95 "_buffer",

96 "_buffer_offset",

97 "_eof",

98 "_waiter",

99 "_eof_waiter",

100 "_exception",

101 "_timer",

102 "_eof_callbacks",

103 "_eof_counter",

104 "_on_chunk_received",

105 "total_bytes",

106 "total_compressed_bytes",

107 )

108

109 def __init__(

110 self,

111 protocol: BaseProtocol,

112 limit: int,

113 *,

114 timer: BaseTimerContext | None = None,

115 loop: asyncio.AbstractEventLoop,

116 ) -> None:

117 self._protocol = protocol

118 self._low_water = limit

119 self._high_water = limit * 2

120 # Use max(4, ...) because there's always at least 1 chunk split remaining

121 # (the current position), so we need low_water >= 2 to allow resume.

122 # limit // 16 gets us a reasonable value of 16k with default 256KiB limit.

123 self._high_water_chunks = max(4, limit // 16)

124 self._low_water_chunks = self._high_water_chunks // 2

125 self._loop = loop

126 self._size = 0

127 self._cursor = 0

128 self._http_chunk_splits: collections.deque[int] | None = None

129 self._buffer: collections.deque[bytes] = collections.deque()

130 self._buffer_offset = 0

131 self._eof = False

132 self._waiter: asyncio.Future[None] | None = None

133 self._eof_waiter: asyncio.Future[None] | None = None

134 self._exception: type[BaseException] | BaseException | None = None

135 self._timer = TimerNoop() if timer is None else timer

136 self._eof_callbacks: list[Callable[[], None]] = []

137 self._eof_counter = 0

138 self._on_chunk_received: (

139 Callable[[bytes], Coroutine[None, None, None]] | None

140 ) = None

141 self.total_bytes = 0

142 self.total_compressed_bytes: int | None = None

143

144 def __repr__(self) -> str:

145 info = [self.__class__.__name__]

146 if self._size:

147 info.append("%d bytes" % self._size)

148 if self._eof:

149 info.append("eof")

150 if self._low_water != DEFAULT_CHUNK_SIZE:

151 info.append("low=%d high=%d" % (self._low_water, self._high_water))

152 if self._waiter:

153 info.append("w=%r" % self._waiter)

154 if self._exception:

155 info.append("e=%r" % self._exception)

156 return "<%s>" % " ".join(info)

157

158 def __aiter__(self) -> AsyncStreamIterator[bytes]:

159 return AsyncStreamIterator(self.readline)

160

161 def iter_chunked(self, n: int) -> AsyncStreamIterator[bytes]:

162 """Returns an asynchronous iterator that yields chunks of size n."""

163 self.set_read_chunk_size(n)

164 return AsyncStreamIterator(lambda: self.read(n))

165

166 def iter_any(self) -> AsyncStreamIterator[bytes]:

167 """Yield all available data as soon as it is received."""

168 return AsyncStreamIterator(self.readany)

169

170 def iter_chunks(self) -> ChunkTupleAsyncStreamIterator:

171 """Yield chunks of data as they are received by the server.

172

173 The yielded objects are tuples

174 of (bytes, bool) as returned by the StreamReader.readchunk method.

175 """

176 return ChunkTupleAsyncStreamIterator(self)

177

178 def get_read_buffer_limits(self) -> tuple[int, int]:

179 return (self._low_water, self._high_water)

180

181 def set_read_chunk_size(self, n: int) -> None:

182 """Raise buffer limits to match the consumer's chunk size."""

183 if n > self._low_water:

184 self._low_water = n

185 self._high_water = n * 2

186

187 def exception(self) -> type[BaseException] | BaseException | None:

188 return self._exception

189

190 def set_exception(

191 self,

192 exc: type[BaseException] | BaseException,

193 exc_cause: BaseException = _EXC_SENTINEL,

194 ) -> None:

195 self._exception = exc

196 self._eof_callbacks.clear()

197

198 waiter = self._waiter

199 if waiter is not None:

200 self._waiter = None

201 set_exception(waiter, exc, exc_cause)

202

203 waiter = self._eof_waiter

204 if waiter is not None:

205 self._eof_waiter = None

206 set_exception(waiter, exc, exc_cause)

207

208 def on_eof(self, callback: Callable[[], None]) -> None:

209 if self._eof:

210 try:

211 callback()

212 except Exception:

213 internal_logger.exception("Exception in eof callback")

214 else:

215 self._eof_callbacks.append(callback)

216

217 def feed_eof(self) -> None:

218 self._eof = True

219

220 waiter = self._waiter

221 if waiter is not None:

222 self._waiter = None

223 set_result(waiter, None)

224

225 waiter = self._eof_waiter

226 if waiter is not None:

227 self._eof_waiter = None

228 set_result(waiter, None)

229

230 # At EOF the parser is done, there won't be unprocessed data.

231 self._protocol.resume_reading(resume_parser=False)

232

233 for cb in self._eof_callbacks:

234 try:

235 cb()

236 except Exception:

237 internal_logger.exception("Exception in eof callback")

238

239 self._eof_callbacks.clear()

240

241 def is_eof(self) -> bool:

242 """Return True if 'feed_eof' was called."""

243 return self._eof

244

245 def at_eof(self) -> bool:

246 """Return True if the buffer is empty and 'feed_eof' was called."""

247 return self._eof and not self._buffer

248

249 async def wait_eof(self) -> None:

250 if self._eof:

251 return

252

253 assert self._eof_waiter is None

254 self._eof_waiter = self._loop.create_future()

255 try:

256 await self._eof_waiter

257 finally:

258 self._eof_waiter = None

259

260 @property

261 def total_raw_bytes(self) -> int:

262 if self.total_compressed_bytes is None:

263 return self.total_bytes

264 return self.total_compressed_bytes

265

266 def unread_data(self, data: bytes) -> None:

267 """rollback reading some data from stream, inserting it to buffer head."""

268 warnings.warn(

269 "unread_data() is deprecated "

270 "and will be removed in future releases (#3260)",

271 DeprecationWarning,

272 stacklevel=2,

273 )

274 if not data:

275 return

276

277 if self._buffer_offset:

278 self._buffer[0] = self._buffer[0][self._buffer_offset :]

279 self._buffer_offset = 0

280 self._size += len(data)

281 self._cursor -= len(data)

282 self._buffer.appendleft(data)

283 self._eof_counter = 0

284

285 def feed_data(self, data: bytes) -> bool:

286 assert not self._eof, "feed_data after feed_eof"

287

288 if not data:

289 return False

290

291 data_len = len(data)

292 self._size += data_len

293 self._buffer.append(data)

294 self.total_bytes += data_len

295

296 waiter = self._waiter

297 if waiter is not None:

298 self._waiter = None

299 set_result(waiter, None)

300

301 if self._size > self._high_water:

302 self._protocol.pause_reading()

303 return False

304

305 def begin_http_chunk_receiving(self) -> None:

306 if self._http_chunk_splits is None:

307 if self.total_bytes:

308 raise RuntimeError(

309 "Called begin_http_chunk_receiving when some data was already fed"

310 )

311 self._http_chunk_splits = collections.deque()

312

313 def end_http_chunk_receiving(self) -> None:

314 if self._http_chunk_splits is None:

315 raise RuntimeError(

316 "Called end_chunk_receiving without calling "

317 "begin_chunk_receiving first"

318 )

319

320 # self._http_chunk_splits contains logical byte offsets from start of

321 # the body transfer. Each offset is the offset of the end of a chunk.

322 # "Logical" means bytes, accessible for a user.

323 # If no chunks containing logical data were received, current position

324 # is difinitely zero.

325 pos = self._http_chunk_splits[-1] if self._http_chunk_splits else 0

326

327 if self.total_bytes == pos:

328 # We should not add empty chunks here. So we check for that.

329 # Note, when chunked + gzip is used, we can receive a chunk

330 # of compressed data, but that data may not be enough for gzip FSM

331 # to yield any uncompressed data. That's why current position may

332 # not change after receiving a chunk.

333 return

334

335 self._http_chunk_splits.append(self.total_bytes)

336

337 # If we get too many small chunks before self._high_water is reached, then any

338 # .read() call becomes computationally expensive, and could block the event loop

339 # for too long, hence an additional self._high_water_chunks here.

340 if len(self._http_chunk_splits) > self._high_water_chunks:

341 self._protocol.pause_reading()

342

343 # wake up readchunk when end of http chunk received

344 waiter = self._waiter

345 if waiter is not None:

346 self._waiter = None

347 set_result(waiter, None)

348

349 async def _wait(self, func_name: str) -> None:

350 if not self._protocol.connected:

351 raise RuntimeError("Connection closed.")

352

353 # StreamReader uses a future to link the protocol feed_data() method

354 # to a read coroutine. Running two read coroutines at the same time

355 # would have an unexpected behaviour. It would not possible to know

356 # which coroutine would get the next data.

357 if self._waiter is not None:

358 raise RuntimeError(

359 "%s() called while another coroutine is "

360 "already waiting for incoming data" % func_name

361 )

362

363 waiter = self._waiter = self._loop.create_future()

364 try:

365 with self._timer:

366 await waiter

367 finally:

368 self._waiter = None

369

370 async def _fire_chunk_received(self, chunk: bytes) -> None:

371 cb = self._on_chunk_received

372 assert cb is not None

373 # Run under the same per-stream timer that _wait() uses, so a hung

374 # trace handler is bounded by sock_read just like a hung socket read would be.

375 with self._timer:

376 await cb(chunk)

377

378 async def readline(self, *, max_line_length: int | None = None) -> bytes:

379 return await self.readuntil(max_size=max_line_length)

380

381 async def readuntil(

382 self, separator: bytes = b"\n", *, max_size: int | None = None

383 ) -> bytes:

384 seplen = len(separator)

385 if seplen == 0:

386 raise ValueError("Separator should be at least one-byte string")

387

388 if self._exception is not None:

389 raise self._exception

390

391 chunk = b""

392 chunk_size = 0

393 not_enough = True

394 max_size = max_size or self._high_water

395

396 while not_enough:

397 while self._buffer and not_enough:

398 offset = self._buffer_offset

399 ichar = self._buffer[0].find(separator, offset) + 1

400 # Read from current offset to found separator or to the end.

401 data = self._read_nowait_chunk(

402 ichar - offset + seplen - 1 if ichar else -1

403 )

404 chunk += data

405 chunk_size += len(data)

406 if ichar:

407 not_enough = False

408

409 if chunk_size > max_size:

410 raise LineTooLong(chunk[:100] + b"...", max_size)

411

412 if self._eof:

413 break

414

415 if not_enough:

416 await self._wait("readuntil")

417

418 if chunk and self._on_chunk_received is not None:

419 await self._fire_chunk_received(chunk)

420 return chunk

421

422 async def read(self, n: int = -1) -> bytes:

423 if self._exception is not None:

424 raise self._exception

425

426 if not n:

427 return b""

428

429 if n < 0:

430 # Reading everything — remove decompression chunk limit.

431 # readany() fires the chunk hook for each block.

432 self.set_read_chunk_size(sys.maxsize)

433 blocks = []

434 while True:

435 block = await self.readany()

436 if not block:

437 break

438 blocks.append(block)

439 return b"".join(blocks)

440

441 self.set_read_chunk_size(n)

442 # TODO: should be `if` instead of `while`

443 # because waiter maybe triggered on chunk end,

444 # without feeding any data

445 while not self._buffer and not self._eof:

446 await self._wait("read")

447

448 chunk = self._read_nowait(n)

449 if chunk and self._on_chunk_received is not None:

450 await self._fire_chunk_received(chunk)

451 return chunk

452

453 async def readany(self) -> bytes:

454 if self._exception is not None:

455 raise self._exception

456

457 # TODO: should be `if` instead of `while`

458 # because waiter maybe triggered on chunk end,

459 # without feeding any data

460 while not self._buffer and not self._eof:

461 await self._wait("readany")

462

463 chunk = self._read_nowait(-1)

464 if chunk and self._on_chunk_received is not None:

465 await self._fire_chunk_received(chunk)

466 return chunk

467

468 async def readchunk(self) -> tuple[bytes, bool]:

469 """Returns a tuple of (data, end_of_http_chunk).

470

471 When chunked transfer

472 encoding is used, end_of_http_chunk is a boolean indicating if the end

473 of the data corresponds to the end of a HTTP chunk , otherwise it is

474 always False.

475 """

476 while True:

477 if self._exception is not None:

478 raise self._exception

479

480 while self._http_chunk_splits:

481 pos = self._http_chunk_splits.popleft()

482 if pos == self._cursor:

483 return (b"", True)

484 if pos > self._cursor:

485 chunk = self._read_nowait(pos - self._cursor)

486 if chunk and self._on_chunk_received is not None:

487 await self._fire_chunk_received(chunk)

488 return (chunk, True)

489 internal_logger.warning(

490 "Skipping HTTP chunk end due to data "

491 "consumption beyond chunk boundary"

492 )

493

494 if self._buffer:

495 chunk = self._read_nowait_chunk(-1)

496 if chunk and self._on_chunk_received is not None:

497 await self._fire_chunk_received(chunk)

498 return (chunk, False)

499 # return (self._read_nowait(-1), False)

500

501 if self._eof:

502 # Special case for signifying EOF.

503 # (b'', True) is not a final return value actually.

504 return (b"", False)

505

506 await self._wait("readchunk")

507

508 async def readexactly(self, n: int) -> bytes:

509 if self._exception is not None:

510 raise self._exception

511

512 blocks: list[bytes] = []

513 while n > 0:

514 block = await self.read(n)

515 if not block:

516 partial = b"".join(blocks)

517 raise asyncio.IncompleteReadError(partial, len(partial) + n)

518 blocks.append(block)

519 n -= len(block)

520

521 return b"".join(blocks)

522

523 def read_nowait(self, n: int = -1) -> bytes:

524 # default was changed to be consistent with .read(-1)

525 #

526 # I believe the most users don't know about the method and

527 # they are not affected.

528 if self._exception is not None:

529 raise self._exception

530

531 if self._waiter and not self._waiter.done():

532 raise RuntimeError(

533 "Called while some coroutine is waiting for incoming data."

534 )

535

536 chunk = self._read_nowait(n)

537 if chunk and (cb := self._on_chunk_received) is not None:

538 # read_nowait is sync but the hook is async; schedule it so the

539 # observability event still fires.

540 # TODO: Save and await this task.

541 asyncio.create_task(cb(chunk)) # type: ignore[unused-awaitable]

542 return chunk

543

544 def _read_nowait_chunk(self, n: int) -> bytes:

545 first_buffer = self._buffer[0]

546 offset = self._buffer_offset

547 if n != -1 and len(first_buffer) - offset > n:

548 data = first_buffer[offset : offset + n]

549 self._buffer_offset += n

550

551 elif offset:

552 self._buffer.popleft()

553 data = first_buffer[offset:]

554 self._buffer_offset = 0

555

556 else:

557 data = self._buffer.popleft()

558

559 data_len = len(data)

560 self._size -= data_len

561 self._cursor += data_len

562

563 chunk_splits = self._http_chunk_splits

564 # Prevent memory leak: drop useless chunk splits

565 while chunk_splits and chunk_splits[0] < self._cursor:

566 chunk_splits.popleft()

567

568 if self._size < self._low_water and (

569 self._http_chunk_splits is None

570 or len(self._http_chunk_splits) < self._low_water_chunks

571 ):

572 self._protocol.resume_reading()

573 return data

574

575 def _read_nowait(self, n: int) -> bytes:

576 """Read not more than n bytes, or whole buffer if n == -1"""

577 self._timer.assert_timeout()

578

579 if n == -1:

580 # Drain only chunks present now; _read_nowait_chunk() can

581 # re-entrantly resume_reading() and refill the buffer.

582 count = len(self._buffer)

583 if count == 1:

584 return self._read_nowait_chunk(-1)

585 return b"".join([self._read_nowait_chunk(-1) for _ in range(count)])

586

587 chunks: list[bytes] = []

588 while self._buffer:

589 chunk = self._read_nowait_chunk(n)

590 chunks.append(chunk)

591 n -= len(chunk)

592 if n == 0:

593 break

594

595 return b"".join(chunks) if chunks else b""

596

597

598class EmptyStreamReader(StreamReader): # lgtm [py/missing-call-to-init]

599

600 __slots__ = ("_read_eof_chunk",)

601

602 def __init__(self) -> None:

603 self._read_eof_chunk = False

604 self.total_bytes = 0

605

606 # Shadow the inherited slot with a property so the EMPTY_PAYLOAD singleton

607 # can't be polluted with a per-response hook that would leak across

608 # requests. EmptyStreamReader never delivers a chunk anyway.

609 @property

610 def _on_chunk_received(self) -> None:

611 return None

612

613 @_on_chunk_received.setter

614 def _on_chunk_received(

615 self, value: Callable[[bytes], Coroutine[None, None, None]] | None

616 ) -> None:

617 raise AttributeError("EmptyStreamReader._on_chunk_received is read-only")

618

619 def __repr__(self) -> str:

620 return "<%s>" % self.__class__.__name__

621

622 def exception(self) -> BaseException | None:

623 return None

624

625 def set_exception(

626 self,

627 exc: type[BaseException] | BaseException,

628 exc_cause: BaseException = _EXC_SENTINEL,

629 ) -> None:

630 pass

631

632 def on_eof(self, callback: Callable[[], None]) -> None:

633 try:

634 callback()

635 except Exception:

636 internal_logger.exception("Exception in eof callback")

637

638 def feed_eof(self) -> None:

639 pass

640

641 def is_eof(self) -> bool:

642 return True

643

644 def at_eof(self) -> bool:

645 return True

646

647 async def wait_eof(self) -> None:

648 return

649

650 def feed_data(self, data: bytes) -> bool:

651 return False

652

653 def set_read_chunk_size(self, n: int) -> None:

654 return

655

656 async def readline(self, *, max_line_length: int | None = None) -> bytes:

657 return b""

658

659 async def read(self, n: int = -1) -> bytes:

660 return b""

661

662 # TODO add async def readuntil

663

664 async def readany(self) -> bytes:

665 return b""

666

667 async def readchunk(self) -> tuple[bytes, bool]:

668 if not self._read_eof_chunk:

669 self._read_eof_chunk = True

670 return (b"", False)

671

672 return (b"", True)

673

674 async def readexactly(self, n: int) -> bytes:

675 raise asyncio.IncompleteReadError(b"", n)

676

677 def read_nowait(self, n: int = -1) -> bytes:

678 return b""

679

680

681EMPTY_PAYLOAD: Final[StreamReader] = EmptyStreamReader()

682

683

684class DataQueue(Generic[_T]):

685 """DataQueue is a general-purpose blocking queue with one reader."""

686

687 def __init__(self, loop: asyncio.AbstractEventLoop) -> None:

688 self._loop = loop

689 self._eof = False

690 self._waiter: asyncio.Future[None] | None = None

691 self._exception: type[BaseException] | BaseException | None = None

692 self._buffer: collections.deque[_T] = collections.deque()

693

694 def __len__(self) -> int:

695 return len(self._buffer)

696

697 def is_eof(self) -> bool:

698 return self._eof

699

700 def at_eof(self) -> bool:

701 return self._eof and not self._buffer

702

703 def exception(self) -> type[BaseException] | BaseException | None:

704 return self._exception

705

706 def set_exception(

707 self,

708 exc: type[BaseException] | BaseException,

709 exc_cause: BaseException = _EXC_SENTINEL,

710 ) -> None:

711 self._eof = True

712 self._exception = exc

713 if (waiter := self._waiter) is not None:

714 self._waiter = None

715 set_exception(waiter, exc, exc_cause)

716

717 def feed_data(self, data: _T) -> None:

718 self._buffer.append(data)

719 if (waiter := self._waiter) is not None:

720 self._waiter = None

721 set_result(waiter, None)

722

723 def feed_eof(self) -> None:

724 self._eof = True

725 if (waiter := self._waiter) is not None:

726 self._waiter = None

727 set_result(waiter, None)

728

729 async def read(self) -> _T:

730 if not self._buffer and not self._eof:

731 assert not self._waiter

732 self._waiter = self._loop.create_future()

733 try:

734 await self._waiter

735 except (asyncio.CancelledError, asyncio.TimeoutError):

736 self._waiter = None

737 raise

738 if self._buffer:

739 return self._buffer.popleft()

740 if self._exception is not None:

741 raise self._exception

742 raise EofStream

743

744 def __aiter__(self) -> AsyncStreamIterator[_T]:

745 return AsyncStreamIterator(self.read)