Merge pull request #25 from simonsobs/dev

Update hook arguments, plugins, add hooks

nextline/spawned/events.py → nextline/events.py

@@ -2,7 +2,7 @@
 from dataclasses import dataclass
 from typing import Optional
 
-from nextline.types import PromptNo, TaskNo, ThreadNo, TraceNo
+from nextline.types import PromptNo, RunNo, TaskNo, ThreadNo, TraceNo
 
 
 @dataclass
@@ -13,6 +13,7 @@ class Event:
 @dataclass
 class OnStartTrace(Event):
     started_at: datetime.datetime
+    run_no: RunNo
     trace_no: TraceNo
     thread_no: ThreadNo
     task_no: Optional[TaskNo]
@@ -21,48 +22,59 @@ class OnStartTrace(Event):
 @dataclass
 class OnEndTrace(Event):
     ended_at: datetime.datetime
+    run_no: RunNo
     trace_no: TraceNo
 
 
 @dataclass
 class OnStartTraceCall(Event):
     started_at: datetime.datetime
+    run_no: RunNo
     trace_no: TraceNo
     file_name: str
     line_no: int
     frame_object_id: int
-    call_event: str
+    event: str
 
 
 @dataclass
 class OnEndTraceCall(Event):
     ended_at: datetime.datetime
+    run_no: RunNo
     trace_no: TraceNo
 
 
 @dataclass
 class OnStartCmdloop(Event):
     started_at: datetime.datetime
+    run_no: RunNo
     trace_no: TraceNo
 
 
 @dataclass
 class OnEndCmdloop(Event):
     ended_at: datetime.datetime
+    run_no: RunNo
     trace_no: TraceNo
 
 
 @dataclass
 class OnStartPrompt(Event):
     started_at: datetime.datetime
+    run_no: RunNo
     trace_no: TraceNo
     prompt_no: PromptNo
     prompt_text: str
+    file_name: str
+    line_no: int
+    frame_object_id: int
+    event: str
 
 
 @dataclass
 class OnEndPrompt(Event):
     ended_at: datetime.datetime
+    run_no: RunNo
     trace_no: TraceNo
     prompt_no: PromptNo
     command: str
@@ -71,5 +83,6 @@ class OnEndPrompt(Event):
 @dataclass
 class OnWriteStdout(Event):
     written_at: datetime.datetime
+    run_no: RunNo
     trace_no: TraceNo
     text: str

nextline/plugin/plugins/__init__.py

@@ -14,7 +14,7 @@
     TraceInfoRegistrar,
     TraceNumbersRegistrar,
 )
-from .session import CommandSender, Result, RunSession, Signal
+from .session import CommandSender, Result, RunSession, Signal, OnEvent
 
 
 def register(hook: PluginManager) -> None:
@@ -31,4 +31,5 @@ def register(hook: PluginManager) -> None:
     hook.register(Result)
     hook.register(Signal)
     hook.register(CommandSender)
+    hook.register(OnEvent)
     hook.register(RunSession)

nextline/plugin/plugins/registrars/prompt_info.py

@@ -2,15 +2,15 @@
 import dataclasses
 from logging import getLogger
 
-from nextline.plugin.spec import Context, hookimpl
-from nextline.spawned import (
+from nextline.events import (
     OnEndPrompt,
     OnEndTrace,
     OnEndTraceCall,
     OnStartPrompt,
     OnStartTrace,
     OnStartTraceCall,
 )
+from nextline.plugin.spec import Context, hookimpl
 from nextline.types import PromptInfo, PromptNo, TraceNo
 
 
@@ -95,7 +95,7 @@ async def on_end_trace_call(self, context: Context, event: OnEndTraceCall) -> No
             trace_no=trace_no,
             prompt_no=PromptNo(-1),
             open=False,
-            event=trace_call.call_event,
+            event=trace_call.event,
             file_name=trace_call.file_name,
             line_no=trace_call.line_no,
             trace_call_end=True,
@@ -118,7 +118,7 @@ async def on_start_prompt(self, context: Context, event: OnStartPrompt) -> None:
             trace_no=trace_no,
             prompt_no=prompt_no,
             open=True,
-            event=trace_call.call_event,
+            event=trace_call.event,
             file_name=trace_call.file_name,
             line_no=trace_call.line_no,
             stdout=event.prompt_text,

nextline/plugin/plugins/registrars/prompt_notice.py

@@ -1,8 +1,8 @@
 from logging import getLogger
 from typing import Optional
 
+from nextline.events import OnEndTraceCall, OnStartPrompt, OnStartTraceCall
 from nextline.plugin.spec import Context, hookimpl
-from nextline.spawned import OnEndTraceCall, OnStartPrompt, OnStartTraceCall
 from nextline.types import PromptNotice, RunNo, TraceNo
 
 
@@ -40,7 +40,7 @@ async def on_start_prompt(self, context: Context, event: OnStartPrompt) -> None:
             trace_no=trace_no,
             prompt_no=prompt_no,
             prompt_text=event.prompt_text,
-            event=trace_call.call_event,
+            event=trace_call.event,
             file_name=trace_call.file_name,
             line_no=trace_call.line_no,
         )

nextline/plugin/plugins/registrars/run_info.py

@@ -1,54 +1,53 @@
 import dataclasses
-import datetime
+from datetime import timezone
 from typing import Optional
 
-from nextline import spawned
 from nextline.plugin.spec import Context, hookimpl
 from nextline.types import RunInfo
-from nextline.utils import ExitedProcess
 
 
 class RunInfoRegistrar:
     def __init__(self) -> None:
-        self._script: Optional[str] = None
         self._run_info: Optional[RunInfo] = None
 
-    @hookimpl
-    async def on_change_script(self, script: str) -> None:
-        self._script = script
-
     @hookimpl
     async def on_initialize_run(self, context: Context) -> None:
         assert context.run_arg
+        if isinstance(context.run_arg.statement, str):
+            script = context.run_arg.statement
+        else:
+            script = None
         self._run_info = RunInfo(
-            run_no=context.run_arg.run_no, state='initialized', script=self._script
+            run_no=context.run_arg.run_no, state='initialized', script=script
         )
         await context.pubsub.publish('run_info', self._run_info)
 
     @hookimpl
     async def on_start_run(self, context: Context) -> None:
         assert self._run_info is not None
+        assert context.running_process
+        assert context.running_process.process_created_at.tzinfo is timezone.utc
+        started_at = context.running_process.process_created_at.replace(tzinfo=None)
         self._run_info = dataclasses.replace(
-            self._run_info,
-            state='running',
-            started_at=datetime.datetime.utcnow(),
+            self._run_info, state='running', started_at=started_at
         )
         await context.pubsub.publish('run_info', self._run_info)
 
     @hookimpl
-    async def on_end_run(
-        self, context: Context, exited_process: ExitedProcess[spawned.RunResult]
-    ) -> None:
+    async def on_end_run(self, context: Context) -> None:
         assert self._run_info is not None
-        run_result = exited_process.returned or spawned.RunResult(ret=None, exc=None)
+        assert context.exited_process
+        run_result = context.exited_process.returned
+        assert run_result
+        assert context.exited_process.process_exited_at.tzinfo is timezone.utc
+        ended_at = context.exited_process.process_exited_at.replace(tzinfo=None)
 
         self._run_info = dataclasses.replace(
             self._run_info,
             state='finished',
             result=run_result.fmt_ret,
             exception=run_result.fmt_exc,
-            ended_at=datetime.datetime.utcnow(),
+            ended_at=ended_at,
         )
         await context.pubsub.publish('run_info', self._run_info)
-
         self._run_info = None

nextline/plugin/plugins/registrars/stdout.py

@@ -1,7 +1,7 @@
 from typing import Optional
 
+from nextline.events import OnWriteStdout
 from nextline.plugin.spec import Context, hookimpl
-from nextline.spawned import OnWriteStdout
 from nextline.types import RunNo, StdoutInfo
 
 

nextline/plugin/plugins/registrars/trace_info.py

@@ -1,8 +1,8 @@
 import dataclasses
 import datetime
 
+from nextline.events import OnEndTrace, OnStartTrace
 from nextline.plugin.spec import Context, hookimpl
-from nextline.spawned import OnEndTrace, OnStartTrace
 from nextline.types import TraceInfo, TraceNo
 
 

nextline/plugin/plugins/registrars/trace_nos.py

@@ -1,5 +1,5 @@
+from nextline.events import OnEndTrace, OnStartTrace
 from nextline.plugin.spec import Context, hookimpl
-from nextline.spawned import OnEndTrace, OnStartTrace
 from nextline.types import TraceNo
 
 

nextline/plugin/plugins/session/__init__.py

@@ -1,8 +1,10 @@
 __all__ = [
+    'OnEvent',
     'CommandSender',
     'Result',
     'RunSession',
     'Signal',
 ]
 
+from .monitor import OnEvent
 from .session import CommandSender, Result, RunSession, Signal

nextline/plugin/plugins/session/monitor.py

@@ -1,54 +1,32 @@
-import asyncio
-import time
-from collections.abc import AsyncIterator
-from contextlib import asynccontextmanager
 from logging import getLogger
 
-from nextline import spawned
-from nextline.plugin.spec import Context
-from nextline.spawned import QueueOut
+from nextline import events
+from nextline.plugin.spec import Context, hookimpl
 
 
-@asynccontextmanager
-async def relay_queue(context: Context, queue: QueueOut) -> AsyncIterator[None]:
-    task = asyncio.create_task(_monitor(context, queue))
-    try:
-        yield
-    finally:
-        up_to = 0.05
-        start = time.process_time()
-        while not queue.empty() and time.process_time() - start < up_to:
-            await asyncio.sleep(0)
-        await asyncio.to_thread(queue.put, None)  # type: ignore
-        await task
-
-
-async def _monitor(context: Context, queue: QueueOut) -> None:
-    while (event := await asyncio.to_thread(queue.get)) is not None:
-        await _on_event(context, event)
-
-
-async def _on_event(context: Context, event: spawned.Event) -> None:
-    ahook = context.hook.ahook
-    match event:
-        case spawned.OnStartTrace():
-            await ahook.on_start_trace(context=context, event=event)
-        case spawned.OnEndTrace():
-            await ahook.on_end_trace(context=context, event=event)
-        case spawned.OnStartTraceCall():
-            await ahook.on_start_trace_call(context=context, event=event)
-        case spawned.OnEndTraceCall():
-            await ahook.on_end_trace_call(context=context, event=event)
-        case spawned.OnStartCmdloop():
-            await ahook.on_start_cmdloop(context=context, event=event)
-        case spawned.OnEndCmdloop():
-            await ahook.on_end_cmdloop(context=context, event=event)
-        case spawned.OnStartPrompt():
-            await ahook.on_start_prompt(context=context, event=event)
-        case spawned.OnEndPrompt():
-            await ahook.on_end_prompt(context=context, event=event)
-        case spawned.OnWriteStdout():
-            await ahook.on_write_stdout(context=context, event=event)
-        case _:
-            logger = getLogger(__name__)
-            logger.warning(f'Unknown event: {event!r}')
+class OnEvent:
+    @hookimpl
+    async def on_event_in_process(self, context: Context, event: events.Event) -> None:
+        ahook = context.hook.ahook
+        match event:
+            case events.OnStartTrace():
+                await ahook.on_start_trace(context=context, event=event)
+            case events.OnEndTrace():
+                await ahook.on_end_trace(context=context, event=event)
+            case events.OnStartTraceCall():
+                await ahook.on_start_trace_call(context=context, event=event)
+            case events.OnEndTraceCall():
+                await ahook.on_end_trace_call(context=context, event=event)
+            case events.OnStartCmdloop():
+                await ahook.on_start_cmdloop(context=context, event=event)
+            case events.OnEndCmdloop():
+                await ahook.on_end_cmdloop(context=context, event=event)
+            case events.OnStartPrompt():
+                await ahook.on_start_prompt(context=context, event=event)
+            case events.OnEndPrompt():
+                await ahook.on_end_prompt(context=context, event=event)
+            case events.OnWriteStdout():
+                await ahook.on_write_stdout(context=context, event=event)
+            case _:
+                logger = getLogger(__name__)
+                logger.warning(f'Unknown event: {event!r}')