ENH: add error filtering; optionally reuse SubprocessTao in the test suite

.pre-commit-config.yaml

@@ -31,7 +31,7 @@ repos:
         args: ["--maxkb=100000"]
 
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.6.1
+    rev: v0.6.4
     hooks:
       - id: ruff
         args: ["--config", "pyproject.toml"]

docs/examples/plot-bokeh.ipynb

@@ -266,12 +266,25 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "import pytao\n",
     "from pytao.plotting import TaoGraphSettings, TaoAxisSettings\n",
     "\n",
     "# Let's use SubprocessTao to make an independent Tao instance in a subprocess.\n",
     "# Now we can use `tao` from above (`optics_matching`) and `erl` here simultaneously.\n",
-    "erl = pytao.SubprocessTao(init_file=\"$ACC_ROOT_DIR/bmad-doc/tao_examples/erl/tao.init\", plot=\"bokeh\")\n",
+    "erl = SubprocessTao(\n",
+    "    init_file=\"$ACC_ROOT_DIR/bmad-doc/tao_examples/erl/tao.init\", \n",
+    "    plot=\"bokeh\", \n",
+    "    # The ERL example startup file customizes plots in a way incompatible with pytao's external plotting.\n",
+    "    # Set \"nostartup=True\" to avoid this.\n",
+    "    nostartup=True,\n",
+    ")\n",
+    "\n",
+    "erl.cmds(\n",
+    "    [\n",
+    "        \"set global track_type = beam\",\n",
+    "        \"set var r56[1]|model = 0.234\",\n",
+    "        \"set var t566[1]|model = 0.567\",\n",
+    "    ]\n",
+    ")\n",
     "\n",
     "erl.plot(\n",
     "    \"alpha\",\n",
@@ -320,6 +333,8 @@
     "        {1: TaoCurveSettings(ele_ref_name=r\"linac.beg\\2\")},\n",
     "        {1: TaoCurveSettings(ele_ref_name=r\"linac.end\\2\")},\n",
     "    ],\n",
+    "    xlim=(-3, 3),\n",
+    "    ylim=(-10, 10),\n",
     "    share_x=False,\n",
     "    include_layout=False,\n",
     "    width=350, # per plot\n",

docs/examples/plot-matplotlib.ipynb

@@ -48,7 +48,7 @@
    "outputs": [],
    "source": [
     "import pytao\n",
-    "from pytao import Tao\n",
+    "from pytao import Tao, SubprocessTao\n",
     "\n",
     "import matplotlib.pyplot as plt"
    ]
@@ -292,7 +292,21 @@
     "\n",
     "# Let's use SubprocessTao to make an independent Tao instance in a subprocess.\n",
     "# Now we can use `tao` from above (`optics_matching`) and `erl` here simultaneously.\n",
-    "erl = pytao.SubprocessTao(init_file=\"$ACC_ROOT_DIR/bmad-doc/tao_examples/erl/tao.init\", plot=\"mpl\")\n",
+    "erl = SubprocessTao(\n",
+    "    init_file=\"$ACC_ROOT_DIR/bmad-doc/tao_examples/erl/tao.init\", \n",
+    "    plot=\"mpl\", \n",
+    "    # The ERL example startup file customizes plots in a way incompatible with pytao's external plotting.\n",
+    "    # Set \"nostartup=True\" to avoid this.\n",
+    "    nostartup=True,\n",
+    ")\n",
+    "\n",
+    "erl.cmds(\n",
+    "    [\n",
+    "        \"set global track_type = beam\",\n",
+    "        \"set var r56[1]|model = 0.234\",\n",
+    "        \"set var t566[1]|model = 0.567\",\n",
+    "    ]\n",
+    ")\n",
     "\n",
     "erl.plot(\n",
     "    \"alpha\",\n",
@@ -341,6 +355,8 @@
     "        {1: TaoCurveSettings(ele_ref_name=r\"linac.beg\\2\")},\n",
     "        {1: TaoCurveSettings(ele_ref_name=r\"linac.end\\2\")},\n",
     "    ],\n",
+    "    xlim=(-3, 3),\n",
+    "    ylim=(-10, 10),\n",
     "    share_x=False,\n",
     "    include_layout=False,\n",
     "    figsize=(6, 6),\n",

pytao/__init__.py

@@ -11,11 +11,11 @@
 from .tao_ctypes import TaoModel, run_tao
 from .tao_ctypes.core import (
     TaoCommandError,
-    TaoException,
     TaoInitializationError,
     TaoSharedLibraryNotFoundError,
 )
 from .tao_ctypes.evaluate import evaluate_tao
+from .tao_ctypes.util import TaoException, filter_tao_messages, filter_tao_messages_context
 from .tao_interface import tao_interface
 from .tao_pexpect import tao_io
 
@@ -35,6 +35,8 @@
     "TaoSharedLibraryNotFoundError",
     "TaoStartup",
     "evaluate_tao",
+    "filter_tao_messages",
+    "filter_tao_messages_context",
     "run_tao",
     "tao_interface",
     "tao_io",

pytao/interface.tpl.py

@@ -400,7 +400,7 @@ def __init__(
             logger.debug("Failed to register cell magic", exc_info=True)
 
     def __repr__(self) -> str:
-        return f"<Tao init={self.init_settings.tao_init!r} so_lib_file={self.so_lib_file!r}>"
+        return f"<{self.__class__.__name__} init={self.init_settings.tao_init!r} so_lib_file={self.so_lib_file!r}>"
 
     @override
     def init(
@@ -595,14 +595,17 @@ def _check_tao_version(self):
                 f"\nSome features may not work as expected.  Please upgrade bmad."
             )
 
-    def _init(self, startup: TaoStartup):
+    def _reset_graph_managers(self) -> None:
         for manager in self._graph_managers.values():
             try:
                 manager.tao_init_hook()
             except Exception:
                 logger.exception(
                     "Tao plot manager re-initialization failure (%s)", type(manager)
                 )
+
+    def _init(self, startup: TaoStartup):
+        self._reset_graph_managers()
         return super().init(startup.tao_init)
 
     def __execute(
@@ -729,6 +732,50 @@ def bunch_data(self, ele_id, *, which="model", ix_bunch=1, verbose=False):
 
         return dat
 
+    def cmds(self, cmds, suppress_lattice_calc=True, suppress_plotting=True, raises=True):
+        """
+        Runs a list of commands, optionally suppressing lattice calculations
+        and plotting updates.
+
+        Parameters
+        ----------
+        cmds : list of str
+            List of string commands.
+        suppress_lattice_calc : bool, default=True
+            Suppress lattice calc when applying the commands.
+        suppress_plotting  : bool, default=True
+            Suppress plotting when applying commands.
+        raises : bool, default=True
+            Raise an exception of [ERROR or [FATAL is detected in the output.
+
+        Returns
+        -------
+        list
+            Results corresponding to the commands
+
+        """
+        # Get globals to detect plotting
+        g = self.tao_global()
+        ploton, laton = g["plot_on"], g["lattice_calc_on"]
+
+        if suppress_plotting and ploton:
+            self.cmd("set global plot_on = F")
+        if suppress_lattice_calc and laton:
+            self.cmd("set global lattice_calc_on = F")
+
+        # Actually apply commands
+        results = []
+        for cmd in cmds:
+            res = self.cmd(cmd, raises=raises)
+            results.append(res)
+
+        if suppress_plotting and ploton:
+            self.cmd("set global plot_on = T")
+        if suppress_lattice_calc and laton:
+            self.cmd("set global lattice_calc_on = T")
+
+        return results
+
     def version(self) -> Optional[datetime.datetime]:
         """Get the date-coded version."""
         cmd = "show version"

pytao/interface_commands.py

@@ -2,7 +2,7 @@
 # AUTOGENERATED FILE - DO NOT MODIFY
 # This file was generated by the script `generate_interface_commands.py`.
 # Any modifications may be overwritten.
-# Generated on: 2024-08-28 10:03:38
+# Generated on: 2024-09-10 09:57:46
 # ==============================================================================
 
 from __future__ import annotations
@@ -407,7 +407,7 @@ def __init__(
             logger.debug("Failed to register cell magic", exc_info=True)
 
     def __repr__(self) -> str:
-        return f"<Tao init={self.init_settings.tao_init!r} so_lib_file={self.so_lib_file!r}>"
+        return f"<{self.__class__.__name__} init={self.init_settings.tao_init!r} so_lib_file={self.so_lib_file!r}>"
 
     @override
     def init(
@@ -602,14 +602,17 @@ def _check_tao_version(self):
                 f"\nSome features may not work as expected.  Please upgrade bmad."
             )
 
-    def _init(self, startup: TaoStartup):
+    def _reset_graph_managers(self) -> None:
         for manager in self._graph_managers.values():
             try:
                 manager.tao_init_hook()
             except Exception:
                 logger.exception(
                     "Tao plot manager re-initialization failure (%s)", type(manager)
                 )
+
+    def _init(self, startup: TaoStartup):
+        self._reset_graph_managers()
         return super().init(startup.tao_init)
 
     def __execute(
@@ -736,6 +739,50 @@ def bunch_data(self, ele_id, *, which="model", ix_bunch=1, verbose=False):
 
         return dat
 
+    def cmds(self, cmds, suppress_lattice_calc=True, suppress_plotting=True, raises=True):
+        """
+        Runs a list of commands, optionally suppressing lattice calculations
+        and plotting updates.
+
+        Parameters
+        ----------
+        cmds : list of str
+            List of string commands.
+        suppress_lattice_calc : bool, default=True
+            Suppress lattice calc when applying the commands.
+        suppress_plotting  : bool, default=True
+            Suppress plotting when applying commands.
+        raises : bool, default=True
+            Raise an exception of [ERROR or [FATAL is detected in the output.
+
+        Returns
+        -------
+        list
+            Results corresponding to the commands
+
+        """
+        # Get globals to detect plotting
+        g = self.tao_global()
+        ploton, laton = g["plot_on"], g["lattice_calc_on"]
+
+        if suppress_plotting and ploton:
+            self.cmd("set global plot_on = F")
+        if suppress_lattice_calc and laton:
+            self.cmd("set global lattice_calc_on = F")
+
+        # Actually apply commands
+        results = []
+        for cmd in cmds:
+            res = self.cmd(cmd, raises=raises)
+            results.append(res)
+
+        if suppress_plotting and ploton:
+            self.cmd("set global plot_on = T")
+        if suppress_lattice_calc and laton:
+            self.cmd("set global lattice_calc_on = T")
+
+        return results
+
     def version(self) -> Optional[datetime.datetime]:
         """Get the date-coded version."""
         cmd = "show version"

pytao/plotting/bokeh.py

@@ -43,6 +43,7 @@
 from typing_extensions import NotRequired, TypedDict
 
 from ..interface_commands import AnyPath
+from ..tao_ctypes.core import TaoCommandError
 from . import floor_plan_shapes, pgplot, util
 from .curves import CurveIndexToCurve, PlotCurveLine, PlotCurveSymbols, TaoCurveSettings
 from .fields import ElementField
@@ -70,7 +71,6 @@
     UnsupportedGraphError,
 )
 from .settings import TaoGraphSettings
-from ..tao_ctypes.core import TaoCommandError
 from .types import FloatVariableInfo
 from .util import Limit, OptionalLimit, fix_grid_limits
 

pytao/plotting/floor_plan_shapes.py

@@ -3,11 +3,11 @@
 import math
 from functools import cached_property
 from typing import List, Optional, Union
-from typing_extensions import Literal
 
 import numpy as np
 import pydantic.dataclasses as dataclasses
 from pydantic import ConfigDict
+from typing_extensions import Literal
 
 from . import util
 from .curves import PlotCurveLine

pytao/plotting/plot.py

@@ -1,10 +1,10 @@
 from __future__ import annotations
 
-from functools import partial
 import logging
 import math
 import typing
 from abc import ABC, abstractmethod
+from functools import partial
 from typing import (
     Any,
     ClassVar,
@@ -46,6 +46,7 @@
     FloorOrbitInfo,
     FloorPlanElementInfo,
     Limit,
+    OptionalLimit,
     PlotCurveInfo,
     PlotGraphInfo,
     PlotHistogramInfo,
@@ -54,11 +55,9 @@
     PlotRegionInfo,
     Point,
     WaveParams,
-    OptionalLimit,
 )
 from .util import fix_grid_limits
 
-
 if typing.TYPE_CHECKING:
     from .. import Tao
 

pytao/plotting/types.py

@@ -1,4 +1,4 @@
-from typing import List, Tuple, Optional
+from typing import List, Optional, Tuple
 
 from typing_extensions import NotRequired, TypedDict
 

pytao/plotting/util.py

@@ -5,9 +5,10 @@
 import sys
 from typing import List, Optional, Sequence, Tuple, Union
 
-from .types import OptionalLimit, Limit
 import numpy as np
 
+from .types import Limit, OptionalLimit
+
 logger = logging.getLogger(__name__)