async-kernel 0.2.0__tar.gz → 0.2.1__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (81) hide show
  1. {async_kernel-0.2.0 → async_kernel-0.2.1}/.github/workflows/pre-commit.yml +1 -1
  2. {async_kernel-0.2.0 → async_kernel-0.2.1}/CHANGELOG.md +33 -0
  3. {async_kernel-0.2.0 → async_kernel-0.2.1}/PKG-INFO +1 -1
  4. {async_kernel-0.2.0 → async_kernel-0.2.1}/_version.py +2 -2
  5. {async_kernel-0.2.0 → async_kernel-0.2.1}/src/async_kernel/caller.py +257 -86
  6. {async_kernel-0.2.0 → async_kernel-0.2.1}/src/async_kernel/kernel.py +7 -5
  7. {async_kernel-0.2.0 → async_kernel-0.2.1}/src/async_kernel/typing.py +3 -2
  8. {async_kernel-0.2.0 → async_kernel-0.2.1}/src/async_kernel/utils.py +6 -1
  9. {async_kernel-0.2.0 → async_kernel-0.2.1}/tests/test_caller.py +288 -88
  10. {async_kernel-0.2.0 → async_kernel-0.2.1}/tests/test_kernel.py +5 -6
  11. {async_kernel-0.2.0 → async_kernel-0.2.1}/.github/dependabot.yaml +0 -0
  12. {async_kernel-0.2.0 → async_kernel-0.2.1}/.github/release.yml +0 -0
  13. {async_kernel-0.2.0 → async_kernel-0.2.1}/.github/workflows/ci.yml +0 -0
  14. {async_kernel-0.2.0 → async_kernel-0.2.1}/.github/workflows/enforce-label.yml +0 -0
  15. {async_kernel-0.2.0 → async_kernel-0.2.1}/.github/workflows/new_release.yml +0 -0
  16. {async_kernel-0.2.0 → async_kernel-0.2.1}/.github/workflows/publish-docs.yml +0 -0
  17. {async_kernel-0.2.0 → async_kernel-0.2.1}/.github/workflows/publish-to-pypi.yml +0 -0
  18. {async_kernel-0.2.0 → async_kernel-0.2.1}/.gitignore +0 -0
  19. {async_kernel-0.2.0 → async_kernel-0.2.1}/.pre-commit-config.yaml +0 -0
  20. {async_kernel-0.2.0 → async_kernel-0.2.1}/.vscode/launch.json +0 -0
  21. {async_kernel-0.2.0 → async_kernel-0.2.1}/.vscode/settings.json +0 -0
  22. {async_kernel-0.2.0 → async_kernel-0.2.1}/.vscode/spellright.dict +0 -0
  23. {async_kernel-0.2.0 → async_kernel-0.2.1}/CONTRIBUTING.md +0 -0
  24. {async_kernel-0.2.0 → async_kernel-0.2.1}/IPYTHON_LICENSE +0 -0
  25. {async_kernel-0.2.0 → async_kernel-0.2.1}/LICENSE +0 -0
  26. {async_kernel-0.2.0 → async_kernel-0.2.1}/README.md +0 -0
  27. {async_kernel-0.2.0 → async_kernel-0.2.1}/cliff.toml +0 -0
  28. {async_kernel-0.2.0 → async_kernel-0.2.1}/docs/about/changelog.md +0 -0
  29. {async_kernel-0.2.0 → async_kernel-0.2.1}/docs/about/contributing.md +0 -0
  30. {async_kernel-0.2.0 → async_kernel-0.2.1}/docs/about/index.md +0 -0
  31. {async_kernel-0.2.0 → async_kernel-0.2.1}/docs/about/license.md +0 -0
  32. {async_kernel-0.2.0 → async_kernel-0.2.1}/docs/caller.ipynb +0 -0
  33. {async_kernel-0.2.0 → async_kernel-0.2.1}/docs/commands.md +0 -0
  34. {async_kernel-0.2.0 → async_kernel-0.2.1}/docs/index.md +0 -0
  35. {async_kernel-0.2.0 → async_kernel-0.2.1}/docs/javascripts/extra.js +0 -0
  36. {async_kernel-0.2.0 → async_kernel-0.2.1}/docs/notebooks/caller.ipynb +0 -0
  37. {async_kernel-0.2.0 → async_kernel-0.2.1}/docs/notebooks/concurrency.ipynb +0 -0
  38. {async_kernel-0.2.0 → async_kernel-0.2.1}/docs/notebooks/index.md +0 -0
  39. {async_kernel-0.2.0 → async_kernel-0.2.1}/docs/notebooks/simple_example.ipynb +0 -0
  40. {async_kernel-0.2.0 → async_kernel-0.2.1}/docs/overrides/main.html +0 -0
  41. {async_kernel-0.2.0 → async_kernel-0.2.1}/docs/reference/asyncshell.md +0 -0
  42. {async_kernel-0.2.0 → async_kernel-0.2.1}/docs/reference/caller.md +0 -0
  43. {async_kernel-0.2.0 → async_kernel-0.2.1}/docs/reference/comm.md +0 -0
  44. {async_kernel-0.2.0 → async_kernel-0.2.1}/docs/reference/command.md +0 -0
  45. {async_kernel-0.2.0 → async_kernel-0.2.1}/docs/reference/index.md +0 -0
  46. {async_kernel-0.2.0 → async_kernel-0.2.1}/docs/reference/kernel.md +0 -0
  47. {async_kernel-0.2.0 → async_kernel-0.2.1}/docs/reference/kernelspec.md +0 -0
  48. {async_kernel-0.2.0 → async_kernel-0.2.1}/docs/reference/typing.md +0 -0
  49. {async_kernel-0.2.0 → async_kernel-0.2.1}/docs/reference/utils.md +0 -0
  50. {async_kernel-0.2.0 → async_kernel-0.2.1}/docs/stylesheets/extra.css +0 -0
  51. {async_kernel-0.2.0 → async_kernel-0.2.1}/docs/usage/index.md +0 -0
  52. {async_kernel-0.2.0 → async_kernel-0.2.1}/hatch_build.py +0 -0
  53. {async_kernel-0.2.0 → async_kernel-0.2.1}/mkdocs.yml +0 -0
  54. {async_kernel-0.2.0 → async_kernel-0.2.1}/pyproject.toml +0 -0
  55. {async_kernel-0.2.0 → async_kernel-0.2.1}/src/async_kernel/__init__.py +0 -0
  56. {async_kernel-0.2.0 → async_kernel-0.2.1}/src/async_kernel/__main__.py +0 -0
  57. {async_kernel-0.2.0 → async_kernel-0.2.1}/src/async_kernel/asyncshell.py +0 -0
  58. {async_kernel-0.2.0 → async_kernel-0.2.1}/src/async_kernel/comm.py +0 -0
  59. {async_kernel-0.2.0 → async_kernel-0.2.1}/src/async_kernel/command.py +0 -0
  60. {async_kernel-0.2.0 → async_kernel-0.2.1}/src/async_kernel/compiler.py +0 -0
  61. {async_kernel-0.2.0 → async_kernel-0.2.1}/src/async_kernel/debugger.py +0 -0
  62. {async_kernel-0.2.0 → async_kernel-0.2.1}/src/async_kernel/iostream.py +0 -0
  63. {async_kernel-0.2.0 → async_kernel-0.2.1}/src/async_kernel/kernelspec.py +0 -0
  64. {async_kernel-0.2.0 → async_kernel-0.2.1}/src/async_kernel/resources/logo-32x32.png +0 -0
  65. {async_kernel-0.2.0 → async_kernel-0.2.1}/src/async_kernel/resources/logo-64x64.png +0 -0
  66. {async_kernel-0.2.0 → async_kernel-0.2.1}/src/async_kernel/resources/logo-svg.svg +0 -0
  67. {async_kernel-0.2.0 → async_kernel-0.2.1}/tests/__init__.py +0 -0
  68. {async_kernel-0.2.0 → async_kernel-0.2.1}/tests/conftest.py +0 -0
  69. {async_kernel-0.2.0 → async_kernel-0.2.1}/tests/references.py +0 -0
  70. {async_kernel-0.2.0 → async_kernel-0.2.1}/tests/test_comm.py +0 -0
  71. {async_kernel-0.2.0 → async_kernel-0.2.1}/tests/test_command.py +0 -0
  72. {async_kernel-0.2.0 → async_kernel-0.2.1}/tests/test_debugger.py +0 -0
  73. {async_kernel-0.2.0 → async_kernel-0.2.1}/tests/test_enter_kernel.py +0 -0
  74. {async_kernel-0.2.0 → async_kernel-0.2.1}/tests/test_iostream.py +0 -0
  75. {async_kernel-0.2.0 → async_kernel-0.2.1}/tests/test_kernel_subclass.py +0 -0
  76. {async_kernel-0.2.0 → async_kernel-0.2.1}/tests/test_kernelspec.py +0 -0
  77. {async_kernel-0.2.0 → async_kernel-0.2.1}/tests/test_message_spec.py +0 -0
  78. {async_kernel-0.2.0 → async_kernel-0.2.1}/tests/test_typing.py +0 -0
  79. {async_kernel-0.2.0 → async_kernel-0.2.1}/tests/test_utils.py +0 -0
  80. {async_kernel-0.2.0 → async_kernel-0.2.1}/tests/utils.py +0 -0
  81. {async_kernel-0.2.0 → async_kernel-0.2.1}/uv.lock +0 -0
@@ -8,7 +8,7 @@ jobs:
8
8
  runs-on: ubuntu-latest
9
9
  steps:
10
10
  - uses: actions/checkout@v5
11
- - uses: actions/setup-python@v5
11
+ - uses: actions/setup-python@v6
12
12
  with:
13
13
  python-version: 3.x
14
14
  - uses: pre-commit/action@v3.0.1
@@ -5,6 +5,36 @@ All notable changes to this project will be documented in this file.
5
5
  The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
6
6
  and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
7
7
 
8
+ ## [0.2.1] - 2025-09-10
9
+
10
+ ### <!-- 0 --> 🏗️ Breaking changes
11
+
12
+ - Maintenance [#105](https://github.com/fleming79/async-kernel/pull/105)
13
+
14
+ ### <!-- 1 --> 🚀 Features
15
+
16
+ - Divide Lock into AsyncLock and ReentrantAsyncLock [#113](https://github.com/fleming79/async-kernel/pull/113)
17
+
18
+ - Improve Lock class [#112](https://github.com/fleming79/async-kernel/pull/112)
19
+
20
+ - Add a context based Lock [#111](https://github.com/fleming79/async-kernel/pull/111)
21
+
22
+ - Add classmethod Caller.wait [#106](https://github.com/fleming79/async-kernel/pull/106)
23
+
24
+ - Add 'shield' option to Caller.as_completed. [#104](https://github.com/fleming79/async-kernel/pull/104)
25
+
26
+ ### <!-- 6 --> 🌀 Miscellaneous
27
+
28
+ - Bump actions/setup-python from 5 to 6 in the actions group [#110](https://github.com/fleming79/async-kernel/pull/110)
29
+
30
+ - Maintenance - Caller refactoring [#109](https://github.com/fleming79/async-kernel/pull/109)
31
+
32
+ - Drop WaitType for Literals directly in Caller.wait. [#108](https://github.com/fleming79/async-kernel/pull/108)
33
+
34
+ - Change Caller._queue_map to a WeakKeyDictionary. [#107](https://github.com/fleming79/async-kernel/pull/107)
35
+
36
+ - Refactor Caller.wait to avoid catching exceptions. [#103](https://github.com/fleming79/async-kernel/pull/103)
37
+
8
38
  ## [0.2.0] - 2025-09-06
9
39
 
10
40
  ### <!-- 0 --> 🏗️ Breaking changes
@@ -29,6 +59,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
29
59
 
30
60
  ### <!-- 6 --> 🌀 Miscellaneous
31
61
 
62
+ - Prepare for release v0.2.0 [#102](https://github.com/fleming79/async-kernel/pull/102)
63
+
32
64
  - Result should raise cancelled error, but was raising and InvalidStateError. [#98](https://github.com/fleming79/async-kernel/pull/98)
33
65
 
34
66
  ## [0.1.4] - 2025-09-03
@@ -207,6 +239,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
207
239
 
208
240
  - Bump the actions group across 1 directory with 2 updates [#3](https://github.com/fleming79/async-kernel/pull/3)
209
241
 
242
+ [0.2.1]: https://github.com/fleming79/async-kernel/compare/v0.2.0..v0.2.1
210
243
  [0.2.0]: https://github.com/fleming79/async-kernel/compare/v0.1.4..v0.2.0
211
244
  [0.1.4]: https://github.com/fleming79/async-kernel/compare/v0.1.3..v0.1.4
212
245
  [0.1.3]: https://github.com/fleming79/async-kernel/compare/v0.1.2..v0.1.3
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: async-kernel
3
- Version: 0.2.0
3
+ Version: 0.2.1
4
4
  Summary: A concurrent python kernel for Jupyter supporting AnyIO, AsyncIO and Trio.
5
5
  Project-URL: Homepage, https://fleming79.github.io/async-kernel
6
6
  Project-URL: Documentation, https://fleming79.github.io/async-kernel
@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
28
28
  commit_id: COMMIT_ID
29
29
  __commit_id__: COMMIT_ID
30
30
 
31
- __version__ = version = '0.2.0'
32
- __version_tuple__ = version_tuple = (0, 2, 0)
31
+ __version__ = version = '0.2.1'
32
+ __version_tuple__ = version_tuple = (0, 2, 1)
33
33
 
34
34
  __commit_id__ = commit_id = None
@@ -14,6 +14,7 @@ from typing import TYPE_CHECKING, Any, ClassVar, Literal, Self, cast, overload
14
14
 
15
15
  import anyio
16
16
  import sniffio
17
+ from anyio.streams.memory import MemoryObjectSendStream
17
18
  from typing_extensions import override
18
19
  from zmq import Context, Socket, SocketType
19
20
 
@@ -32,7 +33,7 @@ if TYPE_CHECKING:
32
33
 
33
34
  from async_kernel.typing import P
34
35
 
35
- __all__ = ["Caller", "Future", "FutureCancelledError", "InvalidStateError"]
36
+ __all__ = ["AsyncLock", "Caller", "Future", "FutureCancelledError", "InvalidStateError", "ReentrantAsyncLock"]
36
37
 
37
38
 
38
39
  class FutureCancelledError(anyio.ClosedResourceError):
@@ -132,23 +133,22 @@ class Future(Awaitable[T]):
132
133
 
133
134
  Args:
134
135
  timeout: Timeout in seconds.
135
- shield: Shield cancellation.
136
+ shield: Shield the future from cancellation.
136
137
  result: Whether the result should be returned.
137
138
  """
138
- if not self._event_done.is_set():
139
- with anyio.fail_after(delay=timeout):
140
- try:
139
+ try:
140
+ if not self.done():
141
+ with anyio.fail_after(timeout):
141
142
  if threading.current_thread() is self.thread:
142
143
  if not self._anyio_event_done:
143
144
  self._anyio_event_done = anyio.Event()
144
145
  await self._anyio_event_done.wait()
145
146
  else:
146
147
  await wait_thread_event(self._event_done)
147
- except BaseException as e:
148
- if not shield:
149
- self.cancel(str(e))
150
- raise
151
- return self.result() if result else None
148
+ return self.result() if result else None
149
+ finally:
150
+ if not self.done() and not shield:
151
+ self.cancel("Cancelled with waiter cancellation.")
152
152
 
153
153
  if TYPE_CHECKING:
154
154
 
@@ -306,10 +306,11 @@ class Caller:
306
306
  _outstanding = 0
307
307
  _to_thread_pool: ClassVar[deque[Self]] = deque()
308
308
  _pool_instances: ClassVar[weakref.WeakSet[Self]] = weakref.WeakSet()
309
- _executor_queue: dict
309
+ _queue_map: weakref.WeakKeyDictionary[Callable[..., Awaitable[Any]], MemoryObjectSendStream[tuple]]
310
310
  _taskgroup: TaskGroup | None = None
311
311
  _callers: deque[tuple[contextvars.Context, tuple[Future, float, float, Callable, tuple, dict]] | Callable[[], Any]]
312
312
  _callers_added: threading.Event
313
+ _stopped_event: threading.Event
313
314
  _stopped = False
314
315
  _protected = False
315
316
  _running = False
@@ -367,7 +368,7 @@ class Caller:
367
368
  inst._callers = deque()
368
369
  inst._callers_added = threading.Event()
369
370
  inst._protected = protected
370
- inst._executor_queue = {}
371
+ inst._queue_map = weakref.WeakKeyDictionary()
371
372
  cls._instances[thread] = inst
372
373
  return inst
373
374
 
@@ -376,9 +377,9 @@ class Caller:
376
377
  return f"Caller<{self.thread.name}>"
377
378
 
378
379
  async def __aenter__(self) -> Self:
379
- self._cancelled_exception_class = anyio.get_cancelled_exc_class()
380
380
  async with contextlib.AsyncExitStack() as stack:
381
381
  self._running = True
382
+ self._stopped_event = threading.Event()
382
383
  self._taskgroup = tg = await stack.enter_async_context(anyio.create_task_group())
383
384
  await tg.start(self._server_loop, tg)
384
385
  self.__stack = stack.pop_all()
@@ -386,7 +387,7 @@ class Caller:
386
387
 
387
388
  async def __aexit__(self, exc_type, exc_value, exc_tb) -> None:
388
389
  if self.__stack is not None:
389
- self.stop()
390
+ self.stop(force=True)
390
391
  await self.__stack.__aexit__(exc_type, exc_value, exc_tb)
391
392
 
392
393
  async def _server_loop(self, tg: TaskGroup, task_status: TaskStatus[None]) -> None:
@@ -397,7 +398,12 @@ class Caller:
397
398
  self.iopub_sockets[self.thread] = socket
398
399
  task_status.started()
399
400
  while not self._stopped:
400
- while len(self._callers):
401
+ if not self._callers:
402
+ self._callers_added.clear()
403
+ await wait_thread_event(self._callers_added)
404
+ while self._callers:
405
+ if self._stopped:
406
+ return
401
407
  job = self._callers.popleft()
402
408
  if isinstance(job, Callable):
403
409
  try:
@@ -407,8 +413,6 @@ class Caller:
407
413
  else:
408
414
  context, args = job
409
415
  context.run(tg.start_soon, self._wrap_call, *args)
410
- self._callers_added.clear()
411
- await wait_thread_event(self._callers_added)
412
416
  finally:
413
417
  self._running = False
414
418
  for job in self._callers:
@@ -416,6 +420,7 @@ class Caller:
416
420
  job[1][0].set_exception(FutureCancelledError())
417
421
  socket.close()
418
422
  self.iopub_sockets.pop(self.thread, None)
423
+ self._stopped_event.set()
419
424
  tg.cancel_scope.cancel()
420
425
 
421
426
  async def _wrap_call(
@@ -442,12 +447,9 @@ class Caller:
442
447
  result: T = await result
443
448
  if fut.cancelled() and not scope.cancel_called:
444
449
  scope.cancel()
445
- if scope.cancel_called:
446
- # await here to allow the cancel scope to be raised/caught.
447
- await anyio.sleep(0)
448
450
  self._outstanding -= 1 # update first for _to_thread_on_done
449
451
  fut.set_result(result)
450
- except self._cancelled_exception_class:
452
+ except anyio.get_cancelled_exc_class():
451
453
  fut.cancel()
452
454
  self._outstanding -= 1 # update first for _to_thread_on_done
453
455
  fut.set_result(cast("T", None)) # This will cancel
@@ -493,13 +495,18 @@ class Caller:
493
495
  if self._protected and not force:
494
496
  return
495
497
  self._stopped = True
498
+ for sender in self._queue_map.values():
499
+ sender.close()
500
+ self._queue_map.clear()
496
501
  self._callers_added.set()
497
502
  self._instances.pop(self.thread, None)
498
503
  if self in self._to_thread_pool:
499
504
  self._to_thread_pool.remove(self)
505
+ if self.thread is not threading.current_thread():
506
+ self._stopped_event.wait()
500
507
 
501
508
  def call_later(
502
- self, func: Callable[P, T | Awaitable[T]], delay: float = 0.0, /, *args: P.args, **kwargs: P.kwargs
509
+ self, delay: float, func: Callable[P, T | Awaitable[T]], /, *args: P.args, **kwargs: P.kwargs
503
510
  ) -> Future[T]:
504
511
  """
505
512
  Schedule func to be called in caller's event loop copying the current context.
@@ -530,7 +537,7 @@ class Caller:
530
537
  *args: Arguments to use with func.
531
538
  **kwargs: Keyword arguments to use with func.
532
539
  """
533
- return self.call_later(func, 0.0, *args, **kwargs)
540
+ return self.call_later(0, func, *args, **kwargs)
534
541
 
535
542
  def call_direct(self, func: Callable[P, Any], /, *args: P.args, **kwargs: P.kwargs) -> None:
536
543
  """
@@ -550,9 +557,9 @@ class Caller:
550
557
  self._callers.append(functools.partial(func, *args, **kwargs))
551
558
  self._callers_added.set()
552
559
 
553
- def has_execution_queue(self, func: Callable) -> bool:
560
+ def queue_exists(self, func: Callable) -> bool:
554
561
  "Returns True if an execution queue exists for `func`."
555
- return func in self._executor_queue
562
+ return func in self._queue_map
556
563
 
557
564
  if TYPE_CHECKING:
558
565
 
@@ -563,7 +570,7 @@ class Caller:
563
570
  /,
564
571
  *args: *PosArgsT,
565
572
  max_buffer_size: NoValue | int = NoValue, # pyright: ignore[reportInvalidTypeForm]
566
- send_nowait: Literal[False],
573
+ wait: Literal[True],
567
574
  ) -> CoroutineType[Any, Any, None]: ...
568
575
  @overload
569
576
  def queue_call(
@@ -572,7 +579,7 @@ class Caller:
572
579
  /,
573
580
  *args: *PosArgsT,
574
581
  max_buffer_size: NoValue | int = NoValue, # pyright: ignore[reportInvalidTypeForm]
575
- send_nowait: Literal[True] | Any = True,
582
+ wait: Literal[False] | Any = False,
576
583
  ) -> None: ...
577
584
 
578
585
  def queue_call(
@@ -581,10 +588,10 @@ class Caller:
581
588
  /,
582
589
  *args: *PosArgsT,
583
590
  max_buffer_size: NoValue | int = NoValue, # pyright: ignore[reportInvalidTypeForm]
584
- send_nowait: bool = True,
591
+ wait: bool = False,
585
592
  ) -> CoroutineType[Any, Any, None] | None:
586
593
  """
587
- Queue the execution of func in queue specific to the function (not thread-safe).
594
+ Queue the execution of `func` with the arguments `*args` in a queue unique to it (not thread-safe).
588
595
 
589
596
  The args are added to a queue associated with the provided `func`. If queue does not already exist for
590
597
  func, a new queue is created with a specified maximum buffer size. The arguments are then sent to the queue,
@@ -595,11 +602,19 @@ class Caller:
595
602
  func: The asynchronous function to execute.
596
603
  *args: The arguments to pass to the function.
597
604
  max_buffer_size: The maximum buffer size for the queue. If NoValue, defaults to [async_kernel.Caller.MAX_BUFFER_SIZE].
598
- send_nowait: Set as False to return a coroutine that is used to send the request.
605
+ wait: Set as True to return a coroutine that will return once the request is sent.
599
606
  Use this to prevent experiencing exceptions if the buffer is full.
607
+
608
+ !!! info
609
+
610
+ The queue will stay open until one of the following occurs.
611
+
612
+ 1. It explicitly closed with the method `queue_close`.
613
+ 1. All strong references are lost the function/method.
614
+
600
615
  """
601
616
  self._check_in_thread()
602
- if not self.has_execution_queue(func):
617
+ if not (sender := self._queue_map.get(func)):
603
618
  max_buffer_size = self.MAX_BUFFER_SIZE if max_buffer_size is NoValue else max_buffer_size
604
619
  sender, queue = anyio.create_memory_object_stream[tuple[*PosArgsT]](max_buffer_size=max_buffer_size)
605
620
 
@@ -608,38 +623,28 @@ class Caller:
608
623
  with contextlib.suppress(anyio.get_cancelled_exc_class()):
609
624
  async with queue as receive_stream:
610
625
  async for args in receive_stream:
626
+ if func not in self._queue_map:
627
+ break
611
628
  try:
612
629
  await func(*args)
613
630
  except Exception as e:
614
631
  self.log.exception("Execution %f failed", func, exc_info=e)
615
632
  finally:
616
- self._executor_queue.pop(execute_loop, None)
633
+ self._queue_map.pop(func, None)
617
634
 
618
- self._executor_queue[func] = {"queue": sender, "future": self.call_soon(execute_loop)}
619
- sender: MemoryObjectSendStream[tuple[*PosArgsT]] = self._executor_queue[func]["queue"]
620
- return sender.send_nowait(args) if send_nowait else sender.send(args)
635
+ self._queue_map[func] = sender
636
+ self.call_soon(execute_loop)
637
+ return sender.send(args) if wait else sender.send_nowait(args)
621
638
 
622
- async def queue_close(self, func: Callable, *, force: bool = False) -> bool:
639
+ def queue_close(self, func: Callable) -> None:
623
640
  """
624
- Close the execution queue associated with func (not thread-safe).
641
+ Close the execution queue associated with func (thread-safe).
625
642
 
626
643
  Args:
627
644
  func: The queue of the function to close.
628
- force: Shutdown without waiting pending tasks to complete.
629
-
630
- Returns:
631
- True if a queue was closed.
632
645
  """
633
- self._check_in_thread()
634
- if queue_map := self._executor_queue.pop(func, None):
635
- if force:
636
- queue_map["future"].cancel()
637
- else:
638
- await queue_map["queue"].aclose()
639
- with contextlib.suppress(FutureCancelledError):
640
- await queue_map["future"]
641
- return True
642
- return False
646
+ if sender := self._queue_map.pop(func, None):
647
+ self.call_direct(sender.close)
643
648
 
644
649
  @classmethod
645
650
  def stop_all(cls, *, _stop_protected: bool = False) -> None:
@@ -756,12 +761,23 @@ class Caller:
756
761
  """A classmethod that returns the current future when called from inside a function scheduled by Caller."""
757
762
  return cls._future_var.get()
758
763
 
764
+ @classmethod
765
+ def all_callers(cls, running_only: bool = True) -> list[Caller]:
766
+ """
767
+ A classmethod to get a list of the callers.
768
+
769
+ Args:
770
+ running_only: Restrict the list to callers that are active (running in an async context).
771
+ """
772
+ return [caller for caller in Caller._instances.values() if caller._running or not running_only]
773
+
759
774
  @classmethod
760
775
  async def as_completed(
761
776
  cls,
762
777
  items: Iterable[Future[T]] | AsyncGenerator[Future[T]],
763
778
  *,
764
779
  max_concurrent: NoValue | int = NoValue, # pyright: ignore[reportInvalidTypeForm]
780
+ shield: bool = False,
765
781
  ) -> AsyncGenerator[Future[T], Any]:
766
782
  """
767
783
  A classmethod iterator to get [Futures][async_kernel.caller.Future] as they complete.
@@ -771,6 +787,7 @@ class Caller:
771
787
  max_concurrent: The maximum number of concurrent futures to monitor at a time.
772
788
  This is useful when `items` is a generator utilising Caller.to_thread.
773
789
  By default this will limit to `Caller.MAX_IDLE_POOL_INSTANCES`.
790
+ shield: Shield existing items from cancellation.
774
791
 
775
792
  !!! tip
776
793
 
@@ -782,62 +799,216 @@ class Caller:
782
799
  futures: set[Future[T]] = set()
783
800
  done = False
784
801
  resume: Event | None = cast("anyio.Event | None", None)
802
+ current_future = cls.current_future()
803
+ if isinstance(items, set | list | tuple):
804
+ max_concurrent_ = 0
805
+ else:
806
+ max_concurrent_ = cls.MAX_IDLE_POOL_INSTANCES if max_concurrent is NoValue else int(max_concurrent)
785
807
 
786
808
  def _on_done(fut: Future[T]) -> None:
787
809
  has_result.append(fut)
788
810
  event_future_ready.set()
789
811
 
790
- async def iter_items(task_status: TaskStatus[None]):
812
+ async def iter_items():
791
813
  nonlocal done, resume
792
- if isinstance(items, set | list | tuple):
793
- max_concurrent_ = 0
794
- else:
795
- max_concurrent_ = cls.MAX_IDLE_POOL_INSTANCES if max_concurrent is NoValue else int(max_concurrent)
796
-
797
814
  gen = items if isinstance(items, AsyncGenerator) else iter(items)
798
- task_status.started()
799
815
  try:
800
816
  while True:
801
817
  fut = await anext(gen) if isinstance(gen, AsyncGenerator) else next(gen)
802
- futures.add(fut)
803
- if fut.done():
804
- has_result.append(fut)
805
- event_future_ready.set()
806
- else:
807
- fut.add_done_callback(_on_done)
808
- if max_concurrent_ and len(futures) == max_concurrent_:
809
- resume = anyio.Event()
810
- await resume.wait()
818
+ if fut is not current_future:
819
+ futures.add(fut)
820
+ if fut.done():
821
+ has_result.append(fut)
822
+ event_future_ready.set()
823
+ else:
824
+ fut.add_done_callback(_on_done)
825
+ if max_concurrent_ and len(futures) == max_concurrent_:
826
+ resume = anyio.Event()
827
+ await resume.wait()
811
828
  except (StopAsyncIteration, StopIteration):
812
829
  return
813
830
  finally:
814
831
  done = True
815
832
  event_future_ready.set()
816
833
 
834
+ fut = cls().call_soon(iter_items)
817
835
  try:
818
- async with anyio.create_task_group() as tg:
819
- await tg.start(iter_items)
820
- while futures or not done:
821
- if has_result:
822
- event_future_ready.clear()
823
- fut = has_result.popleft()
824
- futures.discard(fut)
825
- yield fut
826
- if resume:
827
- resume.set()
828
- continue
829
- if not has_result:
830
- await wait_thread_event(event_future_ready)
836
+ while futures or not done:
837
+ if has_result:
838
+ event_future_ready.clear()
839
+ fut = has_result.popleft()
840
+ futures.discard(fut)
841
+ yield fut
842
+ if resume:
843
+ resume.set()
844
+ else:
845
+ await wait_thread_event(event_future_ready)
831
846
  finally:
847
+ fut.cancel()
832
848
  for fut in futures:
833
- fut.cancel()
849
+ fut.remove_done_callback(_on_done)
850
+ if not shield:
851
+ fut.cancel("Cancelled by as_completed")
834
852
 
835
853
  @classmethod
836
- def all_callers(cls, running_only: bool = True) -> list[Caller]:
854
+ async def wait(
855
+ cls,
856
+ items: Iterable[Future[T]],
857
+ *,
858
+ timeout: float | None = None,
859
+ return_when: Literal["FIRST_COMPLETED", "FIRST_EXCEPTION", "ALL_COMPLETED"] = "ALL_COMPLETED",
860
+ ) -> tuple[set[T], set[Future[T]]]:
837
861
  """
838
- A classmethod to get a list of the callers.
862
+ A classmethod to wait for the futures given by items to complete.
839
863
 
840
- Args:
841
- running_only: Restrict the list to callers that are active (running in an async context).
864
+ Returns two sets of the futures: (done, pending).
865
+
866
+ !!! example
867
+
868
+ ```python
869
+ done, pending = await asyncio.wait(items)
870
+ ```
871
+
872
+ !!! info
873
+
874
+ - This does not raise a TimeoutError!
875
+ - Futures that aren't done when the timeout occurs are returned in the second set.
842
876
  """
843
- return [caller for caller in Caller._instances.values() if caller._running or not running_only]
877
+ done = set()
878
+ if pending := set(items):
879
+ with anyio.move_on_after(timeout):
880
+ async for fut in cls.as_completed(items, shield=True):
881
+ pending.discard(fut)
882
+ done.add(fut)
883
+ if return_when == "FIRST_COMPLETED":
884
+ break
885
+ if return_when == "FIRST_EXCEPTION" and (fut.cancelled() or fut.exception()):
886
+ break
887
+ return done, pending
888
+
889
+
890
+ class AsyncLock:
891
+ """
892
+ Implements a mutex asynchronous lock that is compatible with [async_kernel.caller.Caller][].
893
+
894
+ !!! note
895
+
896
+ - Attempting to lock a 'mutuex' configured lock that is *locked* will raise a [RuntimeError][].
897
+ """
898
+
899
+ _reentrant: ClassVar[bool] = False
900
+ _count: int = 0
901
+ _ctx_count: int = 0
902
+ _ctx_current: int = 0
903
+ _releasing: bool = False
904
+
905
+ def __init__(self):
906
+ self._ctx_var: contextvars.ContextVar[int] = contextvars.ContextVar(f"Lock:{id(self)}", default=0)
907
+ self._queue: deque[tuple[int, Future[Future | None]]] = deque()
908
+
909
+ @override
910
+ def __repr__(self) -> str:
911
+ info = f"🔒{self.count}" if self.count else "🔓"
912
+ return f"{self.__class__.__name__}({info})"
913
+
914
+ async def __aenter__(self) -> Self:
915
+ return await self.acquire()
916
+
917
+ async def __aexit__(self, exc_type, exc, tb) -> None:
918
+ await self.release()
919
+
920
+ @property
921
+ def count(self) -> int:
922
+ "Returns the number of times the locked context has been entered."
923
+ return self._count
924
+
925
+ async def acquire(self) -> Self:
926
+ """
927
+ Acquire a lock.
928
+
929
+ If the lock is reentrant the internal counter increments to share the lock.
930
+ """
931
+ if not self._reentrant and self.is_in_context():
932
+ msg = "Already locked and not reentrant!"
933
+ raise RuntimeError(msg)
934
+ # Get the context.
935
+ if not self._reentrant or not (ctx := self._ctx_var.get()):
936
+ self._ctx_count = ctx = self._ctx_count + 1
937
+ self._ctx_var.set(ctx)
938
+ # Check if we can lock or re-enter an active lock.
939
+ if (not self._releasing) and ((not self.count) or (self._reentrant and self.is_in_context())):
940
+ self._count += 1
941
+ self._ctx_current = ctx
942
+ return self
943
+ # Join the queue.
944
+ k: tuple[int, Future[None | Future[Future[None] | None]]] = ctx, Future()
945
+ self._queue.append(k)
946
+ try:
947
+ fut = await k[1]
948
+ finally:
949
+ if k in self._queue:
950
+ self._queue.remove(k)
951
+ if fut:
952
+ self._ctx_current = ctx
953
+ fut.set_result(None)
954
+ if self._reentrant:
955
+ for k in tuple(self._queue):
956
+ if k[0] == ctx:
957
+ self._queue.remove(k)
958
+ k[1].set_result(None)
959
+ self._count += 1
960
+ self._releasing = False
961
+ return self
962
+
963
+ async def release(self) -> None:
964
+ """
965
+ Decrement the internal counter.
966
+
967
+ If the current depth==1 the lock will be passed to the next queued or released if there isn't one.
968
+ """
969
+ if not self.is_in_context():
970
+ raise InvalidStateError
971
+ if self._count == 1 and self._queue and not self._releasing:
972
+ self._releasing = True
973
+ self._ctx_var.set(0)
974
+ try:
975
+ fut = Future()
976
+ k = self._queue.popleft()
977
+ k[1].set_result(fut)
978
+ await k[1]
979
+ except Exception:
980
+ self._releasing = False
981
+ else:
982
+ self._count -= 1
983
+ if self._count == 0:
984
+ self._ctx_current = 0
985
+
986
+ def is_in_context(self) -> bool:
987
+ "Returns `True` if the current context has the lock."
988
+ return bool(self._count and self._ctx_current and (self._ctx_var.get() == self._ctx_current))
989
+
990
+
991
+ class ReentrantAsyncLock(AsyncLock):
992
+ """
993
+ Implements a Reentrant asynchronous lock compatible with [async_kernel.caller.Caller][].
994
+
995
+
996
+ !!! example
997
+
998
+ ```python
999
+ # Inside a coroutine running inside a thread where a [asyncio.caller.Caller][] instance is running.
1000
+
1001
+ lock = ReentrantAsyncLock(reentrant=True) # a reentrant lock
1002
+ async with lock:
1003
+ async with lock:
1004
+ Caller().to_thread(...) # The lock is shared with the thread.
1005
+ ```
1006
+
1007
+ !!! note
1008
+
1009
+ - The lock context can be exitied in any order.
1010
+ - A 'reentrant' lock can *release* control to another context and then re-enter later for
1011
+ tasks or threads called from a locked thread maintaining the same reentrant context.
1012
+ """
1013
+
1014
+ _reentrant: ClassVar[bool] = True