Skip to content

Commit

Permalink
propagate exceptions from cancel scope to nursery __aexit__ manually
Browse files Browse the repository at this point in the history
To further reduce extraneous exception frames, CancelScope gets
create() constructor and close() methods, which nursery __aexit__
uses directly rather than going through a scope context manager.
  • Loading branch information
belm0 committed Sep 1, 2018
1 parent c52a274 commit 1cb40da
Showing 1 changed file with 33 additions and 31 deletions.
64 changes: 33 additions & 31 deletions trio/_core/_run.py
Expand Up @@ -123,6 +123,16 @@ class CancelScope:
cancel_called = attr.ib(default=False)
cancelled_caught = attr.ib(default=False)

@staticmethod
def create(deadline, shield):
task = _core.current_task()
scope = CancelScope()
scope._scope_task = task
scope._add_task(task)
scope.deadline = deadline
scope.shield = shield
return scope

def __repr__(self):
return "<cancel scope object at {:#x}>".format(id(self))

Expand Down Expand Up @@ -211,6 +221,12 @@ def _exc_filter(self, exc):
return None
return exc

def _close(self, exc):
self._remove_task(self._scope_task)
if exc is not None:
filtered_exc = MultiError.filter(self._exc_filter, exc)
return filtered_exc

# Note we explicitly avoid @contextmanager since it adds extraneous stack
# frames to exceptions.
@enable_ki_protection
Expand All @@ -219,30 +235,21 @@ def __enter__(self):

@enable_ki_protection
def __exit__(self, etype, exc, tb):
try:
if exc is not None:
filtered_exc = MultiError.filter(self._exc_filter, exc)
if filtered_exc is None:
return True
elif filtered_exc == exc:
return False
else:
raise filtered_exc
finally:
self._remove_task(self._scope_task)
filtered_exc = self._close(exc)
if filtered_exc is None:
return True
elif filtered_exc == exc:
return False
else:
raise filtered_exc


def open_cancel_scope(*, deadline=inf, shield=False):
"""Returns a context manager which creates a new cancellation scope.
"""

scope = CancelScope()
scope._scope_task = _core.current_task()
scope._add_task(scope._scope_task)
scope.deadline = deadline
scope.shield = shield
return scope
return CancelScope.create(deadline, shield)


################################################################
Expand Down Expand Up @@ -348,27 +355,22 @@ class NurseryManager:

@enable_ki_protection
async def __aenter__(self):
self._scope_manager = open_cancel_scope()
scope = self._scope_manager.__enter__()
self._nursery = Nursery(current_task(), scope)
self._scope = CancelScope.create(deadline=inf, shield=False)
self._nursery = Nursery(current_task(), self._scope)
return self._nursery

@enable_ki_protection
async def __aexit__(self, etype, exc, tb):
new_exc = await self._nursery._nested_child_finished(exc)
if new_exc:
try:
if self._scope_manager.__exit__(
type(new_exc), new_exc, new_exc.__traceback__
):
return True
except BaseException as scope_manager_exc:
if scope_manager_exc == exc:
return False
raise # scope_manager_exc
raise new_exc
scope_exc = self._scope._close(new_exc)
if scope_exc is None:
return True
elif scope_exc == exc:
return False
else:
raise scope_exc
else:
self._scope_manager.__exit__(None, None, None)
return True

def __enter__(self):
Expand Down

0 comments on commit 1cb40da

Please sign in to comment.