3
3
4
4
__all__ = ["TaskGroup" ]
5
5
6
- import weakref
7
-
8
6
from . import events
9
7
from . import exceptions
10
8
from . import tasks
@@ -19,8 +17,7 @@ def __init__(self):
19
17
self ._loop = None
20
18
self ._parent_task = None
21
19
self ._parent_cancel_requested = False
22
- self ._tasks = weakref .WeakSet ()
23
- self ._unfinished_tasks = 0
20
+ self ._tasks = set ()
24
21
self ._errors = []
25
22
self ._base_error = None
26
23
self ._on_completed_fut = None
@@ -29,8 +26,6 @@ def __repr__(self):
29
26
info = ['' ]
30
27
if self ._tasks :
31
28
info .append (f'tasks={ len (self ._tasks )} ' )
32
- if self ._unfinished_tasks :
33
- info .append (f'unfinished={ self ._unfinished_tasks } ' )
34
29
if self ._errors :
35
30
info .append (f'errors={ len (self ._errors )} ' )
36
31
if self ._aborting :
@@ -93,7 +88,7 @@ async def __aexit__(self, et, exc, tb):
93
88
# can be cancelled multiple times if our parent task
94
89
# is being cancelled repeatedly (or even once, when
95
90
# our own cancellation is already in progress)
96
- while self ._unfinished_tasks :
91
+ while self ._tasks :
97
92
if self ._on_completed_fut is None :
98
93
self ._on_completed_fut = self ._loop .create_future ()
99
94
@@ -114,7 +109,7 @@ async def __aexit__(self, et, exc, tb):
114
109
115
110
self ._on_completed_fut = None
116
111
117
- assert self ._unfinished_tasks == 0
112
+ assert not self ._tasks
118
113
119
114
if self ._base_error is not None :
120
115
raise self ._base_error
@@ -141,15 +136,14 @@ async def __aexit__(self, et, exc, tb):
141
136
def create_task (self , coro , * , name = None , context = None ):
142
137
if not self ._entered :
143
138
raise RuntimeError (f"TaskGroup { self !r} has not been entered" )
144
- if self ._exiting and self ._unfinished_tasks == 0 :
139
+ if self ._exiting and not self ._tasks :
145
140
raise RuntimeError (f"TaskGroup { self !r} is finished" )
146
141
if context is None :
147
142
task = self ._loop .create_task (coro )
148
143
else :
149
144
task = self ._loop .create_task (coro , context = context )
150
145
tasks ._set_task_name (task , name )
151
146
task .add_done_callback (self ._on_task_done )
152
- self ._unfinished_tasks += 1
153
147
self ._tasks .add (task )
154
148
return task
155
149
@@ -169,10 +163,9 @@ def _abort(self):
169
163
t .cancel ()
170
164
171
165
def _on_task_done (self , task ):
172
- self ._unfinished_tasks -= 1
173
- assert self ._unfinished_tasks >= 0
166
+ self ._tasks .discard (task )
174
167
175
- if self ._on_completed_fut is not None and not self ._unfinished_tasks :
168
+ if self ._on_completed_fut is not None and not self ._tasks :
176
169
if not self ._on_completed_fut .done ():
177
170
self ._on_completed_fut .set_result (True )
178
171
0 commit comments