jmri_scripting.runner

  1"""``ScriptRunner`` — load + execute Python scripts against a JmriShim.
  2
  3Each script gets:
  4
  5* its own globals dict seeded from :meth:`JmriShim.globals`;
  6* a worker thread it runs on (so sync shim wrappers can bridge to the
  7  event loop without deadlocking);
  8* a unique name for status / stop lookup.
  9
 10A script can be:
 11
 121. A plain ``.py`` file that runs top-to-bottom (statements at module level).
 132. A file that defines ``class MyScript(AbstractAutomaton): ...`` and
 14   calls ``MyScript().start()`` — the runner intercepts ``start`` so the
 15   automaton runs on the worker thread.
 163. A file that defines ``def main(shim): ...`` — called with the shim.
 17
 18For (2), :func:`run_script` automatically detects ``AbstractAutomaton``
 19subclasses and instantiates them. The script doesn't have to import
 20``AbstractAutomaton`` — it's in the injected globals.
 21"""
 22from __future__ import annotations
 23
 24import asyncio
 25import inspect
 26import pathlib
 27import sys
 28import threading
 29import time
 30import traceback
 31from dataclasses import dataclass, field
 32from typing import TYPE_CHECKING, Any, TextIO
 33
 34from jmri_scripting import jython_shim
 35from jmri_scripting.automaton import AbstractAutomaton
 36
 37if TYPE_CHECKING:
 38    from jmri_core.hal.connection import SystemConnection
 39
 40    from jmri_scripting.shim import JmriShim
 41
 42
 43# Per-thread stdout/stderr capture. See :class:`_ThreadDispatchStream`.
 44_DEFAULT_OUTPUT_BUFFER_BYTES = 64 * 1024
 45
 46
 47# ``jmri-autoload: false`` (or ``no``, ``off``, ``0``) anywhere in the
 48# first 20 lines opts the file out of the startup autoloader while
 49# leaving it available for on-demand Run via the /scripts UI. Case-
 50# insensitive on the value; the key spelling is fixed to make grep
 51# straightforward.
 52import re as _re  # noqa: E402
 53_AUTOLOAD_OPTOUT_RE = _re.compile(
 54    r"#\s*jmri-autoload\s*:\s*(false|no|off|0)\b",
 55    _re.IGNORECASE,
 56)
 57_OPTOUT_SCAN_LINES = 20
 58
 59
 60def _autoload_opt_out(source: str) -> bool:
 61    """Return True when *source*'s header carries the opt-out marker.
 62
 63    Restricted to the first :data:`_OPTOUT_SCAN_LINES` so a script
 64    that happens to contain the string in a large docstring or data
 65    block doesn't accidentally opt out."""
 66    head = "\n".join(source.splitlines()[:_OPTOUT_SCAN_LINES])
 67    return bool(_AUTOLOAD_OPTOUT_RE.search(head))
 68
 69
 70class _BoundedTextBuffer:
 71    """Text buffer with FIFO truncation once *max_bytes* is exceeded.
 72
 73    Not thread-safe by design: one script thread owns one buffer, and
 74    reads come from the asyncio thread only after ``finished`` flips or
 75    when the caller explicitly copies :meth:`getvalue`. String concat
 76    at 64 KB is fine on modern hardware — a script printing megabytes
 77    of chatter isn't the workload we're optimising for.
 78    """
 79
 80    __slots__ = ("_data", "_max")
 81
 82    def __init__(self, max_bytes: int = _DEFAULT_OUTPUT_BUFFER_BYTES) -> None:
 83        self._max = max_bytes
 84        self._data: str = ""
 85
 86    def write(self, text: str) -> int:
 87        self._data += text
 88        overshoot = len(self._data) - self._max
 89        if overshoot > 0:
 90            # Drop the oldest bytes; keep a leading truncation marker so
 91            # readers can tell the buffer wrapped.
 92            marker = f"... ({overshoot} bytes elided) ...\n"
 93            self._data = marker + self._data[overshoot + len(marker):]
 94        return len(text)
 95
 96    def flush(self) -> None:  # noqa: D401 — file-like protocol
 97        return
 98
 99    def isatty(self) -> bool:
100        return False
101
102    def writable(self) -> bool:
103        return True
104
105    def getvalue(self) -> str:
106        return self._data
107
108
109class _ThreadDispatchStream:
110    """Route ``write`` calls to a per-thread target, else to *fallback*.
111
112    Installed once at module import time as ``sys.stdout`` (and again
113    as ``sys.stderr``). When a script thread starts, the runner calls
114    :meth:`register` to attach that thread's :class:`_BoundedTextBuffer`;
115    on exit it calls :meth:`unregister`. Writes from other threads
116    (asyncio loop, driver monitors, logging) pass through to the
117    original ``sys.__stdout__`` unchanged.
118
119    Thread-safe: the target map is guarded by a lock, and each
120    ``write`` snapshots the current target before delegating so an
121    unregister racing with a write can't corrupt state.
122    """
123
124    def __init__(self, fallback: TextIO) -> None:
125        self._fallback = fallback
126        self._targets: dict[int, TextIO] = {}
127        self._lock = threading.Lock()
128
129    def register(self, target: TextIO) -> None:
130        with self._lock:
131            self._targets[threading.get_ident()] = target
132
133    def unregister(self) -> None:
134        with self._lock:
135            self._targets.pop(threading.get_ident(), None)
136
137    def _pick(self) -> TextIO:
138        with self._lock:
139            return self._targets.get(threading.get_ident(), self._fallback)
140
141    def write(self, data: str) -> int:
142        return self._pick().write(data)
143
144    def flush(self) -> None:
145        self._pick().flush()
146
147    def isatty(self) -> bool:
148        try:
149            return self._fallback.isatty()
150        except (AttributeError, ValueError):
151            return False
152
153    def writable(self) -> bool:
154        return True
155
156    def fileno(self) -> int:
157        return self._fallback.fileno()
158
159
160def _install_dispatch_streams() -> tuple[_ThreadDispatchStream, _ThreadDispatchStream]:
161    """Idempotently swap :data:`sys.stdout` and :data:`sys.stderr` for
162    :class:`_ThreadDispatchStream` proxies. Returns the two proxies so
163    the runner can register/unregister on them.
164
165    Called from :meth:`ScriptRunner.__init__` (not at module import)
166    so that test frameworks like ``pytest`` that replace ``sys.stdout``
167    per test still get their capture wrapped — the proxy uses whatever
168    ``sys.stdout`` is at install time as its fallback for non-script
169    threads.
170    """
171    stdout = sys.stdout
172    stderr = sys.stderr
173    if isinstance(stdout, _ThreadDispatchStream):
174        stdout_proxy = stdout
175    else:
176        stdout_proxy = _ThreadDispatchStream(stdout)
177        sys.stdout = stdout_proxy
178    if isinstance(stderr, _ThreadDispatchStream):
179        stderr_proxy = stderr
180    else:
181        stderr_proxy = _ThreadDispatchStream(stderr)
182        sys.stderr = stderr_proxy
183    return stdout_proxy, stderr_proxy
184
185
186@dataclass
187class ScriptStatus:
188    """One running (or finished) script."""
189
190    name: str
191    path: pathlib.Path | None
192    started_at: float
193    automaton: AbstractAutomaton | None = None
194    thread: threading.Thread | None = None
195    error: str | None = None
196    finished: bool = False
197    # Per-script captured stdout/stderr. Populated on the script's
198    # worker thread via the dispatch streams above; readable from any
199    # thread — the buffer's underlying string is immutable per read
200    # since ``getvalue`` returns a snapshot.
201    _stdout_buffer: _BoundedTextBuffer = field(
202        default_factory=_BoundedTextBuffer,
203    )
204    _stderr_buffer: _BoundedTextBuffer = field(
205        default_factory=_BoundedTextBuffer,
206    )
207
208    @property
209    def stdout(self) -> str:
210        """Captured ``print()`` output. Truncated FIFO once ~64 KB."""
211        return self._stdout_buffer.getvalue()
212
213    @property
214    def stderr(self) -> str:
215        """Captured ``sys.stderr`` output (raised tracebacks, warnings)."""
216        return self._stderr_buffer.getvalue()
217
218
219class ScriptRunner:
220    """Owns the script lifecycle for one ServerContext."""
221
222    def __init__(self, shim: JmriShim) -> None:
223        self._shim = shim
224        self._scripts: dict[str, ScriptStatus] = {}
225        self._lock = threading.Lock()
226        # Install the stdout/stderr dispatch streams lazily on runner
227        # construction so ``sys.stdout``'s current identity (which may
228        # be pytest's capture, a systemd journal wrapper, or the real
229        # terminal) becomes the fallback for non-script threads.
230        self._stdout_proxy, self._stderr_proxy = _install_dispatch_streams()
231        # Register the synthetic ``jmri.*`` module tree so Jython-shaped
232        # scripts (``from jmri import InstanceManager``) import
233        # unchanged. Idempotent — modules are only registered once
234        # per process.
235        jython_shim.install()
236
237    @classmethod
238    def for_connection(
239        cls, loop: asyncio.AbstractEventLoop, connection: SystemConnection
240    ) -> ScriptRunner:
241        from jmri_scripting.shim import JmriShim  # noqa: PLC0415
242        return cls(JmriShim(loop, connection))
243
244    @property
245    def shim(self) -> JmriShim:
246        return self._shim
247
248    # ----- public API ----------------------------------------------------
249
250    def run_file(self, path: str | pathlib.Path, *, name: str | None = None) -> ScriptStatus:
251        """Compile and start *path*. Returns immediately; script runs in background."""
252        p = pathlib.Path(path)
253        source = p.read_text(encoding="utf-8")
254        name = name or p.stem
255        return self._start(name, source, p)
256
257    def load_directory(
258        self, path: str | pathlib.Path, *, glob: str = "*.py",
259    ) -> list[ScriptStatus]:
260        """Autoload every ``*.py`` under *path* and start it.
261
262        Used on server startup to resume scripts the user parked in the
263        profile's ``python/`` directory — this is the profile-scoped
264        persistence layer.
265
266        **Three opt-outs from autoload**, in order of aggressiveness:
267
268        1. **``_``-prefix** (``_helpers.py``) — private module. Skipped
269           entirely; not run and not visible under Python's ``import``
270           system as a top-level module either.
271        2. **``lib_``-prefix** (``lib_stuff.py``) — importable helper.
272           Skipped from auto-execute but the directory is on
273           ``sys.path`` so ``main.py`` can ``import lib_stuff``.
274        3. **Magic comment** (``# jmri-autoload: false`` in the first
275           few lines) — the file is a regular module, importable by
276           its own name, but the autoloader leaves it alone. Ideal
277           for scripts you want to keep on disk for the on-demand
278           Run button but don't want firing at every server start.
279
280        Read errors (permissions, encoding) skip the file silently.
281        Scripts that fail to compile still get a status entry with
282        ``error`` populated so the UI can surface them.
283        """
284        root = pathlib.Path(path)
285        if not root.is_dir():
286            return []
287        # Register the directory on sys.path once so scripts can
288        # ``import`` peer files as regular Python modules — the
289        # recursive-scripting affordance. Idempotent via ``in`` check.
290        root_str = str(root.resolve())
291        if root_str not in sys.path:
292            sys.path.insert(0, root_str)
293        started: list[ScriptStatus] = []
294        for p in sorted(root.glob(glob)):
295            if p.name.startswith("_") or not p.is_file():
296                continue
297            # ``lib_*.py`` are library modules for other scripts to
298            # ``import``. Don't auto-execute them.
299            if p.name.startswith("lib_"):
300                continue
301            try:
302                source = p.read_text(encoding="utf-8")
303            except OSError:
304                continue
305            if _autoload_opt_out(source):
306                # The magic-comment opt-out. File stays on disk +
307                # visible to the /scripts UI's Run button, but not
308                # autoexecuted at server startup.
309                continue
310            started.append(self._start(p.stem, source, p))
311        return started
312
313    def run_source(self, source: str, *, name: str) -> ScriptStatus:
314        """Compile and start *source* (inline script body)."""
315        return self._start(name, source, None)
316
317    def stop(self, name: str) -> bool:
318        """Signal the named script to stop. Returns True if it was running."""
319        with self._lock:
320            status = self._scripts.get(name)
321            if status is None:
322                return False
323            if status.automaton is not None:
324                status.automaton.stop()
325            # Plain scripts can't be stopped mid-statement; they'll finish
326            # naturally on the next handle() boundary if they're using an
327            # automaton, or run to completion otherwise.
328            return True
329
330    def list(self) -> list[ScriptStatus]:
331        with self._lock:
332            return list(self._scripts.values())
333
334    def get(self, name: str) -> ScriptStatus | None:
335        with self._lock:
336            return self._scripts.get(name)
337
338    def stop_all(self, *, join_timeout: float = 2.0) -> None:
339        with self._lock:
340            statuses = list(self._scripts.values())
341        for s in statuses:
342            if s.automaton is not None:
343                s.automaton.stop()
344        for s in statuses:
345            if s.thread is not None and s.thread.is_alive():
346                s.thread.join(timeout=join_timeout)
347
348    # ----- internals -----------------------------------------------------
349
350    def _start(self, name: str, source: str, path: pathlib.Path | None) -> ScriptStatus:
351        # Re-verify the dispatch streams are still installed. Some test
352        # runners (pytest) and long-running processes (uvicorn --reload)
353        # swap ``sys.stdout`` after our ``__init__`` first installed the
354        # proxy — re-installing here wraps whatever is current so a
355        # freshly-started script still routes through us.
356        if sys.stdout is not self._stdout_proxy or sys.stderr is not self._stderr_proxy:
357            self._stdout_proxy, self._stderr_proxy = _install_dispatch_streams()
358        status = ScriptStatus(
359            name=name, path=path, started_at=time.time(),
360        )
361        with self._lock:
362            # If a previous script with the same name is around, stop it first.
363            existing = self._scripts.get(name)
364            if existing and existing.automaton is not None:
365                existing.automaton.stop()
366            self._scripts[name] = status
367
368        thread = threading.Thread(
369            target=self._exec_script,
370            args=(name, source, path, status),
371            name=f"jmri-script:{name}",
372            daemon=True,
373        )
374        status.thread = thread
375        thread.start()
376        return status
377
378    def _exec_script(
379        self,
380        name: str,
381        source: str,
382        path: pathlib.Path | None,
383        status: ScriptStatus,
384    ) -> None:
385        # Route this thread's stdout/stderr into the status buffers so
386        # ``print()`` and any raised traceback surface in the web UI
387        # rather than getting lost in the server's log stream. The
388        # dispatch streams were installed on runner construction; here
389        # we just register/unregister this thread's target.
390        self._stdout_proxy.register(status._stdout_buffer)
391        self._stderr_proxy.register(status._stderr_buffer)
392        # Bind the shim to this thread so ``jmri.InstanceManager``
393        # can find the active connection during script execution.
394        jython_shim._bind_shim(self._shim)
395        try:
396            self._exec_script_body(name, source, path, status)
397        finally:
398            jython_shim._unbind_shim()
399            self._stdout_proxy.unregister()
400            self._stderr_proxy.unregister()
401
402    def _exec_script_body(
403        self,
404        name: str,
405        source: str,
406        path: pathlib.Path | None,
407        status: ScriptStatus,
408    ) -> None:
409        globals_dict = self._shim.globals() | {
410            "AbstractAutomaton": AbstractAutomaton,
411            "__name__": "__jmri_script__",
412            "__file__": str(path) if path else f"<{name}>",
413        }
414        try:
415            code = compile(source, str(path) if path else f"<{name}>", "exec")
416            exec(code, globals_dict)  # noqa: S102 — script execution is the feature
417        except Exception:
418            status.error = traceback.format_exc()
419            status.finished = True
420            return
421
422        # After top-level execution, look for an automaton subclass to run.
423        automaton_cls = _find_automaton_class(globals_dict)
424        if automaton_cls is not None:
425            automaton = automaton_cls(name=name)
426            automaton.shim = self._shim
427            with self._lock:
428                status.automaton = automaton
429            try:
430                automaton.run()
431            except Exception:
432                status.error = traceback.format_exc()
433        # Or look for a `main(shim)` callable.
434        else:
435            main_fn = globals_dict.get("main")
436            if callable(main_fn):
437                try:
438                    sig = inspect.signature(main_fn)
439                    if len(sig.parameters) > 0:
440                        main_fn(self._shim)
441                    else:
442                        main_fn()
443                except Exception:
444                    status.error = traceback.format_exc()
445        status.finished = True
446
447
448def _find_automaton_class(globals_dict: dict[str, Any]) -> type[AbstractAutomaton] | None:
449    """Find the first ``AbstractAutomaton`` subclass defined in the script."""
450    for value in globals_dict.values():
451        if (
452            isinstance(value, type)
453            and issubclass(value, AbstractAutomaton)
454            and value is not AbstractAutomaton
455            and value.__module__ in ("__jmri_script__", "<string>")
456        ):
457            return value
458    return None