cpython/Lib/asyncio/taskgroups.py at 3.15 · python/cpython · GitHub

Repository navigation

History

100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

116

117

118

119

120

121

122

123

124

125

126

127

128

129

130

131

132

133

134

135

136

137

138

139

140

141

142

143

144

145

146

147

148

149

150

151

152

153

154

155

156

157

158

159

160

161

162

163

164

165

166

167

168

169

170

171

172

173

174

175

176

177

178

179

180

181

182

183

184

185

186

187

188

189

190

191

192

193

194

195

196

197

198

199

200

201

202

203

204

205

206

207

208

209

210

211

212

213

214

215

216

217

218

219

220

221

222

223

224

225

226

227

228

229

230

231

232

233

234

235

236

237

238

239

240

241

242

243

244

245

246

247

248

249

250

251

252

253

254

255

256

257

258

259

260

261

262

263

264

265

266

267

268

269

270

271

272

273

274

275

276

277

278

279

280

281

282

283

284

285

286

287

288

289

290

291

292

293

294

295

296

297

298

299

300

301

302

303

304

305

306

307

308

309

310

311

312

313

314

315

316

317

318

319

320

321

322

323

324

325

326

327

328

329

330

331

332

333

334

# Adapted with permission from the EdgeDB project;

# license: PSFL.

__all__ = ("TaskGroup",)

from . import events

from . import exceptions

from . import futures

from . import tasks

class TaskGroup:

"""Asynchronous context manager for managing groups of tasks.

Example use:

async with asyncio.TaskGroup() as group:

task1 = group.create_task(some_coroutine(...))

task2 = group.create_task(other_coroutine(...))

print("Both tasks have completed now.")

All tasks are awaited when the context manager exits.

Any exceptions other than `asyncio.CancelledError` raised within

a task will cancel all remaining tasks and wait for them to exit.

The exceptions are then combined and raised as an `ExceptionGroup`.

"""

def __init__(self):

self._entered = False

self._exiting = False

self._aborting = False

self._loop = None

self._parent_task = None

self._parent_cancel_requested = False

self._tasks = set()

self._errors = []

self._base_error = None

self._on_completed_fut = None

self._cancel_on_enter = False

def __repr__(self):

info = ['']

if self._tasks:

info.append(f'tasks={len(self._tasks)}')

if self._errors:

info.append(f'errors={len(self._errors)}')

if self._aborting:

info.append('cancelling')

elif self._entered:

info.append('entered')

info_str = ' '.join(info)

return f'<TaskGroup{info_str}>'

async def __aenter__(self):

if self._entered:

raise RuntimeError(

f"TaskGroup {self!r} has already been entered")

if self._loop is None:

self._loop = events.get_running_loop()

self._parent_task = tasks.current_task(self._loop)

if self._parent_task is None:

raise RuntimeError(

f'TaskGroup {self!r} cannot determine the parent task')

self._entered = True

if self._cancel_on_enter:

self.cancel()

return self

async def __aexit__(self, et, exc, tb):

tb = None

try:

return await self._aexit(et, exc)

finally:

# Exceptions are heavy objects that can have object

# cycles (bad for GC); let's not keep a reference to

# a bunch of them. It would be nicer to use a try/finally

# in __aexit__ directly but that introduced some diff noise

self._parent_task = None

self._errors = None

self._base_error = None

exc = None

async def _aexit(self, et, exc):

self._exiting = True

if (exc is not None and

self._is_base_error(exc) and

self._base_error is None):

self._base_error = exc

if et is not None and issubclass(et, exceptions.CancelledError):

propagate_cancellation_error = exc

else:

propagate_cancellation_error = None

if et is not None:

if not self._aborting:

# Our parent task is being cancelled:

# async with TaskGroup() as g:

# g.create_task(...)

# await ... # <- CancelledError

# or there's an exception in "async with":

# async with TaskGroup() as g:

# g.create_task(...)

# 1 / 0

self._abort()

# We use while-loop here because "self._on_completed_fut"

# can be cancelled multiple times if our parent task

# is being cancelled repeatedly (or even once, when

# our own cancellation is already in progress)

pending_cancellation_error = None

while self._tasks:

if self._on_completed_fut is None:

self._on_completed_fut = self._loop.create_future()

try:

await self._on_completed_fut

except exceptions.CancelledError as ex:

pending_cancellation_error = ex

if not self._aborting:

# Our parent task is being cancelled:

# async def wrapper():

# async with TaskGroup() as g:

# g.create_task(foo)

# "wrapper" is being cancelled while "foo" is

# still running.

propagate_cancellation_error = ex

self._abort()

self._on_completed_fut = None

assert not self._tasks

if self._base_error is not None:

try:

raise self._base_error

finally:

exc = None

if self._parent_cancel_requested:

# If this flag is set we *must* call uncancel().

if self._parent_task.uncancel() == 0:

# If there are no pending cancellations left,

# don't propagate CancelledError.

propagate_cancellation_error = None

elif propagate_cancellation_error is None:

# gh-155433: the remaining cancellation is not ours, don't drop it

propagate_cancellation_error = pending_cancellation_error

# Propagate CancelledError if there is one, except if there

# are other errors -- those have priority.

try:

if propagate_cancellation_error is not None and not self._errors:

try:

raise propagate_cancellation_error

finally:

exc = None

finally:

propagate_cancellation_error = None

if et is not None and not issubclass(et, exceptions.CancelledError):

self._errors.append(exc)

if self._errors:

# If the parent task is being cancelled from the outside

# of the taskgroup, un-cancel and re-cancel the parent task,

# which will keep the cancel count stable.

if self._parent_task.cancelling():

self._parent_task.uncancel()

self._parent_task.cancel()

try:

# If the *only* error is a GeneratorExit from the body

# of the group, then instead of raising an

# ExceptionGroup we raise GeneratorExit. This ensures

# that async generators that use TaskGroup properly

# swallow the exception on `aclose()` while ensuring

# that no exceptions from subtasks are swallowed.

if (

et is not None

and issubclass(et, GeneratorExit)

and len(self._errors) == 1

raise exc

else:

raise BaseExceptionGroup(

'unhandled errors in a TaskGroup',

self._errors,

) from None

finally:

exc = None

# Suppress any remaining exception (exceptions deserving to be raised

# were raised above).

return True

def create_task(self, coro, **kwargs):

"""Create a new task in this group and return it.

Similar to `asyncio.create_task`.

"""

if not self._entered:

coro.close()

raise RuntimeError(f"TaskGroup {self!r} has not been entered")

if self._exiting and not self._tasks:

coro.close()

raise RuntimeError(f"TaskGroup {self!r} is finished")

if self._aborting:

coro.close()

raise RuntimeError(f"TaskGroup {self!r} is shutting down")

task = self._loop.create_task(coro, **kwargs)

futures.future_add_to_awaited_by(task, self._parent_task)

# Always schedule the done callback even if the task is

# already done (e.g. if the coro was able to complete eagerly),

# otherwise if the task completes with an exception then it will cancel

# the current task too early. gh-128550, gh-128588

self._tasks.add(task)

task.add_done_callback(self._on_task_done)

# gh-155418: an eager task can cancel the group before joining _tasks

if self._aborting and not task.done():

task.cancel()

try:

return task

finally:

# gh-128552: prevent a refcycle of

# task.exception().__traceback__->TaskGroup.create_task->task

del task

# Since Python 3.8 Tasks propagate all exceptions correctly,

# except for KeyboardInterrupt and SystemExit which are

# still considered special.

def _is_base_error(self, exc: BaseException) -> bool:

assert isinstance(exc, BaseException)

return isinstance(exc, (SystemExit, KeyboardInterrupt))

def _abort(self):

self._aborting = True

for t in self._tasks:

if not t.done():

t.cancel()

def _on_task_done(self, task):

self._tasks.discard(task)

futures.future_discard_from_awaited_by(task, self._parent_task)

if self._on_completed_fut is not None and not self._tasks:

if not self._on_completed_fut.done():

self._on_completed_fut.set_result(True)

if task.cancelled():

return

exc = task.exception()

if exc is None:

return

self._errors.append(exc)

if self._is_base_error(exc) and self._base_error is None:

self._base_error = exc

if self._parent_task.done():

# Not sure if this case is possible, but we want to handle

# it anyways.

self._loop.call_exception_handler({

'message': f'Task {task!r} has errored out but its parent '

f'task {self._parent_task} is already completed',

'exception': exc,

'task': task,

})

return

if not self._aborting and not self._parent_cancel_requested:

# If parent task *is not* being cancelled, it means that we want

# to manually cancel it to abort whatever is being run right now

# in the TaskGroup. But we want to mark parent task as

# "not cancelled" later in __aexit__. Example situation that

# we need to handle:

# async def foo():

# try:

# async with TaskGroup() as g:

# g.create_task(crash_soon())

# await something # <- this needs to be canceled

# # by the TaskGroup, e.g.

# # foo() needs to be cancelled

# except Exception:

# # Ignore any exceptions raised in the TaskGroup

# pass

# await something_else # this line has to be called

# # after TaskGroup is finished.

self._abort()

self._parent_cancel_requested = True

self._parent_task.cancel()

def cancel(self):

"""Cancel the task group

`cancel()` will be called on any tasks in the group that aren't yet

done, as well as the parent (body) of the group. This will cause

the task group context manager to exit *without*

`asyncio.CancelledError` being raised.

If `cancel()` is called before entering the task group, the group

will be cancelled upon entry. This is useful for patterns where

one piece of code passes an unused TaskGroup instance to another in

order to have the ability to cancel anything run within the group.

`cancel()` is idempotent and may be called after the task group has

already exited.

"""

if not self._entered:

self._cancel_on_enter = True

return

if self._exiting and not self._tasks:

return

if not self._aborting:

self._abort()

if self._parent_task and not self._parent_cancel_requested:

self._parent_cancel_requested = True

self._parent_task.cancel()

cpython/Lib/asyncio/taskgroups.py at 3.15 · python/cpython · GitHub

Repository navigation

Expand file tree

Breadcrumbs

taskgroups.py

Latest commit

History

Breadcrumbs

taskgroups.py

File metadata and controls