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

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 n = -1

400 if chunk and seplen > 1:

401 # The separator may straddle data already read and this chunk.

402 # tail and head are each shorter than the separator, so a match

403 # in tail + head must span both; n counts its bytes in head.

404 tail = chunk[1 - seplen :]

405 head = self._buffer[0][offset : offset + seplen - 1]

406 ichar = (tail + head).find(separator)

407 if ichar != -1:

408 n = ichar + seplen - len(tail)

409 if n == -1:

410 ichar = self._buffer[0].find(separator, offset)

411 if ichar != -1:

412 n = ichar - offset + seplen

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

414 data = self._read_nowait_chunk(n)

415 chunk += data

416 chunk_size += len(data)

417 if n != -1:

418 not_enough = False

419

420 if chunk_size > max_size:

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

422

423 if self._eof:

424 break

425

426 if not_enough:

427 await self._wait("readuntil")

428

429 if chunk and self._on_chunk_received is not None:

430 await self._fire_chunk_received(chunk)

431 return chunk

432

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

434 if self._exception is not None:

435 raise self._exception

436

437 if not n:

438 return b""

439

440 if n < 0:

441 # Reading everything — remove decompression chunk limit.

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

443 self.set_read_chunk_size(sys.maxsize)

444 blocks = []

445 while True:

446 block = await self.readany()

447 if not block:

448 break

449 blocks.append(block)

450 return b"".join(blocks)

451

452 self.set_read_chunk_size(n)

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

454 # because waiter maybe triggered on chunk end,

455 # without feeding any data

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

457 await self._wait("read")

458

459 chunk = self._read_nowait(n)

460 if chunk and self._on_chunk_received is not None:

461 await self._fire_chunk_received(chunk)

462 return chunk

463

464 async def readany(self) -> bytes:

465 if self._exception is not None:

466 raise self._exception

467

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

469 # because waiter maybe triggered on chunk end,

470 # without feeding any data

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

472 await self._wait("readany")

473

474 chunk = self._read_nowait(-1)

475 if chunk and self._on_chunk_received is not None:

476 await self._fire_chunk_received(chunk)

477 return chunk

478

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

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

481

482 When chunked transfer

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

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

485 always False.

486 """

487 while True:

488 if self._exception is not None:

489 raise self._exception

490

491 while self._http_chunk_splits:

492 pos = self._http_chunk_splits.popleft()

493 if pos == self._cursor:

494 return (b"", True)

495 if pos > self._cursor:

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

497 if chunk and self._on_chunk_received is not None:

498 await self._fire_chunk_received(chunk)

499 return (chunk, True)

500 internal_logger.warning(

501 "Skipping HTTP chunk end due to data "

502 "consumption beyond chunk boundary"

503 )

504

505 if self._buffer:

506 chunk = self._read_nowait_chunk(-1)

507 if chunk and self._on_chunk_received is not None:

508 await self._fire_chunk_received(chunk)

509 return (chunk, False)

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

511

512 if self._eof:

513 # Special case for signifying EOF.

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

515 return (b"", False)

516

517 await self._wait("readchunk")

518

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

520 if self._exception is not None:

521 raise self._exception

522

523 blocks: list[bytes] = []

524 while n > 0:

525 block = await self.read(n)

526 if not block:

527 partial = b"".join(blocks)

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

529 blocks.append(block)

530 n -= len(block)

531

532 return b"".join(blocks)

533

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

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

536 #

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

538 # they are not affected.

539 if self._exception is not None:

540 raise self._exception

541

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

543 raise RuntimeError(

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

545 )

546

547 chunk = self._read_nowait(n)

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

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

550 # observability event still fires.

551 # TODO: Save and await this task.

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

553 return chunk

554

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

556 first_buffer = self._buffer[0]

557 offset = self._buffer_offset

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

559 data = first_buffer[offset : offset + n]

560 self._buffer_offset += n

561

562 elif offset:

563 self._buffer.popleft()

564 data = first_buffer[offset:]

565 self._buffer_offset = 0

566

567 else:

568 data = self._buffer.popleft()

569

570 data_len = len(data)

571 self._size -= data_len

572 self._cursor += data_len

573

574 chunk_splits = self._http_chunk_splits

575 # Prevent memory leak: drop useless chunk splits

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

577 chunk_splits.popleft()

578

579 if self._size < self._low_water and (

580 self._http_chunk_splits is None

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

582 ):

583 self._protocol.resume_reading()

584 return data

585

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

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

588 self._timer.assert_timeout()

589

590 if n == -1:

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

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

593 count = len(self._buffer)

594 if count == 1:

595 return self._read_nowait_chunk(-1)

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

597

598 chunks: list[bytes] = []

599 while self._buffer:

600 chunk = self._read_nowait_chunk(n)

601 chunks.append(chunk)

602 n -= len(chunk)

603 if n == 0:

604 break

605

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

607

608

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

610

611 __slots__ = ("_read_eof_chunk",)

612

613 def __init__(self) -> None:

614 self._read_eof_chunk = False

615 self.total_bytes = 0

616

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

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

619 # requests. EmptyStreamReader never delivers a chunk anyway.

620 @property

621 def _on_chunk_received(self) -> None:

622 return None

623

624 @_on_chunk_received.setter

625 def _on_chunk_received(

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

627 ) -> None:

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

629

630 def __repr__(self) -> str:

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

632

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

634 return None

635

636 def set_exception(

637 self,

638 exc: type[BaseException] | BaseException,

639 exc_cause: BaseException = _EXC_SENTINEL,

640 ) -> None:

641 pass

642

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

644 try:

645 callback()

646 except Exception:

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

648

649 def feed_eof(self) -> None:

650 pass

651

652 def is_eof(self) -> bool:

653 return True

654

655 def at_eof(self) -> bool:

656 return True

657

658 async def wait_eof(self) -> None:

659 return

660

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

662 return False

663

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

665 return

666

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

668 return b""

669

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

671 return b""

672

673 # TODO add async def readuntil

674

675 async def readany(self) -> bytes:

676 return b""

677

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

679 if not self._read_eof_chunk:

680 self._read_eof_chunk = True

681 return (b"", False)

682

683 return (b"", True)

684

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

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

687

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

689 return b""

690

691

692EMPTY_PAYLOAD: Final[StreamReader] = EmptyStreamReader()

693

694

695class DataQueue(Generic[_T]):

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

697

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

699 self._loop = loop

700 self._eof = False

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

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

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

704

705 def __len__(self) -> int:

706 return len(self._buffer)

707

708 def is_eof(self) -> bool:

709 return self._eof

710

711 def at_eof(self) -> bool:

712 return self._eof and not self._buffer

713

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

715 return self._exception

716

717 def set_exception(

718 self,

719 exc: type[BaseException] | BaseException,

720 exc_cause: BaseException = _EXC_SENTINEL,

721 ) -> None:

722 self._eof = True

723 self._exception = exc

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

725 self._waiter = None

726 set_exception(waiter, exc, exc_cause)

727

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

729 self._buffer.append(data)

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

731 self._waiter = None

732 set_result(waiter, None)

733

734 def feed_eof(self) -> None:

735 self._eof = True

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

737 self._waiter = None

738 set_result(waiter, None)

739

740 async def read(self) -> _T:

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

742 assert not self._waiter

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

744 try:

745 await self._waiter

746 except (asyncio.CancelledError, TimeoutError):

747 self._waiter = None

748 raise

749 if self._buffer:

750 return self._buffer.popleft()

751 if self._exception is not None:

752 raise self._exception

753 raise EofStream

754

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

756 return AsyncStreamIterator(self.read)