Model Base Classes

The model base classes provide serialization (JSON, JSON.GZ, msgpack, YAML), equality checking, and (for settable models) the ability to generate and apply Tao set commands.

TaoBaseModel

Common base for all PyTao Pydantic models. Provides write() and from_file() for serialization to JSON, compressed JSON (.json.gz), msgpack, and YAML.

pytao.model.TaoBaseModel

Bases: BaseModel

A helper base class which allows for creating/updating an instance with Tao objects.

Methods:

pytao.model.TaoBaseModel.from_file classmethod

from_file(filename, *, format=None)

Load Tao model data from a previously-written file.

Parameters:

Name	Type	Description	Default
filename	str or Path		required

Returns:

Type	Description
TaoModel

Source code in pytao/model/base.py

@classmethod
def from_file(
    cls,
    filename: str | pathlib.Path,
    *,
    format: ArchiveFormatLike | None = None,
) -> Self:
    """
    Load Tao model data from a previously-written file.

    Parameters
    ----------
    filename : str or pathlib.Path

    Returns
    -------
    TaoModel
    """
    return load_model(pathlib.Path(filename).resolve(), cls, format=format)

pytao.model.TaoBaseModel.to_dict

to_dict(include_metadata=False, **kwargs)

Serialize the model to a dictionary.

Parameters:

Name	Type	Description	Default
include_metadata	bool	If True, includes base metadata fields like command_args and __class_name__ in the serialized output. Default is False.	False
**kwargs		Additional arguments to pass to Pydantic's model_dump.	{}

Returns:

Type	Description
dict[str, Any]

Source code in pytao/model/base.py

def to_dict(self, include_metadata: bool = False, **kwargs) -> dict[str, Any]:
    """
    Serialize the model to a dictionary.

    Parameters
    ----------
    include_metadata : bool, optional
        If True, includes base metadata fields like `command_args` and
        `__class_name__` in the serialized output. Default is False.
    **kwargs
        Additional arguments to pass to Pydantic's `model_dump`.

    Returns
    -------
    dict[str, Any]
    """
    res = self.model_dump(**kwargs)
    if include_metadata:
        return res
    return _strip_metadata(res, {"__class_name__", "command_args"})

pytao.model.TaoBaseModel.write

write(filename, *, exclude_defaults=True, backup_existing=True, datefmt=DEFAULT_DATEFMT, format=None, indent=False, sort_keys=False)

Write the model data to a file in JSON or YAML format.

Parameters:

Name	Type	Description	Default
filename	str or Path	The path to the file where the model data should be written. The file format is determined by the extension.	required
exclude_defaults	bool	Exclude model defaults from the output. Defaults to True.	True
backup_existing	bool	If a file exists at the target, rename it using datefmt as a backup. Defaults to True.	True
datefmt	str	The date format for the backup file.	DEFAULT_DATEFMT
format	ArchiveFormat or None	File format. If not specified, determined by file extension.	None
indent	bool	Indent the output, if applicable based on the file format.	False
sort_keys	bool	Sort the output keys, if applicable based on the file format.	False

Source code in pytao/model/base.py

def write(
    self,
    filename: str | pathlib.Path,
    *,
    exclude_defaults: bool = True,
    backup_existing: bool = True,
    datefmt: str = DEFAULT_DATEFMT,
    format: ArchiveFormatLike | None = None,
    indent: bool = False,
    sort_keys: bool = False,
):
    """
    Write the model data to a file in JSON or YAML format.

    Parameters
    ----------
    filename : str or pathlib.Path
        The path to the file where the model data should be written.
        The file format is determined by the extension.
    exclude_defaults : bool, optional
        Exclude model defaults from the output. Defaults to True.
    backup_existing : bool, optional
        If a file exists at the target, rename it using `datefmt` as a backup.
        Defaults to True.
    datefmt : str, optional
        The date format for the backup file.
    format : ArchiveFormat or None, optional
        File format.  If not specified, determined by file extension.
    indent : bool, optional
        Indent the output, if applicable based on the file format.
    sort_keys : bool, optional
        Sort the output keys, if applicable based on the file format.
    """

    return dump_model(
        filename,
        self,
        exclude_defaults=exclude_defaults,
        backup_existing=backup_existing,
        datefmt=datefmt,
        format=format,
        indent=indent,
        sort_keys=sort_keys,
    )

TaoModel

Read-only model populated by querying a Tao instance. Each subclass defines a _tao_command_attr_ that maps to a specific Tao pipe command.

pytao.model.TaoModel

Bases: TaoBaseModel

A helper base class which allows for creating/updating an instance with Tao objects.

Methods:

pytao.model.TaoModel.from_tao classmethod

from_tao(tao, **kwargs)

Create this structure by querying Tao for its current values.

Parameters:

Name	Type	Description	Default
tao	Tao		required
**kwargs		Keyword arguments to pass to the relevant tao command. An ele_id keyword may be an integer, string, or ElementID.	{}

Source code in pytao/model/base.py

@classmethod
def from_tao(cls: type[Self], tao: Tao, **kwargs) -> Self:
    """
    Create this structure by querying Tao for its current values.

    Parameters
    ----------
    tao : Tao
    **kwargs
        Keyword arguments to pass to the relevant `tao` command.
        An `ele_id` keyword may be an integer, string, or ElementID.
    """
    cmd_kwargs = cls._tao_command_kwargs(kwargs)
    data = cls._query_tao(tao, cmd_kwargs)
    data = cls._process_tao_data(data)
    return cls(command_args=cmd_kwargs, **data)

pytao.model.TaoModel.query

query(tao)

Query Tao again to generate a new instance of this model.

Source code in pytao/model/base.py

def query(self, tao: Tao) -> Self:
    """Query Tao again to generate a new instance of this model."""
    return self.from_tao(tao, **self.command_args)

TaoSettableModel

Extends TaoModel with the ability to generate set commands and apply changes back to Tao. Supports differential updates (only changed fields), context-manager usage, and error-tolerant application.

pytao.model.TaoSettableModel

Bases: TaoModel

A helper base class which allows for setting Tao parameters based on instance attributes.

Attributes

pytao.model.TaoSettableModel.set_commands property

set_commands

Get all Tao 'set' commands to apply this configuration.

Returns:

Type	Description
list of str

pytao.model.TaoSettableModel.settable_fields property

settable_fields

Names of all 'settable' (modifiable) fields.

Methods:

pytao.model.TaoSettableModel.get_set_commands

get_set_commands(tao=None)

Generate a list of set commands to apply this configuration to tao.

Parameters:

Name	Type	Description	Default
tao	Tao or None	An instance of the Tao class. If provided, only differing configuration parameters will be included in the list of set commands. If None, all attributes to be set will be used.	None

Returns:

Name	Type	Description
cmds	list of str

Source code in pytao/model/base.py

def get_set_commands(self, tao: Tao | None = None) -> list[str]:
    """
    Generate a list of set commands to apply this configuration to `tao`.

    Parameters
    ----------
    tao : Tao or None, optional
        An instance of the Tao class.
        If provided, only differing
        configuration parameters will be included in the list of set
        commands.
        If `None`, all attributes to be set will be used.

    Returns
    -------
    cmds : list of str
    """
    cmds = []
    if tao is not None:
        to_set = self._get_changed_attributes(tao)
    else:
        to_set = self._all_attributes_to_set

    set_name = self._tao_set_name_ or self._tao_command_attr_
    ix_uni = self.command_args.get("ix_uni", None)
    for item in to_set:
        tao_attr_name = _fix_tao_attr_name(item)
        if ix_uni:
            tao_attr_name = f"{ix_uni}@{tao_attr_name}"
        set_value = self._make_set_value(item.value)
        cmds.append(f"set {set_name} {tao_attr_name} = {set_value}")

    if self.model_extra:
        logger.warning(
            "Unhandled extra fields - code regeneration required. Class '%s': %s.",
            type(self).__name__,
            ", ".join(sorted(self.model_extra)),
        )
    return cmds

pytao.model.TaoSettableModel.set

set(tao, *, allow_errors=False, only_changed=False, suppress_plotting=True, suppress_lattice_calc=True, log='DEBUG', exclude_matches=None)

Apply this configuration to Tao.

Parameters:

Name	Type	Description	Default
tao	Tao	The Tao instance to which the configuration will be applied.	required
allow_errors	bool	Allow individual commands to raise errors.	False
only_changed	bool	Only apply changes that differ from the current configuration in Tao.	False
suppress_plotting	bool	Suppress any plotting updates during the commands.	True
suppress_lattice_calc	bool	Suppress lattice calculations during the commands.	True
log	str	The log level to use during the configuration application.	"DEBUG"

Returns:

Name	Type	Description
success	bool	Returns True if the configuration was applied without errors.

Source code in pytao/model/base.py

def set(
    self,
    tao: Tao,
    *,
    allow_errors: bool = False,
    only_changed: bool = False,
    suppress_plotting: bool = True,
    suppress_lattice_calc: bool = True,
    log: str = "DEBUG",
    exclude_matches: list[str] | None = None,
) -> bool:
    """
    Apply this configuration to Tao.

    Parameters
    ----------
    tao : Tao
        The Tao instance to which the configuration will be applied.
    allow_errors : bool, default=False
        Allow individual commands to raise errors.
    only_changed : bool, default=False
        Only apply changes that differ from the current configuration in Tao.
    suppress_plotting : bool, default=True
        Suppress any plotting updates during the commands.
    suppress_lattice_calc : bool, default=True
        Suppress lattice calculations during the commands.
    log : str, default="DEBUG"
        The log level to use during the configuration application.

    Returns
    -------
    success : bool
        Returns True if the configuration was applied without errors.
    """
    cmds = self.get_set_commands(tao=tao if only_changed else None)
    if not cmds:
        return True

    success = True

    tao_global = cast(dict[str, Any], tao.tao_global())
    plot_on = tao_global["plot_on"]
    lat_calc_on = tao_global["lattice_calc_on"]

    exclude_matches = list(exclude_matches or [])

    if suppress_plotting and plot_on:
        tao.cmd("set global plot_on = F")
    if suppress_lattice_calc and lat_calc_on:
        tao.cmd("set global lattice_calc_on = F")

    exclude = [re.compile(line, flags=re.IGNORECASE) for line in exclude_matches]

    log_level: int = getattr(logging, log.upper())

    try:
        for cmd in self.get_set_commands(tao=tao if only_changed else None):
            if any(regex.match(cmd) for regex in exclude):
                continue
            try:
                logger.log(log_level, f"Tao> {cmd}")
                for line in tao.cmd(cmd):
                    logger.log(log_level, f"{line}")
            except Exception as ex:
                if not allow_errors:
                    raise
                success = False
                reason = textwrap.indent(str(ex), "  ")
                logger.error(f"{cmd!r} failed with:\n{reason}")
    finally:
        if suppress_plotting and plot_on:
            tao.cmd("set global plot_on = T")
        if suppress_lattice_calc and lat_calc_on:
            tao.cmd("set global lattice_calc_on = T")

    return success

pytao.model.TaoSettableModel.set_context

set_context(tao)

Apply this configuration to Tao only for the given with block.

Examples:

Set an initial value for a parameter:

>>> new_state.param = 1
>>> new_state.set()

Then temporarily set it to another value, just for the with block:

>>> new_state.param = 3
>>> with new_state.set_context(tao):
...     assert new_state.query(tao).param == 3

After the with block, param will be reset to its previous value:

>>> assert new_state.query(tao).param == 1

Source code in pytao/model/base.py

@contextlib.contextmanager
def set_context(self, tao: Tao):
    """
    Apply this configuration to Tao **only** for the given ``with`` block.

    Examples
    --------

    Set an initial value for a parameter:

    >>> new_state.param = 1
    >>> new_state.set()

    Then temporarily set it to another value, just for the `with` block:

    >>> new_state.param = 3
    >>> with new_state.set_context(tao):
    ...     assert new_state.query(tao).param == 3

    After the ``with`` block, ``param`` will be reset to its previous
    value:

    >>> assert new_state.query(tao).param == 1
    """
    pre_state = self.query(tao)
    for cmd in self.set_commands:
        tao.cmd(cmd)
    yield pre_state
    pre_state.set(tao)

Serialization Helpers

pytao.model.load_model_data

load_model_data(filename, *, format=None, raw=False)

Read the raw model data from a file in JSON, YAML, or msgpack format.

Parameters:

Name	Type	Description	Default
filename	str or Path	The path to the file where the model data should be written. The file format is determined by the extension.	required
format	ArchiveFormat or None	File format. If not specified, determined by file extension.	None
raw	bool	If False, any data not stored in its native format may be returned as-is from the source format. Specifically, when using the MessagePack format, NDArray instances may be returned as an encoded dictionary. If True, all conversions will be applied recursively to the data.	False

Source code in pytao/model/base.py

def load_model_data(
    filename: str | pathlib.Path,
    *,
    format: ArchiveFormatLike | None = None,
    raw: bool = False,
):
    """
    Read the raw model data from a file in JSON, YAML, or msgpack format.

    Parameters
    ----------
    filename : str or pathlib.Path
        The path to the file where the model data should be written.
        The file format is determined by the extension.
    format : ArchiveFormat or None, optional
        File format.  If not specified, determined by file extension.
    raw : bool, optional
        If False, any data not stored in its native format may be returned
        as-is from the source format.  Specifically, when using the MessagePack
        format, NDArray instances may be returned as an encoded dictionary.
        If True, all conversions will be applied recursively to the data.
    """
    fname = pathlib.Path(filename)

    format = ArchiveFormat.from_format_or_file(fname, format)

    if format == ArchiveFormat.yaml:
        import yaml  # NOTE: yaml is not a required dependency

        try:
            loader = yaml.CSafeLoader
        except AttributeError:
            loader = yaml.SafeLoader

        with open(fname, "rt") as fp:
            return yaml.load(fp, Loader=loader)
    elif format == ArchiveFormat.msgpack:
        import ormsgpack

        data = ormsgpack.unpackb(fname.read_bytes(), option=ormsgpack.OPT_NON_STR_KEYS)
        if not raw:
            _msgpack_restore_ndarrays(data)
        return data
    elif format == ArchiveFormat.json_gz:
        with gzip.open(fname, "rb") as fp:
            return orjson.loads(fp.read())
    elif format == ArchiveFormat.json:
        return orjson.loads(fname.read_bytes())
    raise NotImplementedError(format)

pytao.model.load_model

load_model(filename, cls, *, format=None)

Read the model from a file in JSON, YAML, or msgpack format.

Parameters:

Name	Type	Description	Default
filename	str or Path	The path to the file where the model data should be written. The file format is determined by the extension.	required
cls	pydantic.BaseModel class		required
format	ArchiveFormat or None	File format. If not specified, determined by file extension.	None

Source code in pytao/model/base.py

def load_model(
    filename: str | pathlib.Path,
    cls: type[T],
    *,
    format: ArchiveFormatLike | None = None,
) -> T:
    """
    Read the model from a file in JSON, YAML, or msgpack format.

    Parameters
    ----------
    filename : str or pathlib.Path
        The path to the file where the model data should be written.
        The file format is determined by the extension.
    cls : pydantic.BaseModel class

    format : ArchiveFormat or None, optional
        File format.  If not specified, determined by file extension.
    """
    data = load_model_data(filename, format=format, raw=True)
    return cls.model_validate(data)

pytao.model.dump_model

dump_model(filename, model, *, exclude_defaults=True, backup_existing=True, datefmt=DEFAULT_DATEFMT, format=None, indent=False, sort_keys=False)

Write the model data to a file in JSON, YAML, or msgpack format.

Parameters:

Name	Type	Description	Default
filename	str or Path	The path to the file where the model data should be written. The file format is determined by the extension.	required
model	BaseModel		required
exclude_defaults	bool	Exclude model defaults from the output. Defaults to True.	True
backup_existing	bool	If a file exists at the target, rename it using datefmt as a backup. Defaults to True.	True
datefmt	str	The date format for the backup file.	DEFAULT_DATEFMT
format	ArchiveFormat or None	File format. If not specified, determined by file extension.	None
indent	bool	Indent the output, if applicable based on the file format.	False
sort_keys	bool	Sort the output keys, if applicable based on the file format.	False

Source code in pytao/model/base.py

def dump_model(
    filename: str | pathlib.Path,
    model: pydantic.BaseModel,
    *,
    exclude_defaults: bool = True,
    backup_existing: bool = True,
    datefmt: str = DEFAULT_DATEFMT,
    format: ArchiveFormatLike | None = None,
    indent: bool = False,
    sort_keys: bool = False,
):
    """
    Write the model data to a file in JSON, YAML, or msgpack format.

    Parameters
    ----------
    filename : str or pathlib.Path
        The path to the file where the model data should be written.
        The file format is determined by the extension.
    model : pydantic.BaseModel

    exclude_defaults : bool, optional
        Exclude model defaults from the output. Defaults to True.
    backup_existing : bool, optional
        If a file exists at the target, rename it using `datefmt` as a backup.
        Defaults to True.
    datefmt : str, optional
        The date format for the backup file.
    format : ArchiveFormat or None, optional
        File format.  If not specified, determined by file extension.
    indent : bool, optional
        Indent the output, if applicable based on the file format.
    sort_keys : bool, optional
        Sort the output keys, if applicable based on the file format.
    """
    fname = pathlib.Path(filename)

    format = ArchiveFormat.from_format_or_file(fname, format)

    if backup_existing:
        date_coded_rename(fname, datefmt=datefmt)

    if format == ArchiveFormat.msgpack:
        import ormsgpack

        data = model.model_dump(
            exclude_defaults=exclude_defaults,
            exclude_computed_fields=True,
            mode="python",
        )
        dumped = ormsgpack.packb(
            data,
            default=_msgpack_default,
            option=ormsgpack.OPT_NON_STR_KEYS,
        )
        pathlib.Path(fname).write_bytes(dumped)
        return data

    data = model.model_dump(
        exclude_defaults=exclude_defaults,
        exclude_computed_fields=True,
        mode="json",
    )
    if format == ArchiveFormat.yaml:
        import yaml  # NOTE: yaml is not a required dependency

        try:
            dumper = yaml.CSafeDumper
        except AttributeError:
            dumper = yaml.SafeDumper

        with fname.open("wt") as fp:
            yaml.dump(data, fp, Dumper=dumper)
    elif format in (ArchiveFormat.json_gz, ArchiveFormat.json):
        options = 0
        if indent:
            options |= orjson.OPT_INDENT_2
        if sort_keys:
            options |= orjson.OPT_SORT_KEYS
        dumped = orjson.dumps(data, option=options)

        if format == ArchiveFormat.json_gz:
            with gzip.open(fname, "wb") as fp:
                fp.write(dumped)
        else:
            pathlib.Path(fname).write_bytes(dumped)
    else:
        raise NotImplementedError(format)

    return data