diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index 62dc955..00a1296 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -20,7 +20,7 @@ jobs: fail-fast: false matrix: os: [ubuntu-latest, macos-latest, windows-latest] - python-version: [ '3.10', '3.10', '3.11', '3.12', '3.13', '3.14' ] + python-version: [ '3.10', '3.11', '3.12', '3.13', '3.14' ] steps: - name: Checkout @@ -33,7 +33,7 @@ jobs: - name: Install akernel run: | - pip install . --group test + pip install . --group dev - name: Check style and types run: | diff --git a/MANIFEST.in b/MANIFEST.in deleted file mode 100644 index 1aba38f..0000000 --- a/MANIFEST.in +++ /dev/null @@ -1 +0,0 @@ -include LICENSE diff --git a/README.md b/README.md index d264f9e..21944d8 100644 --- a/README.md +++ b/README.md @@ -65,9 +65,11 @@ cell with the following code: ```python # cell 1 +from anyio import sleep + for i in range(10): print("cell 1:", i) - await asyncio.sleep(1) + await sleep(1) ``` Since this cell is `async` (it has an `await`), it will not block the execution of other cells. @@ -77,11 +79,11 @@ So you can run another cell concurrently, provided that this cell is also cooper # cell 2 for j in range(10): print("cell 2:", j) - await asyncio.sleep(1) + await sleep(1) ``` If cell 2 was blocking, cell 1 would pause until cell 2 was finished. You can see that by changing -`await asyncio.sleep(1)` into `time.sleep(1)` in cell 2. +`await sleep(1)` into `time.sleep(1)` in cell 2. You can make a cell wait for the previous one to be finished with: diff --git a/examples/concurrency.ipynb b/examples/concurrency.ipynb index 33cd4c7..54884a2 100644 --- a/examples/concurrency.ipynb +++ b/examples/concurrency.ipynb @@ -15,6 +15,8 @@ "metadata": {}, "outputs": [], "source": [ + "from anyio import sleep\n", + "\n", "__unchain_execution__()" ] }, @@ -26,7 +28,7 @@ "outputs": [], "source": [ "for i in range(5):\n", - " await asyncio.sleep(1)\n", + " await sleep(1)\n", " print(f\"{i=}\")" ] }, @@ -38,7 +40,7 @@ "outputs": [], "source": [ "for j in range(5):\n", - " await asyncio.sleep(1)\n", + " await sleep(1)\n", " print(f\"{j=}\")" ] }, @@ -50,7 +52,7 @@ "outputs": [], "source": [ "for k in range(5):\n", - " await asyncio.sleep(1)\n", + " await sleep(1)\n", " print(f\"{k=}\")" ] }, @@ -80,7 +82,7 @@ "outputs": [], "source": [ "for i in range(5):\n", - " await asyncio.sleep(1)\n", + " await sleep(1)\n", " print(f\"{i=}\")" ] }, @@ -92,7 +94,7 @@ "outputs": [], "source": [ "for j in range(5):\n", - " await asyncio.sleep(1)\n", + " await sleep(1)\n", " print(f\"{j=}\")" ] }, @@ -104,7 +106,7 @@ "outputs": [], "source": [ "for k in range(5):\n", - " await asyncio.sleep(1)\n", + " await sleep(1)\n", " print(f\"{k=}\")" ] } diff --git a/plugins/akernel_task/fps_akernel_task/akernel_task.py b/plugins/akernel_task/fps_akernel_task/akernel_task.py index 4e88019..a0349dd 100644 --- a/plugins/akernel_task/fps_akernel_task/akernel_task.py +++ b/plugins/akernel_task/fps_akernel_task/akernel_task.py @@ -8,8 +8,9 @@ class AKernelTask(_Kernel): - def __init__(self, *args, **kwargs): + def __init__(self, *args, execute_in_thread: bool = False, **kwargs): super().__init__() + self.execute_in_thread = execute_in_thread async def start(self, *, task_status: TaskStatus[None] = TASK_STATUS_IGNORED) -> None: async with ( @@ -37,6 +38,7 @@ async def start(self, *, task_status: TaskStatus[None] = TASK_STATUS_IGNORED) -> self._to_stdin_receive_stream, self._from_stdin_send_stream, self._from_iopub_send_stream, + execute_in_thread=self.execute_in_thread, ) self.task_group.start_soon(self.kernel.start) task_status.started() diff --git a/plugins/akernel_task/fps_akernel_task/main.py b/plugins/akernel_task/fps_akernel_task/main.py index 7f99c85..bb7aa73 100644 --- a/plugins/akernel_task/fps_akernel_task/main.py +++ b/plugins/akernel_task/fps_akernel_task/main.py @@ -1,5 +1,7 @@ from __future__ import annotations +from functools import partial + from fps import Module from jupyverse_api.kernel import KernelFactory @@ -9,6 +11,20 @@ class AKernelTaskModule(Module): + def __init__(self, *args, execute_in_thread: bool = False, **kwargs) -> None: + super().__init__(*args, **kwargs) + self.execute_in_thread = execute_in_thread + + async def prepare(self) -> None: + kernels = await self.get(Kernels) + kernels.register_kernel_factory( + "akernel", KernelFactory(partial(AKernelTask, execute_in_thread=self.execute_in_thread)) + ) + + +class AKernelThreadTaskModule(Module): async def prepare(self) -> None: kernels = await self.get(Kernels) - kernels.register_kernel_factory("akernel", KernelFactory(AKernelTask)) + kernels.register_kernel_factory( + "akernel-thread", KernelFactory(partial(AKernelTask, execute_in_thread=True)) + ) diff --git a/plugins/akernel_task/pyproject.toml b/plugins/akernel_task/pyproject.toml index 38245a7..bf4da91 100644 --- a/plugins/akernel_task/pyproject.toml +++ b/plugins/akernel_task/pyproject.toml @@ -29,5 +29,5 @@ text = "MIT" Homepage = "https://github.com/davidbrochart/akernel" [project.entry-points] -"fps.modules" = {akernel_task = "fps_akernel_task.main:AKernelTaskModule"} -"jupyverse.modules" = {akernel_task = "fps_akernel_task.main:AKernelTaskModule"} +"fps.modules" = {akernel_task = "fps_akernel_task.main:AKernelTaskModule", akernelthread_task = "fps_akernel_task.main:AKernelThreadTaskModule"} +"jupyverse.modules" = {akernel_task = "fps_akernel_task.main:AKernelTaskModule", akernelthread_task = "fps_akernel_task.main:AKernelThreadTaskModule"} diff --git a/pyproject.toml b/pyproject.toml index 06ae9bb..3bc628e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -23,6 +23,8 @@ classifiers = [ ] keywords = [ "jupyter" ] dependencies = [ + "anyio >=4.12.0,<5.0.0", + "anyioutils >=0.7.4,<0.8.0", "python-dateutil", "colorama", "gast >=0.6.0, <0.7.0", @@ -30,11 +32,11 @@ dependencies = [ ] [dependency-groups] -test = [ +dev = [ + "trio", "mypy", "ruff", "pytest", - "pytest-asyncio", "pytest-rerunfailures", "types-python-dateutil", "kernel_driver >=0.0.7", @@ -44,8 +46,8 @@ test = [ [project.optional-dependencies] subprocess = [ - "zmq-anyio >=0.3.9,<0.4.0", - "typer >=0.4.0", + "zmq-anyio >=0.3.13,<0.4.0", + "cyclopts >=4.22.2,<5.0.0", ] react = [ @@ -57,7 +59,7 @@ cache = [ ] [project.scripts] -akernel = "akernel.akernel:cli" +akernel = "akernel.akernel:app" [tool.hatch.build.targets.wheel] ignore-vcs = true @@ -65,6 +67,7 @@ packages = ["src/akernel"] [tool.hatch.build.targets.wheel.shared-data] "share/jupyter/kernels/akernel/kernel.json" = "share/jupyter/kernels/akernel/kernel.json" +"share/jupyter/kernels/akernel-thread/kernel.json" = "share/jupyter/kernels/akernel-thread/kernel.json" [project.urls] Homepage = "https://github.com/davidbrochart/akernel" @@ -75,3 +78,6 @@ exclude = ["examples"] [tool.uv.sources] fps-akernel-task = { workspace = true } + +[tool.pytest.ini_options] +anyio_mode = "auto" diff --git a/share/jupyter/kernels/akernel-thread/kernel.json b/share/jupyter/kernels/akernel-thread/kernel.json new file mode 100644 index 0000000..6569676 --- /dev/null +++ b/share/jupyter/kernels/akernel-thread/kernel.json @@ -0,0 +1,11 @@ +{ + "argv": [ + "akernel", + "launch", + "--execute-in-thread", + "-f", + "{connection_file}" + ], + "display_name": "Python 3 (akernel-thread)", + "language": "python" +} diff --git a/src/akernel/akernel.py b/src/akernel/akernel.py index 9f8955e..27733ea 100644 --- a/src/akernel/akernel.py +++ b/src/akernel/akernel.py @@ -1,26 +1,30 @@ from __future__ import annotations import json -from typing import Optional, cast +from typing import Annotated, cast -import typer from anyio import create_memory_object_stream, create_task_group, run, sleep_forever +from cyclopts import App, Parameter from .connect import connect_channel from .kernel import Kernel from .kernelspec import write_kernelspec -cli = typer.Typer() +app = App() -@cli.command() +@app.command() def install( - mode: str = typer.Argument("", help="Mode of the kernel to install."), - cache_dir: Optional[str] = typer.Option( - None, "-c", help="Path to the cache directory, if mode is 'cache'." - ), -): + mode: str = "", + cache_dir: str | None = None, +) -> None: + """Install the kernel. + + Args: + mode: Mode of the kernel to install. + cache_dir: Path to the cache directory, if mode is 'cache'. + """ kernel_name = "akernel" if mode: modes = mode.split("-") @@ -31,27 +35,48 @@ def install( write_kernelspec(kernel_name, mode, display_name, cache_dir) -@cli.command() +@app.command() def launch( - mode: str = typer.Argument("", help="Mode of the kernel to launch."), - cache_dir: Optional[str] = typer.Option( - None, "-c", help="Path to the cache directory, if mode is 'cache'." - ), - connection_file: str = typer.Option(..., "-f", help="Path to the connection file."), + connection_file: Annotated[str, Parameter(alias=["-f"])], + mode: str = "", + cache_dir: str | None = None, + execute_in_thread: bool = False, ): - akernel = AKernel(mode, cache_dir, connection_file) + """Launch the kernel. + + Args: + mode: Mode of the kernel to launch. + cache_dir: Path to the cache directory, if mode is 'cache'. + connection_file: Path to the connection file. + execute_in_thread: Whether to run user code in a thread. + """ + akernel = AKernel(mode, cache_dir, connection_file, execute_in_thread) run(akernel.start) class AKernel: - def __init__(self, mode, cache_dir, connection_file): - self._to_shell_send_stream, self._to_shell_receive_stream = create_memory_object_stream[list[bytes]]() - self._from_shell_send_stream, self._from_shell_receive_stream = create_memory_object_stream[list[bytes]]() - self._to_control_send_stream, self._to_control_receive_stream = create_memory_object_stream[list[bytes]]() - self._from_control_send_stream, self._from_control_receive_stream = create_memory_object_stream[list[bytes]]() - self._to_stdin_send_stream, self._to_stdin_receive_stream = create_memory_object_stream[list[bytes]]() - self._from_stdin_send_stream, self._from_stdin_receive_stream = create_memory_object_stream[list[bytes]]() - self._from_iopub_send_stream, self._from_iopub_receive_stream = create_memory_object_stream[list[bytes]](max_buffer_size=float("inf")) + def __init__(self, mode, cache_dir, connection_file, execute_in_thread): + self._to_shell_send_stream, self._to_shell_receive_stream = create_memory_object_stream[ + list[bytes] + ]() + self._from_shell_send_stream, self._from_shell_receive_stream = create_memory_object_stream[ + list[bytes] + ]() + self._to_control_send_stream, self._to_control_receive_stream = create_memory_object_stream[ + list[bytes] + ]() + self._from_control_send_stream, self._from_control_receive_stream = ( + create_memory_object_stream[list[bytes]]() + ) + self._to_stdin_send_stream, self._to_stdin_receive_stream = create_memory_object_stream[ + list[bytes] + ]() + self._from_stdin_send_stream, self._from_stdin_receive_stream = create_memory_object_stream[ + list[bytes] + ]() + self._from_iopub_send_stream, self._from_iopub_receive_stream = create_memory_object_stream[ + list[bytes] + ](max_buffer_size=float("inf")) self.kernel = Kernel( self._to_shell_receive_stream, self._from_shell_send_stream, @@ -62,6 +87,7 @@ def __init__(self, mode, cache_dir, connection_file): self._from_iopub_send_stream, mode, cache_dir, + execute_in_thread, ) with open(connection_file) as f: connection_cfg = json.load(f) @@ -136,4 +162,4 @@ async def from_iopub(self) -> None: if __name__ == "__main__": - cli() + app() diff --git a/src/akernel/code.py b/src/akernel/code.py index 40a76a8..5cab63a 100644 --- a/src/akernel/code.py +++ b/src/akernel/code.py @@ -151,7 +151,7 @@ def get_async_code(self) -> str: def get_async_bytecode(self, task_i: int) -> CodeType: tree = self.get_async_ast() - #tree = gast.gast_to_ast(gtree) + # tree = gast.gast_to_ast(gtree) bytecode = compile(tree, filename=f"", mode="exec") return bytecode diff --git a/src/akernel/execution.py b/src/akernel/execution.py index a2618d4..74c8397 100644 --- a/src/akernel/execution.py +++ b/src/akernel/execution.py @@ -2,6 +2,7 @@ import hashlib import pickle +import sys from typing import List, Dict, Tuple, Any from colorama import Fore, Style # type: ignore @@ -27,8 +28,8 @@ def pre_execute( transform = Transform(code, task_i, react) async_bytecode = transform.get_async_bytecode(task_i) exec(async_bytecode, globals_, locals_) - except SyntaxError as e: - exception = e + except SyntaxError as exc: + exception = exc filename = exception.filename if filename == "": filename = f"{Fore.CYAN}Cell{Style.RESET_ALL} {Fore.GREEN}{execution_count}" @@ -151,9 +152,11 @@ async def execute( try: result = await locals_["__async_cell__"]() except KeyboardInterrupt: + with open("log.txt", "a") as f: f.write(f"interrupt\n") interrupted = True - except Exception as e: - traceback = get_traceback(code, e) + except Exception: + exc_type, exception, traceback = sys.exc_info() + traceback = get_traceback(code, exception, traceback) else: cache_execution(cache, cache_info, globals_, result) diff --git a/src/akernel/kernel.py b/src/akernel/kernel.py index 3c2f09a..51ac271 100644 --- a/src/akernel/kernel.py +++ b/src/akernel/kernel.py @@ -1,14 +1,22 @@ -from __future__ import annotations - -import asyncio -import sys import platform -import json +import signal +import sys from io import StringIO from contextvars import ContextVar -from typing import Dict, Any, List, Union, Awaitable, cast - -from anyio import Event, create_task_group, sleep +from typing import Dict, Any, List, Awaitable + +from anyio import ( + Event, + create_memory_object_stream, + create_task_group, + from_thread, + get_cancelled_exc_class, + open_signal_receiver, + run, + sleep, + to_thread, +) +from anyioutils import Task, create_task import comm # type: ignore from akernel.comm.manager import CommManager from akernel.display import display @@ -38,7 +46,7 @@ class Kernel: comm_manager: CommManager kernel_mode: str cell_done: Dict[int, Event] - running_cells: Dict[int, asyncio.Task] + running_cells: Dict[int, Task] _source_map: Dict[str, str] task_i: int execution_count: int @@ -62,6 +70,7 @@ def __init__( from_iopub_send_stream, kernel_mode: str = "", cache_dir: str | None = None, + execute_in_thread: bool = False, ): global KERNEL KERNEL = self @@ -78,6 +87,7 @@ def __init__( self.kernel_mode = kernel_mode self.cache_dir = cache_dir + self.execute_in_thread = execute_in_thread self._concurrent_kernel = None self._multi_kernel = None self._cache_kernel = None @@ -102,6 +112,10 @@ def __init__( else: self.cache = None self.stop_event = Event() + if execute_in_thread: + import threading + + self._stop_event = threading.Event() self.key = "0" def chain_execution(self) -> None: @@ -139,14 +153,16 @@ def init_kernel(self, namespace): return self.globals[namespace] = { - "ainput": self.ainput, - "asyncio": asyncio, "print": self.print, "__task__": self.task, "__chain_execution__": self.chain_execution, "__unchain_execution__": self.unchain_execution, "_": None, } + if self.execute_in_thread: + self.globals[namespace]["input"] = self.input + else: + self.globals[namespace]["ainput"] = self.ainput self.locals[namespace] = {} if self.react_kernel: code = ( @@ -168,8 +184,55 @@ def interrupt(self): task.cancel() self.running_cells = {} + async def thread_execute(self): + while True: + message = await to_thread.run_sync(self.to_thread_queue.get) + if message is None: + return + + parent, idents, async_cell = message + PARENT_VAR.set(parent) + IDENTS_VAR.set(idents) + result = None + exception = None + traceback = None + try: + result = await async_cell() + except get_cancelled_exc_class(): + return + except Exception: + exc_type, exception, traceback = sys.exc_info() + from_thread.run_sync( + self.from_thread_send_stream.send_nowait, (result, exception, traceback) + ) + + async def thread_main(self) -> None: + async with create_task_group() as tg: + tg.start_soon(self.thread_execute) + await to_thread.run_sync(self._stop_event.wait) + tg.cancel_scope.cancel() + + def thread(self) -> None: + run(self.thread_main) + + async def signal_handler(self): + with open_signal_receiver(signal.SIGINT) as signals: + async for signum in signals: + if signum == signal.SIGINT: + self.interrupt() + with open("log.txt", "a") as f: f.write("interrupt\n") + async def start(self) -> None: async with create_task_group() as self.task_group: + self.task_group.start_soon(self.signal_handler) + if self.execute_in_thread: + from queue import Queue + + self.to_thread_queue = Queue() + self.from_thread_send_stream, self.from_thread_receive_stream = ( + create_memory_object_stream(max_buffer_size=1) + ) + self.task_group.start_soon(to_thread.run_sync, self.thread) msg = self.create_message("status", content={"execution_state": self.execution_state}) to_send = serialize(msg, self.key) await self.from_iopub_send_stream.send(to_send) @@ -188,15 +251,21 @@ async def start(self) -> None: async def _start(self) -> None: self.task_group.start_soon(self.listen_shell) self.task_group.start_soon(self.listen_control) - while True: - # run until shutdown request - await self.stop_event.wait() - if self.restart: - # kernel restart - self.stop_event = Event() - else: - # kernel shutdown - break + try: + while True: + # run until shutdown request + await self.stop_event.wait() + if self.restart: + # kernel restart + self.stop_event = Event() + else: + # kernel shutdown + break + except get_cancelled_exc_class(): + if self.execute_in_thread: + self._stop_event.set() + self.to_thread_queue.put(None) + raise async def listen_shell(self) -> None: while True: @@ -205,7 +274,10 @@ async def listen_shell(self) -> None: # if there was a blocking cell execution, and it was interrupted, # let's ignore all the following execution requests until the pipe # is empty - if self.interrupted and self.to_shell_receive_stream.statistics().tasks_waiting_send == 0: + if ( + self.interrupted + and self.to_shell_receive_stream.statistics().tasks_waiting_send == 0 + ): self.interrupted = False msg_list = await self.to_shell_receive_stream.receive() idents, msg_list = feed_identities(msg_list) @@ -289,7 +361,7 @@ async def listen_shell(self) -> None: exception=exception, ) else: - task = asyncio.create_task( + task = create_task( self.execute_and_finish( idents, parent, @@ -297,9 +369,10 @@ async def listen_shell(self) -> None: self.execution_count, code, cache_info, - ) + ), + self.task_group, ) - self.cell_done[self.task_i] = asyncio.Event() + self.cell_done[self.task_i] = Event() self.running_cells[self.task_i] = task self.task_i += 1 self.execution_count += 1 @@ -371,22 +444,41 @@ async def execute_and_finish( if self._chain_execution and prev_task_i in self.cell_done: await self.cell_done[prev_task_i].wait() del self.cell_done[prev_task_i] - PARENT_VAR.set(parent) - IDENTS_VAR.set(idents) parent_header = parent["header"] traceback, exception = [], None namespace = self.get_namespace(parent_header) self._source_map[f""] = code try: - result = await self.locals[namespace][f"__async_cell{task_i}__"]() + with open("log.txt", "a") as f: f.write("execute\n") + if self.execute_in_thread: + self.to_thread_queue.put( + (parent, idents, self.locals[namespace][f"__async_cell{task_i}__"]) + ) + result, exception, traceback = await self.from_thread_receive_stream.receive() + else: + PARENT_VAR.set(parent) + IDENTS_VAR.set(idents) + result = await self.locals[namespace][f"__async_cell{task_i}__"]() except KeyboardInterrupt: + with open("log.txt", "a") as f: f.write("interrupt\n") self.interrupt() - except Exception as e: - exception = e - traceback = get_traceback(code, e, execution_count, self._source_map) + except Exception: + with open("log.txt", "a") as f: f.write("interrupt\n") + if self.execute_in_thread: + raise + else: + exc_type, exception, traceback = sys.exc_info() + traceback = get_traceback(code, exception, traceback, execution_count, self._source_map) else: - await self.show_result(result, self.globals[namespace], parent_header) - cache_execution(self.cache, cache_info, self.globals[namespace], result) + if self.execute_in_thread: + if exception is not None: + traceback = get_traceback(code, exception, traceback, execution_count) + else: + await self.show_result(result, self.globals[namespace], parent_header) + cache_execution(self.cache, cache_info, self.globals[namespace], result) + else: + await self.show_result(result, self.globals[namespace], parent_header) + cache_execution(self.cache, cache_info, self.globals[namespace], result) finally: self.cell_done[task_i].set() del self.locals[namespace][f"__async_cell{task_i}__"] @@ -455,8 +547,26 @@ def task(self, cell_i: int = -1) -> Awaitable: else: i = cell_i if i in self.running_cells: - return self.running_cells[i] - return asyncio.sleep(0) + return self.running_cells[i].wait() + return sleep(0) + + def input(self, prompt: str = "") -> Any: + parent = PARENT_VAR.get() + idents = IDENTS_VAR.get() + if parent["content"]["allow_stdin"]: + msg = self.create_message( + "input_request", + parent_header=parent["header"], + content={"prompt": prompt, "password": False}, + address=idents[0], + ) + to_send = serialize(msg, self.key) + from_thread.run_sync(self.from_stdin_send_stream.send_nowait, to_send) + msg_list = from_thread.run(self.to_stdin_receive_stream.receive) + idents, msg_list = feed_identities(msg_list) + msg = deserialize(msg_list) + if msg["content"]["status"] == "ok": + return msg["content"]["value"] async def ainput(self, prompt: str = "") -> Any: parent = PARENT_VAR.get() @@ -469,11 +579,10 @@ async def ainput(self, prompt: str = "") -> Any: address=idents[0], ) to_send = serialize(msg, self.key) - await self.from_stdin_send_stream.send(to_send) + self.from_stdin_send_stream.send_nowait(to_send) msg_list = await self.to_stdin_receive_stream.receive() idents, msg_list = feed_identities(msg_list) msg = deserialize(msg_list) - idents, msg = res if msg["content"]["status"] == "ok": return msg["content"]["value"] @@ -502,7 +611,10 @@ def print( content={"name": name, "text": text}, ) to_send = serialize(msg, self.key) - self.from_iopub_send_stream.send_nowait(to_send) + if self.execute_in_thread: + from_thread.run_sync(self.from_iopub_send_stream.send_nowait, to_send) + else: + self.from_iopub_send_stream.send_nowait(to_send) def create_message( self, diff --git a/src/akernel/message.py b/src/akernel/message.py index 0ca14dc..2078ab9 100644 --- a/src/akernel/message.py +++ b/src/akernel/message.py @@ -29,7 +29,7 @@ def utcnow() -> datetime: def feed_identities(msg_list: list[bytes]) -> tuple[list[bytes], list[bytes]]: idx = msg_list.index(DELIM) idents = msg_list[:idx] or [b"foo"] - return idents , msg_list[idx + 1 :] # noqa + return idents, msg_list[idx + 1 :] # noqa def create_message_header(msg_type: str, session_id: str, msg_cnt: int) -> dict[str, Any]: diff --git a/src/akernel/traceback.py b/src/akernel/traceback.py index c136830..2aad177 100644 --- a/src/akernel/traceback.py +++ b/src/akernel/traceback.py @@ -1,15 +1,13 @@ from __future__ import annotations -import sys import types -from typing import Dict, Optional, cast +from traceback import extract_tb, format_list +from typing import Dict, cast from colorama import Fore, Style # type: ignore -def get_traceback(code: str, exception, execution_count: int = 0, source_map: Optional[Dict[str, str]] = None): - exc_info = sys.exc_info() - tb = cast(types.TracebackType, exc_info[2]) +def get_traceback(code: str, exception, tb: types.TracebackType, execution_count: int = 0, source_map: Dict[str, str] | None = None): while True: if tb.tb_next is None: break @@ -24,31 +22,39 @@ def get_traceback(code: str, exception, execution_count: int = 0, source_map: Op break stack.reverse() traceback = ["Traceback (most recent call last):"] + get_frame = False + print for frame in stack: - filename = frame.f_code.co_filename - if filename.startswith(" str: return nocolor_tb -@pytest.mark.asyncio -async def test_execute_assign(all_modes): +async def test_execute_assign(): code = dedent( """ a = 1 @@ -47,8 +46,7 @@ async def test_execute_assign(all_modes): assert g == {"a": 1} -@pytest.mark.asyncio -async def test_execute_assign_in_try(all_modes): +async def test_execute_assign_in_try(): code = dedent( """ try: @@ -61,8 +59,7 @@ async def test_execute_assign_in_try(all_modes): assert g == {"a": 1} -@pytest.mark.asyncio -async def test_execute_invalid_syntax(all_modes): +async def test_execute_invalid_syntax(): code = dedent( """ ab cd @@ -88,8 +85,7 @@ async def test_execute_invalid_syntax(all_modes): assert tb_str(t) == expected -@pytest.mark.asyncio -async def test_execute_not_defined(all_modes): +async def test_execute_not_defined(): code = dedent( """ a @@ -107,8 +103,7 @@ async def test_execute_not_defined(all_modes): assert tb_str(t) == expected -@pytest.mark.asyncio -async def test_execute_import_error(all_modes): +async def test_execute_import_error(): code = dedent( """ from .foo import bar @@ -127,12 +122,11 @@ async def test_execute_import_error(all_modes): assert tb_str(t) == excepted -@pytest.mark.asyncio -async def test_execute_async(all_modes): +async def test_execute_async(): code = dedent( """ - import asyncio - await asyncio.sleep(0) + import anyio + await anyio.sleep(0) a = 1 """ ).strip() @@ -140,7 +134,8 @@ async def test_execute_async(all_modes): assert g["a"] == 1 -@pytest.mark.asyncio + +@pytest.mark.skip(reason="React mode not suported") async def test_execute_react_op(): code = dedent( """ @@ -154,7 +149,7 @@ async def test_execute_react_op(): assert g["a"].v == 3 -@pytest.mark.asyncio +@pytest.mark.skip(reason="React mode not suported") async def test_execute_react_func(): code = dedent( """ @@ -169,7 +164,7 @@ async def test_execute_react_func(): assert g["a"].v == sin(2) + 1 -@pytest.mark.asyncio +@pytest.mark.skip(reason="Cache mode not suported") async def test_execute_cache(): cache = {} globals_ = {} diff --git a/tests/test_kernel.py b/tests/test_kernel.py index 08929c8..876ab2c 100644 --- a/tests/test_kernel.py +++ b/tests/test_kernel.py @@ -1,8 +1,10 @@ import os +import platform import sys -import asyncio import signal import re +from functools import partial +from anyio import create_task_group, sleep from pathlib import Path from textwrap import dedent @@ -11,21 +13,22 @@ TIMEOUT = 5 -KERNELSPEC_PATH = str(Path(sys.prefix) / "share" / "jupyter" / "kernels" / "akernel" / "kernel.json") +KERNELSPEC_PATH = str( + Path(sys.prefix) / "share" / "jupyter" / "kernels" / "akernel" / "kernel.json" +) ANSI_ESCAPE = re.compile(r"\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])") def interrupt_kernel(kernel_process): - if sys.platform.startswith("win"): + if platform.system() == "Windows": os.kill(kernel_process.pid, signal.CTRL_C_EVENT) else: kernel_process.send_signal(signal.SIGINT) -@pytest.mark.asyncio -async def test_syntax_error(capfd, all_modes): +async def test_syntax_error(capfd): kd = KernelDriver(kernelspec_path=KERNELSPEC_PATH, log=False) await kd.start(startup_timeout=TIMEOUT) await kd.execute("foo bar", timeout=TIMEOUT) @@ -52,8 +55,7 @@ async def test_syntax_error(capfd, all_modes): assert ANSI_ESCAPE.sub("", err).strip() == expected -@pytest.mark.asyncio -async def test_name_not_defined(capfd, all_modes): +async def test_name_not_defined(capfd): kd = KernelDriver(kernelspec_path=KERNELSPEC_PATH, log=False) await kd.start(startup_timeout=TIMEOUT) await kd.execute("foo", timeout=TIMEOUT) @@ -71,8 +73,7 @@ async def test_name_not_defined(capfd, all_modes): ) -@pytest.mark.asyncio -async def test_hello_world(capfd, all_modes): +async def test_hello_world(capfd): kd = KernelDriver(kernelspec_path=KERNELSPEC_PATH, log=False) await kd.start(startup_timeout=TIMEOUT) await kd.execute("print('Hello World!')", timeout=TIMEOUT) @@ -82,8 +83,7 @@ async def test_hello_world(capfd, all_modes): assert out == "Hello World!\n" -@pytest.mark.asyncio -async def test_global_variable(capfd, all_modes): +async def test_global_variable(capfd): kd = KernelDriver(kernelspec_path=KERNELSPEC_PATH, log=False) await kd.start(startup_timeout=TIMEOUT) await kd.execute("a = 1", timeout=TIMEOUT) @@ -96,113 +96,122 @@ async def test_global_variable(capfd, all_modes): assert out == "1\n3\n" -@pytest.mark.asyncio -async def test_concurrent_cells(capfd, all_modes): - kd = KernelDriver(kernelspec_path=KERNELSPEC_PATH, log=False) - await kd.start(startup_timeout=TIMEOUT) - asyncio.create_task(kd.execute("__unchain_execution__()", timeout=TIMEOUT)) - asyncio.create_task(kd.execute("await asyncio.sleep(0.2)\nprint('done1')", timeout=TIMEOUT)) - asyncio.create_task(kd.execute("await asyncio.sleep(0.1)\nprint('done2')", timeout=TIMEOUT)) - await asyncio.sleep(0.5) - await kd.stop() +async def test_concurrent_cells(capfd): + async with create_task_group() as tg: + kd = KernelDriver(kernelspec_path=KERNELSPEC_PATH, log=False) + await kd.start(startup_timeout=TIMEOUT) + tg.start_soon(partial(kd.execute, "__unchain_execution__()", timeout=TIMEOUT)) + tg.start_soon(partial(kd.execute, "from anyio import sleep", timeout=TIMEOUT)) + tg.start_soon(partial(kd.execute, "await sleep(0.2)\nprint('done1')", timeout=TIMEOUT)) + tg.start_soon(partial(kd.execute, "await sleep(0.1)\nprint('done2')", timeout=TIMEOUT)) + await sleep(0.5) + await kd.stop() out, err = capfd.readouterr() assert out == "done2\ndone1\n" -@pytest.mark.asyncio -async def test_chained_cells(capfd, all_modes): - kd = KernelDriver(kernelspec_path=KERNELSPEC_PATH, log=False) - await kd.start(startup_timeout=TIMEOUT) - asyncio.create_task(kd.execute("await asyncio.sleep(0.2)\nprint('done1')", timeout=TIMEOUT)) - asyncio.create_task( - kd.execute( - "await __task__()\nawait asyncio.sleep(0.1)\nprint('done2')", - timeout=TIMEOUT, +async def test_chained_cells(capfd): + async with create_task_group() as tg: + kd = KernelDriver(kernelspec_path=KERNELSPEC_PATH, log=False) + await kd.start(startup_timeout=TIMEOUT) + tg.start_soon(partial(kd.execute, "from anyio import sleep", timeout=TIMEOUT)) + tg.start_soon(partial(kd.execute, "await sleep(0.2)\nprint('done1')", timeout=TIMEOUT)) + tg.start_soon( + partial( + kd.execute, + "await __task__()\nawait sleep(0.1)\nprint('done2')", + timeout=TIMEOUT, + ) ) - ) - await asyncio.sleep(1) - await kd.stop() - - out, err = capfd.readouterr() - assert out == "done1\ndone2\n" - - -@pytest.mark.asyncio -async def test_interrupt_async(capfd, all_modes): - kd = KernelDriver(kernelspec_path=KERNELSPEC_PATH, log=False) - await kd.start(startup_timeout=TIMEOUT) - asyncio.create_task(kd.execute("__unchain_execution__()", timeout=TIMEOUT)) - expected = [] - n0, n1 = 2, 3 - for i0 in range(n0): - for i1 in range(n1): - asyncio.create_task( - kd.execute( - f"print('{i0} {i1} before')\nawait asyncio.sleep(1)\nprint('{i0} {i1} after')", - timeout=TIMEOUT, + await sleep(1) + await kd.stop() + + out, err = capfd.readouterr() + assert out == "done1\ndone2\n" + + +async def test_interrupt_async(capfd): + async with create_task_group() as tg: + kd = KernelDriver(kernelspec_path=KERNELSPEC_PATH, log=False) + await kd.start(startup_timeout=TIMEOUT) + await kd.execute("__unchain_execution__()", timeout=TIMEOUT) + expected = [] + n0, n1 = 2, 3 + for i0 in range(n0): + for i1 in range(n1): + tg.start_soon( + partial( + kd.execute, + f"from anyio import sleep\nprint('{i0} {i1} before')\nawait sleep(1)\nprint('{i0} {i1} after')", + timeout=TIMEOUT, + ) ) + expected.append(f"{i0} {i1} before") + await sleep(0.1) + interrupt_kernel(kd.kernel_process) + await sleep(0.1) + await kd.stop() + + expected = "\n".join(expected) + "\n" + out, err = capfd.readouterr() + print(expected) + assert out == expected + + +async def test_interrupt_chained(capfd): + async with create_task_group() as tg: + kd = KernelDriver(kernelspec_path=KERNELSPEC_PATH, log=False) + await kd.start(startup_timeout=TIMEOUT) + tg.start_soon( + partial( + kd.execute, + "from anyio import sleep\nprint('before 0')\nawait sleep(1)\nprint('after 0')", + timeout=TIMEOUT, ) - expected.append(f"{i0} {i1} before") - await asyncio.sleep(0.1) - interrupt_kernel(kd.kernel_process) - await asyncio.sleep(0.1) - await kd.stop() - - expected = "\n".join(expected) + "\n" - out, err = capfd.readouterr() - assert out == expected - - -@pytest.mark.asyncio -async def test_interrupt_chained(capfd, all_modes): - kd = KernelDriver(kernelspec_path=KERNELSPEC_PATH, log=False) - await kd.start(startup_timeout=TIMEOUT) - asyncio.create_task( - kd.execute( - "print('before 0')\nawait asyncio.sleep(1)\nprint('after 0')", - timeout=TIMEOUT, ) - ) - asyncio.create_task( - kd.execute( - "await __task__()\nprint('before 1')\nawait asyncio.sleep(1)\nprint('after 1')", - timeout=TIMEOUT, + tg.start_soon( + partial( + kd.execute, + "await __task__()\nprint('before 1')\nawait sleep(1)\nprint('after 1')", + timeout=TIMEOUT, + ) ) - ) - await asyncio.sleep(0.1) - interrupt_kernel(kd.kernel_process) - await asyncio.sleep(0.1) - await kd.stop() - - out, err = capfd.readouterr() - assert out == "before 0\n" - - -@pytest.mark.asyncio -async def test_interrupt_blocking(capfd, all_modes): - kd = KernelDriver(kernelspec_path=KERNELSPEC_PATH, log=False) - await kd.start(startup_timeout=TIMEOUT) - asyncio.create_task( - kd.execute( - "import time\nprint('before 0')\ntime.sleep(1)\nprint('after 0')", - timeout=TIMEOUT, + await sleep(0.1) + interrupt_kernel(kd.kernel_process) + await sleep(0.1) + await kd.stop() + + out, err = capfd.readouterr() + assert out == "before 0\n" + + +async def test_interrupt_blocking(capfd): + async with create_task_group() as tg: + kd = KernelDriver(kernelspec_path=KERNELSPEC_PATH, log=False) + await kd.start(startup_timeout=TIMEOUT) + tg.start_soon( + partial( + kd.execute, + "import time\nprint('before 0')\ntime.sleep(1)\nprint('after 0')", + timeout=TIMEOUT, + ) ) - ) - asyncio.create_task( - kd.execute("print('before 1')\ntime.sleep(1)\nprint('after 1')", timeout=TIMEOUT) - ) - await asyncio.sleep(0.1) - interrupt_kernel(kd.kernel_process) - await asyncio.sleep(0.1) - await kd.stop() + tg.start_soon( + partial( + kd.execute, "print('before 1')\ntime.sleep(1)\nprint('after 1')", timeout=TIMEOUT + ) + ) + await sleep(0.1) + interrupt_kernel(kd.kernel_process) + await sleep(0.1) + await kd.stop() - out, err = capfd.readouterr() - assert out == "before 0\n" + out, err = capfd.readouterr() + assert out == "before 0\n" -@pytest.mark.asyncio -async def test_repr(capfd, all_modes): +async def test_repr(capfd ): kd = KernelDriver(kernelspec_path=KERNELSPEC_PATH, log=False) await kd.start(startup_timeout=TIMEOUT) await kd.execute("1 + 2", timeout=TIMEOUT) @@ -212,8 +221,7 @@ async def test_repr(capfd, all_modes): assert out == "3\n" -@pytest.mark.asyncio -async def test_globals(capfd, all_modes): +async def test_globals(capfd): code = dedent( """\ a = 1