bpo-29302: Implement contextlib.AsyncExitStack. (#4790) · python/cpython@1aa094f (original) (raw)

`@@ -7,7 +7,7 @@

7

8

`all = ["asynccontextmanager", "contextmanager", "closing", "nullcontext",

9

`"AbstractContextManager", "AbstractAsyncContextManager",

10

"ContextDecorator", "ExitStack",

10

"AsyncExitStack", "ContextDecorator", "ExitStack",

11

`"redirect_stdout", "redirect_stderr", "suppress"]

12

13

`@@ -365,85 +365,102 @@ def exit(self, exctype, excinst, exctb):

365

`return exctype is not None and issubclass(exctype, self._exceptions)

366

367

368

Inspired by discussions on http://bugs.python.org/issue13585

369

class ExitStack(AbstractContextManager):

370

"""Context manager for dynamic management of a stack of exit callbacks

368

class _BaseExitStack:

369

"""A base class for ExitStack and AsyncExitStack."""

371

370

372

For example:

371

@staticmethod

372

def _create_exit_wrapper(cm, cm_exit):

373

def _exit_wrapper(exc_type, exc, tb):

374

return cm_exit(cm, exc_type, exc, tb)

375

return _exit_wrapper

373

376

374

with ExitStack() as stack:

375

files = [stack.enter_context(open(fname)) for fname in filenames]

376

All opened files will automatically be closed at the end of

377

the with statement, even if attempts to open files later

378

in the list raise an exception

377

@staticmethod

378

def _create_cb_wrapper(callback, *args, **kwds):

379

def _exit_wrapper(exc_type, exc, tb):

380

callback(*args, **kwds)

381

return _exit_wrapper

379

382

380

"""

381

383

`def init(self):

382

384

`self._exit_callbacks = deque()

383

385

384

386

`def pop_all(self):

385

"""Preserve the context stack by transferring it to a new instance"""

387

"""Preserve the context stack by transferring it to a new instance."""

386

388

`new_stack = type(self)()

387

389

`new_stack._exit_callbacks = self._exit_callbacks

388

390

`self._exit_callbacks = deque()

389

391

`return new_stack

390

392

391

def _push_cm_exit(self, cm, cm_exit):

392

"""Helper to correctly register callbacks to exit methods"""

393

def _exit_wrapper(*exc_details):

394

return cm_exit(cm, *exc_details)

395

_exit_wrapper.self = cm

396

self.push(_exit_wrapper)

397

-

398

393

`def push(self, exit):

399

"""Registers a callback with the standard exit method signature

400

-

401

Can suppress exceptions the same way exit methods can.

394

"""Registers a callback with the standard exit method signature.

402

395

396

Can suppress exceptions the same way exit method can.

403

397

` Also accepts any object with an exit method (registering a call

404

to the method instead of the object itself)

398

to the method instead of the object itself).

405

399

` """

406

400

`# We use an unbound method rather than a bound method to follow

407

the standard lookup behaviour for special methods

401

the standard lookup behaviour for special methods.

408

402

`_cb_type = type(exit)

403

+

409

404

`try:

410

405

`exit_method = _cb_type.exit

411

406

`except AttributeError:

412

Not a context manager, so assume its a callable

413

self._exit_callbacks.append(exit)

407

Not a context manager, so assume it's a callable.

408

self._push_exit_callback(exit)

414

409

`else:

415

410

`self._push_cm_exit(exit, exit_method)

416

return exit # Allow use as a decorator

417

-

418

def callback(self, callback, *args, **kwds):

419

"""Registers an arbitrary callback and arguments.

420

-

421

Cannot suppress exceptions.

422

"""

423

def _exit_wrapper(exc_type, exc, tb):

424

callback(*args, **kwds)

425

We changed the signature, so using @wraps is not appropriate, but

426

setting wrapped may still help with introspection

427

_exit_wrapper.wrapped = callback

428

self.push(_exit_wrapper)

429

return callback # Allow use as a decorator

411

return exit # Allow use as a decorator.

430

412

431

413

`def enter_context(self, cm):

432

"""Enters the supplied context manager

414

"""Enters the supplied context manager.

433

415

434

416

` If successful, also pushes its exit method as a callback and

435

417

` returns the result of the enter method.

436

418

` """

437

We look up the special methods on the type to match the with statement

419

We look up the special methods on the type to match the with

420

statement.

438

421

`_cm_type = type(cm)

439

422

`_exit = _cm_type.exit

440

423

`result = _cm_type.enter(cm)

441

424

`self._push_cm_exit(cm, _exit)

442

425

`return result

443

426

444

def close(self):

445

"""Immediately unwind the context stack"""

446

self.exit(None, None, None)

427

def callback(self, callback, *args, **kwds):

428

"""Registers an arbitrary callback and arguments.

429

+

430

Cannot suppress exceptions.

431

"""

432

_exit_wrapper = self._create_cb_wrapper(callback, *args, **kwds)

433

+

434

We changed the signature, so using @wraps is not appropriate, but

435

setting wrapped may still help with introspection.

436

_exit_wrapper.wrapped = callback

437

self._push_exit_callback(_exit_wrapper)

438

return callback # Allow use as a decorator

439

+

440

def _push_cm_exit(self, cm, cm_exit):

441

"""Helper to correctly register callbacks to exit methods."""

442

_exit_wrapper = self._create_exit_wrapper(cm, cm_exit)

443

_exit_wrapper.self = cm

444

self._push_exit_callback(_exit_wrapper, True)

445

+

446

def _push_exit_callback(self, callback, is_sync=True):

447

self._exit_callbacks.append((is_sync, callback))

448

+

449

+

450

Inspired by discussions on http://bugs.python.org/issue13585

451

class ExitStack(_BaseExitStack, AbstractContextManager):

452

"""Context manager for dynamic management of a stack of exit callbacks.

453

+

454

For example:

455

with ExitStack() as stack:

456

files = [stack.enter_context(open(fname)) for fname in filenames]

457

All opened files will automatically be closed at the end of

458

the with statement, even if attempts to open files later

459

in the list raise an exception.

460

"""

461

+

462

def enter(self):

463

return self

447

464

448

465

`def exit(self, *exc_details):

449

466

`received_exc = exc_details[0] is not None

`@@ -470,7 +487,8 @@ def _fix_exception_context(new_exc, old_exc):

470

487

`suppressed_exc = False

471

488

`pending_raise = False

472

489

`while self._exit_callbacks:

473

cb = self._exit_callbacks.pop()

490

is_sync, cb = self._exit_callbacks.pop()

491

assert is_sync

474

492

`try:

475

493

`if cb(*exc_details):

476

494

`suppressed_exc = True

`@@ -493,6 +511,147 @@ def _fix_exception_context(new_exc, old_exc):

493

511

`raise

494

512

`return received_exc and suppressed_exc

495

513

514

def close(self):

515

"""Immediately unwind the context stack."""

516

self.exit(None, None, None)

517

+

518

+

519

Inspired by discussions on https://bugs.python.org/issue29302

520

class AsyncExitStack(_BaseExitStack, AbstractAsyncContextManager):

521

"""Async context manager for dynamic management of a stack of exit

522

callbacks.

523

+

524

For example:

525

async with AsyncExitStack() as stack:

526

connections = [await stack.enter_async_context(get_connection())

527

for i in range(5)]

528

All opened connections will automatically be released at the

529

end of the async with statement, even if attempts to open a

530

connection later in the list raise an exception.

531

"""

532

+

533

@staticmethod

534

def _create_async_exit_wrapper(cm, cm_exit):

535

async def _exit_wrapper(exc_type, exc, tb):

536

return await cm_exit(cm, exc_type, exc, tb)

537

return _exit_wrapper

538

+

539

@staticmethod

540

def _create_async_cb_wrapper(callback, *args, **kwds):

541

async def _exit_wrapper(exc_type, exc, tb):

542

await callback(*args, **kwds)

543

return _exit_wrapper

544

+

545

async def enter_async_context(self, cm):

546

"""Enters the supplied async context manager.

547

+

548

If successful, also pushes its aexit method as a callback and

549

returns the result of the aenter method.

550

"""

551

_cm_type = type(cm)

552

_exit = _cm_type.aexit

553

result = await _cm_type.aenter(cm)

554

self._push_async_cm_exit(cm, _exit)

555

return result

556

+

557

def push_async_exit(self, exit):

558

"""Registers a coroutine function with the standard aexit method

559

signature.

560

+

561

Can suppress exceptions the same way aexit method can.

562

Also accepts any object with an aexit method (registering a call

563

to the method instead of the object itself).

564

"""

565

_cb_type = type(exit)

566

try:

567

exit_method = _cb_type.aexit

568

except AttributeError:

569

Not an async context manager, so assume it's a coroutine function

570

self._push_exit_callback(exit, False)

571

else:

572

self._push_async_cm_exit(exit, exit_method)

573

return exit # Allow use as a decorator

574

+

575

def push_async_callback(self, callback, *args, **kwds):

576

"""Registers an arbitrary coroutine function and arguments.

577

+

578

Cannot suppress exceptions.

579

"""

580

_exit_wrapper = self._create_async_cb_wrapper(callback, *args, **kwds)

581

+

582

We changed the signature, so using @wraps is not appropriate, but

583

setting wrapped may still help with introspection.

584

_exit_wrapper.wrapped = callback

585

self._push_exit_callback(_exit_wrapper, False)

586

return callback # Allow use as a decorator

587

+

588

async def aclose(self):

589

"""Immediately unwind the context stack."""

590

await self.aexit(None, None, None)

591

+

592

def _push_async_cm_exit(self, cm, cm_exit):

593

"""Helper to correctly register coroutine function to aexit

594

method."""

595

_exit_wrapper = self._create_async_exit_wrapper(cm, cm_exit)

596

_exit_wrapper.self = cm

597

self._push_exit_callback(_exit_wrapper, False)

598

+

599

async def aenter(self):

600

return self

601

+

602

async def aexit(self, *exc_details):

603

received_exc = exc_details[0] is not None

604

+

605

We manipulate the exception state so it behaves as though

606

we were actually nesting multiple with statements

607

frame_exc = sys.exc_info()[1]

608

def _fix_exception_context(new_exc, old_exc):

609

Context may not be correct, so find the end of the chain

610

while 1:

611

exc_context = new_exc.context

612

if exc_context is old_exc:

613

Context is already set correctly (see issue 20317)

614

return

615

if exc_context is None or exc_context is frame_exc:

616

break

617

new_exc = exc_context

618

Change the end of the chain to point to the exception

619

we expect it to reference

620

new_exc.context = old_exc

621

+

622

Callbacks are invoked in LIFO order to match the behaviour of

623

nested context managers

624

suppressed_exc = False

625

pending_raise = False

626

while self._exit_callbacks:

627

is_sync, cb = self._exit_callbacks.pop()

628

try:

629

if is_sync:

630

cb_suppress = cb(*exc_details)

631

else:

632

cb_suppress = await cb(*exc_details)

633

+

634

if cb_suppress:

635

suppressed_exc = True

636

pending_raise = False

637

exc_details = (None, None, None)

638

except:

639

new_exc_details = sys.exc_info()

640

simulate the stack of exceptions by setting the context

641

_fix_exception_context(new_exc_details[1], exc_details[1])

642

pending_raise = True

643

exc_details = new_exc_details

644

if pending_raise:

645

try:

646

bare "raise exc_details[1]" replaces our carefully

647

set-up context

648

fixed_ctx = exc_details[1].context

649

raise exc_details[1]

650

except BaseException:

651

exc_details[1].context = fixed_ctx

652

raise

653

return received_exc and suppressed_exc

654

+

496

655

497

656

`class nullcontext(AbstractContextManager):

498

657

`"""Context manager that does no additional processing.