Add `debug_mode: bool` control to task mngr

Allows dynamically importing `pdbp` when enabled and a way for
eventually linking with `tractor`'s own debug mode flag.
oco_supervisor_prototype
Tyler Goodlet 2023-05-19 15:51:47 -04:00
parent e0c888fd5c
commit c7e27ad09d
1 changed files with 18 additions and 7 deletions

View File

@ -23,6 +23,7 @@ from contextlib import (
asynccontextmanager as acm, asynccontextmanager as acm,
contextmanager as cm, contextmanager as cm,
) )
from functools import partial
from typing import ( from typing import (
Generator, Generator,
Any, Any,
@ -32,7 +33,6 @@ from outcome import (
Outcome, Outcome,
acapture, acapture,
) )
import pdbp
from msgspec import Struct from msgspec import Struct
import trio import trio
from trio._core._run import ( from trio._core._run import (
@ -206,6 +206,8 @@ class TaskManagerNursery(Struct):
def add_task_handle_and_crash_handling( def add_task_handle_and_crash_handling(
nursery: Nursery, nursery: Nursery,
debug_mode: bool = False,
) -> Generator[ ) -> Generator[
Any, Any,
Outcome, Outcome,
@ -246,7 +248,9 @@ def add_task_handle_and_crash_handling(
# a REPL on std errors. # a REPL on std errors.
except Exception as err: except Exception as err:
print(f'{task.name} crashed, entering debugger!') print(f'{task.name} crashed, entering debugger!')
pdbp.xpm() if debug_mode:
import pdbp
pdbp.xpm()
raise raise
finally: finally:
@ -255,11 +259,15 @@ def add_task_handle_and_crash_handling(
@acm @acm
async def open_nursery( async def open_nursery(
task_manager = None, task_manager: Generator[Any, Outcome, None] | None = None,
**kwargs,
**lowlevel_nursery_kwargs,
): ):
async with trio.open_nursery(**kwargs) as nurse: async with trio.open_nursery(**lowlevel_nursery_kwargs) as nurse:
yield TaskManagerNursery(nurse, task_manager=task_manager) yield TaskManagerNursery(
nurse,
task_manager=task_manager,
)
async def sleep_then_return_val(val: str): async def sleep_then_return_val(val: str):
@ -284,7 +292,10 @@ if __name__ == '__main__':
async def main(): async def main():
async with open_nursery( async with open_nursery(
task_manager=add_task_handle_and_crash_handling, task_manager=partial(
add_task_handle_and_crash_handling,
debug_mode=True,
),
) as sn: ) as sn:
for _ in range(3): for _ in range(3):
outcome, _ = await sn.start_soon(trio.sleep_forever) outcome, _ = await sn.start_soon(trio.sleep_forever)