Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/aiohttp/streams.py: 20%

1import asyncio

2import collections

3import sys

4import warnings

5from collections.abc import Awaitable, Callable

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 "total_bytes",

105 "total_compressed_bytes",

106 )

107

108 def __init__(

109 self,

110 protocol: BaseProtocol,

111 limit: int,

112 *,

113 timer: BaseTimerContext | None = None,

114 loop: asyncio.AbstractEventLoop | None = None,

115 ) -> None:

116 self._protocol = protocol

117 self._low_water = limit

118 self._high_water = limit * 2

119 if loop is None:

120 loop = asyncio.get_event_loop()

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

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

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

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

125 self._low_water_chunks = self._high_water_chunks // 2

126 self._loop = loop

127 self._size = 0

128 self._cursor = 0

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

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

131 self._buffer_offset = 0

132 self._eof = False

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

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

135 self._exception: BaseException | None = None

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

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

138 self._eof_counter = 0

139 self.total_bytes = 0

140 self.total_compressed_bytes: int | None = None

141

142 def __repr__(self) -> str:

143 info = [self.__class__.__name__]

144 if self._size:

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

146 if self._eof:

147 info.append("eof")

148 if self._low_water != DEFAULT_CHUNK_SIZE:

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

150 if self._waiter:

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

152 if self._exception:

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

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

155

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

157 return AsyncStreamIterator(self.readline)

158

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

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

161 self.set_read_chunk_size(n)

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

163

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

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

166 return AsyncStreamIterator(self.readany)

167

168 def iter_chunks(self) -> ChunkTupleAsyncStreamIterator:

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

170

171 The yielded objects are tuples

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

173 """

174 return ChunkTupleAsyncStreamIterator(self)

175

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

177 return (self._low_water, self._high_water)

178

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

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

181 if n > self._low_water:

182 self._low_water = n

183 self._high_water = n * 2

184

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

186 return self._exception

187

188 def set_exception(

189 self,

190 exc: BaseException,

191 exc_cause: BaseException = _EXC_SENTINEL,

192 ) -> None:

193 self._exception = exc

194 self._eof_callbacks.clear()

195

196 waiter = self._waiter

197 if waiter is not None:

198 self._waiter = None

199 set_exception(waiter, exc, exc_cause)

200

201 waiter = self._eof_waiter

202 if waiter is not None:

203 self._eof_waiter = None

204 set_exception(waiter, exc, exc_cause)

205

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

207 if self._eof:

208 try:

209 callback()

210 except Exception:

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

212 else:

213 self._eof_callbacks.append(callback)

214

215 def feed_eof(self) -> None:

216 self._eof = True

217

218 waiter = self._waiter

219 if waiter is not None:

220 self._waiter = None

221 set_result(waiter, None)

222

223 waiter = self._eof_waiter

224 if waiter is not None:

225 self._eof_waiter = None

226 set_result(waiter, None)

227

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

229 self._protocol.resume_reading(resume_parser=False)

230

231 for cb in self._eof_callbacks:

232 try:

233 cb()

234 except Exception:

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

236

237 self._eof_callbacks.clear()

238

239 def is_eof(self) -> bool:

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

241 return self._eof

242

243 def at_eof(self) -> bool:

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

245 return self._eof and not self._buffer

246

247 async def wait_eof(self) -> None:

248 if self._eof:

249 return

250

251 assert self._eof_waiter is None

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

253 try:

254 await self._eof_waiter

255 finally:

256 self._eof_waiter = None

257

258 @property

259 def total_raw_bytes(self) -> int:

260 if self.total_compressed_bytes is None:

261 return self.total_bytes

262 return self.total_compressed_bytes

263

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

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

266 warnings.warn(

267 "unread_data() is deprecated "

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

269 DeprecationWarning,

270 stacklevel=2,

271 )

272 if not data:

273 return

274

275 if self._buffer_offset:

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

277 self._buffer_offset = 0

278 self._size += len(data)

279 self._cursor -= len(data)

280 self._buffer.appendleft(data)

281 self._eof_counter = 0

282

283 # TODO: size is ignored, remove the param later

284 def feed_data(self, data: bytes, size: int = 0) -> bool:

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

286

287 if not data:

288 return False

289

290 data_len = len(data)

291 self._size += data_len

292 self._buffer.append(data)

293 self.total_bytes += data_len

294

295 waiter = self._waiter

296 if waiter is not None:

297 self._waiter = None

298 set_result(waiter, None)

299

300 if self._size > self._high_water:

301 self._protocol.pause_reading()

302 return False

303

304 def begin_http_chunk_receiving(self) -> None:

305 if self._http_chunk_splits is None:

306 if self.total_bytes:

307 raise RuntimeError(

308 "Called begin_http_chunk_receiving when some data was already fed"

309 )

310 self._http_chunk_splits = collections.deque()

311

312 def end_http_chunk_receiving(self) -> None:

313 if self._http_chunk_splits is None:

314 raise RuntimeError(

315 "Called end_chunk_receiving without calling "

316 "begin_chunk_receiving first"

317 )

318

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

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

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

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

323 # is difinitely zero.

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

325

326 if self.total_bytes == pos:

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

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

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

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

331 # not change after receiving a chunk.

332 return

333

334 self._http_chunk_splits.append(self.total_bytes)

335

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

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

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

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

340 self._protocol.pause_reading()

341

342 # wake up readchunk when end of http chunk received

343 waiter = self._waiter

344 if waiter is not None:

345 self._waiter = None

346 set_result(waiter, None)

347

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

349 if not self._protocol.connected:

350 raise RuntimeError("Connection closed.")

351

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

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

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

355 # which coroutine would get the next data.

356 if self._waiter is not None:

357 raise RuntimeError(

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

359 "already waiting for incoming data" % func_name

360 )

361

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

363 try:

364 with self._timer:

365 await waiter

366 finally:

367 self._waiter = None

368

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

370 return await self.readuntil(max_size=max_line_length)

371

372 async def readuntil(

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

374 ) -> bytes:

375 seplen = len(separator)

376 if seplen == 0:

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

378

379 if self._exception is not None:

380 raise self._exception

381

382 chunk = b""

383 chunk_size = 0

384 not_enough = True

385 max_size = max_size or self._high_water

386

387 while not_enough:

388 while self._buffer and not_enough:

389 offset = self._buffer_offset

390 n = -1

391 if chunk and seplen > 1:

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

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

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

395 tail = chunk[1 - seplen :]

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

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

398 if ichar != -1:

399 n = ichar + seplen - len(tail)

400 if n == -1:

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

402 if ichar != -1:

403 n = ichar - offset + seplen

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

405 data = self._read_nowait_chunk(n)

406 chunk += data

407 chunk_size += len(data)

408 if n != -1:

409 not_enough = False

410

411 if chunk_size > max_size:

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

413

414 if self._eof:

415 break

416

417 if not_enough:

418 await self._wait("readuntil")

419

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 # migration problem; with DataQueue you have to catch

427 # EofStream exception, so common way is to run payload.read() inside

428 # infinite loop. what can cause real infinite loop with StreamReader

429 # lets keep this code one major release.

430 if __debug__:

431 if self._eof and not self._buffer:

432 self._eof_counter = getattr(self, "_eof_counter", 0) + 1

433 if self._eof_counter > 5:

434 internal_logger.warning(

435 "Multiple access to StreamReader in eof state, "

436 "might be infinite loop.",

437 stack_info=True,

438 )

439

440 if not n:

441 return b""

442

443 if n < 0:

444 # Reading everything — remove decompression chunk limit.

445 self.set_read_chunk_size(sys.maxsize)

446 blocks = []

447 while True:

448 block = await self.readany()

449 if not block:

450 break

451 blocks.append(block)

452 return b"".join(blocks)

453

454 self.set_read_chunk_size(n)

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

456 # because waiter maybe triggered on chunk end,

457 # without feeding any data

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

459 await self._wait("read")

460

461 return self._read_nowait(n)

462

463 async def readany(self) -> bytes:

464 if self._exception is not None:

465 raise self._exception

466

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

468 # because waiter maybe triggered on chunk end,

469 # without feeding any data

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

471 await self._wait("readany")

472

473 return self._read_nowait(-1)

474

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

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

477

478 When chunked transfer

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

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

481 always False.

482 """

483 while True:

484 if self._exception is not None:

485 raise self._exception

486

487 while self._http_chunk_splits:

488 pos = self._http_chunk_splits.popleft()

489 if pos == self._cursor:

490 return (b"", True)

491 if pos > self._cursor:

492 return (self._read_nowait(pos - self._cursor), True)

493 internal_logger.warning(

494 "Skipping HTTP chunk end due to data "

495 "consumption beyond chunk boundary"

496 )

497

498 if self._buffer:

499 return (self._read_nowait_chunk(-1), False)

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

501

502 if self._eof:

503 # Special case for signifying EOF.

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

505 return (b"", False)

506

507 await self._wait("readchunk")

508

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

510 if self._exception is not None:

511 raise self._exception

512

513 blocks: list[bytes] = []

514 while n > 0:

515 block = await self.read(n)

516 if not block:

517 partial = b"".join(blocks)

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

519 blocks.append(block)

520 n -= len(block)

521

522 return b"".join(blocks)

523

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

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

526 #

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

528 # they are not affected.

529 if self._exception is not None:

530 raise self._exception

531

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

533 raise RuntimeError(

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

535 )

536

537 return self._read_nowait(n)

538

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

540 first_buffer = self._buffer[0]

541 offset = self._buffer_offset

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

543 data = first_buffer[offset : offset + n]

544 self._buffer_offset += n

545

546 elif offset:

547 self._buffer.popleft()

548 data = first_buffer[offset:]

549 self._buffer_offset = 0

550

551 else:

552 data = self._buffer.popleft()

553

554 data_len = len(data)

555 self._size -= data_len

556 self._cursor += data_len

557

558 chunk_splits = self._http_chunk_splits

559 # Prevent memory leak: drop useless chunk splits

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

561 chunk_splits.popleft()

562

563 if self._size < self._low_water and (

564 self._http_chunk_splits is None

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

566 ):

567 self._protocol.resume_reading()

568 return data

569

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

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

572 self._timer.assert_timeout()

573

574 if n == -1:

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

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

577 count = len(self._buffer)

578 if count == 1:

579 return self._read_nowait_chunk(-1)

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

581

582 chunks: list[bytes] = []

583 while self._buffer:

584 chunk = self._read_nowait_chunk(n)

585 chunks.append(chunk)

586 n -= len(chunk)

587 if n == 0:

588 break

589

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

591

592

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

594

595 __slots__ = ("_read_eof_chunk",)

596

597 def __init__(self) -> None:

598 self._read_eof_chunk = False

599 self.total_bytes = 0

600

601 def __repr__(self) -> str:

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

603

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

605 return None

606

607 def set_exception(

608 self,

609 exc: BaseException,

610 exc_cause: BaseException = _EXC_SENTINEL,

611 ) -> None:

612 pass

613

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

615 try:

616 callback()

617 except Exception:

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

619

620 def feed_eof(self) -> None:

621 pass

622

623 def is_eof(self) -> bool:

624 return True

625

626 def at_eof(self) -> bool:

627 return True

628

629 async def wait_eof(self) -> None:

630 return

631

632 def feed_data(self, data: bytes, n: int = 0) -> bool:

633 return False

634

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

636 return

637

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

639 return b""

640

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

642 return b""

643

644 # TODO add async def readuntil

645

646 async def readany(self) -> bytes:

647 return b""

648

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

650 if not self._read_eof_chunk:

651 self._read_eof_chunk = True

652 return (b"", False)

653

654 return (b"", True)

655

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

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

658

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

660 return b""

661

662

663EMPTY_PAYLOAD: Final[StreamReader] = EmptyStreamReader()

664

665

666class DataQueue(Generic[_T]):

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

668

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

670 self._loop = loop

671 self._eof = False

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

673 self._exception: BaseException | None = None

674 self._buffer: collections.deque[tuple[_T, int]] = collections.deque()

675

676 def __len__(self) -> int:

677 return len(self._buffer)

678

679 def is_eof(self) -> bool:

680 return self._eof

681

682 def at_eof(self) -> bool:

683 return self._eof and not self._buffer

684

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

686 return self._exception

687

688 def set_exception(

689 self,

690 exc: BaseException,

691 exc_cause: BaseException = _EXC_SENTINEL,

692 ) -> None:

693 self._eof = True

694 self._exception = exc

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

696 self._waiter = None

697 set_exception(waiter, exc, exc_cause)

698

699 def feed_data(self, data: _T, size: int = 0) -> None:

700 self._buffer.append((data, size))

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

702 self._waiter = None

703 set_result(waiter, None)

704

705 def feed_eof(self) -> None:

706 self._eof = True

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

708 self._waiter = None

709 set_result(waiter, None)

710

711 async def read(self) -> _T:

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

713 assert not self._waiter

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

715 try:

716 await self._waiter

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

718 self._waiter = None

719 raise

720 if self._buffer:

721 data, _ = self._buffer.popleft()

722 return data

723 if self._exception is not None:

724 raise self._exception

725 raise EofStream

726

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

728 return AsyncStreamIterator(self.read)

729

730

731class FlowControlDataQueue(DataQueue[_T]):

732 """FlowControlDataQueue resumes and pauses an underlying stream.

733

734 It is a destination for parsed data.

735

736 This class is deprecated and will be removed in version 4.0.

737 """

738

739 def __init__(

740 self, protocol: BaseProtocol, limit: int, *, loop: asyncio.AbstractEventLoop

741 ) -> None:

742 super().__init__(loop=loop)

743 self._size = 0

744 self._protocol = protocol

745 self._limit = limit * 2

746

747 def feed_data(self, data: _T, size: int = 0) -> None:

748 super().feed_data(data, size)

749 self._size += size

750

751 if self._size > self._limit and not self._protocol._reading_paused:

752 self._protocol.pause_reading()

753

754 async def read(self) -> _T:

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

756 assert not self._waiter

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

758 try:

759 await self._waiter

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

761 self._waiter = None

762 raise

763 if self._buffer:

764 data, size = self._buffer.popleft()

765 self._size -= size

766 if self._size < self._limit and self._protocol._reading_paused:

767 self._protocol.resume_reading()

768 return data

769 if self._exception is not None:

770 raise self._exception

771 raise EofStream