"""Code generation for annotated ophyd device classes using Jinja2 templates."""
import json
from dataclasses import dataclass, field
from inspect import Signature
from logging import Logger
from pathlib import Path
import autoflake
import black
from frappy.client import get_datatype
from frappy.datatypes import CommandType, DataType, EnumType, StructOf
from jinja2 import Environment, PackageLoader, select_autoescape
from ophyd_async.core import SignalR, SignalRW
from ophyd_async.core import StandardReadableFormat as Format
from secop_ophyd.SECoPDevices import (
IGNORED_PROPS,
ParameterType,
PropertyType,
class_from_interface,
)
from secop_ophyd.SECoPSignal import secop_dtype_obj_from_json
from secop_ophyd.util import (
SECoPdtype,
build_command_signature,
command_dtype_to_annotation_str,
secop_enum_name_to_python,
)
[docs]
def internalize_name(name: str) -> str:
"""how to create internal names"""
if name.startswith("_"):
return name[1:]
return name
[docs]
@dataclass
class EnumMember:
"""Represents an enum member with name and value."""
name: str # Python identifier (e.g., "LOW")
value: str # Original SECoP string (e.g., "Low Energy")
description: str | None = None # Optional description
[docs]
@dataclass
class EnumClass:
"""Represents an enum class definition."""
name: str # Enum class name (e.g., "TemperatureRegulatorModeEnum")
members: list[EnumMember]
description: str | None = None # Optional enum description
base_enum_class: str = "StrictEnum" # "StrictEnum" or "SupersetEnum"
def _build_enum_class(
enum_class_name: str, members_dict: dict, description: str
) -> "EnumClass | None":
"""Build an EnumClass from a SECoP enum members dict ({name: value}),
shared by the parameter and command enum-detection code paths.
Returns None if there are no members."""
if not members_dict:
return None
enum_members = [
EnumMember(name=secop_enum_name_to_python(member_name), value=member_name)
for member_name in members_dict
]
return EnumClass(
name=enum_class_name, members=enum_members, description=description
)
[docs]
@dataclass
class ModuleAttribute:
"""Represents a module attribute with name, type, and optional description."""
name: str
type: str
[docs]
@dataclass
class PropertyAttribute:
"""Represents a module property attribute with name, type"""
name: str
type_param: str | None = (
None # Optional type parameter like float for SignalRW[float]
)
path_annotation: str | None = (
None # Annotation like ParamPath(...) or PropPath(...)
)
type: str = "SignalR" # Default to SignalR for properties
[docs]
@dataclass
class ParameterAttribute:
"""Represents a module parameter attribute with name, type, and
optional description."""
name: str
type: str
type_param: str | None = (
None # Optional type parameter like float for SignalRW[float]
)
description: str | None = None # Optional description from SECoP or docstrings
path_annotation: str | None = (
None # Annotation like ParamPath(...) or PropPath(...)
)
format_annotation: str | None = None # StandardReadableFormat.CONFIG_SIGNAL, etc.
[docs]
@dataclass
class CommandAttribute:
"""Represents a SECoP command exposed as a Command/TriggerableCommand
class-level annotation (the '<name>' attribute)."""
name: str # SECoP command name, e.g. "test_cmd"; also the attribute name
command_type: str = "Command" # "Command" or "TriggerableCommand"
arg_type: str | None = None # e.g. "dict[str, Any]"; None means no argument
return_type: str | None = None # e.g. "int"; None means no result
description: str | None = None
# Human-readable call signature, e.g. "execute(*, preset: PresetEnum, ...) -> None",
# rendered as an attribute docstring so IDEs can show it on hover -- the
# Command[[ArgT], ResT] annotation itself can't express per-field keyword
# arguments for struct commands (ParamSpec only encodes a positional type list).
call_signature: str | None = None
[docs]
class Method:
"""Represents a class method with signature and description.
This class supports both old-style initialization (for backward compatibility)
and new-style dataclass initialization.
"""
def __init__(self, cmd_name: str, description: str, cmd_sign: Signature) -> None:
"""Initialize Method (backward compatibility constructor).
Args:
cmd_name: Name of the command
description: Description of the command
cmd_sign: Signature of the command
"""
raw_sig_str: str = str(cmd_sign)
raw_sig_str = raw_sig_str.replace("typing.", "")
if "self" in raw_sig_str:
sig_str = raw_sig_str
else:
sig_str = "(self, " + raw_sig_str[1:]
self.name: str = cmd_name
self.signature: str = sig_str
self.description: str = description
[docs]
@dataclass
class ModuleClass:
"""Represents a module class to be generated."""
name: str
bases: list[str]
parameters: list[ParameterAttribute] = field(default_factory=list)
properties: list[PropertyAttribute] = field(default_factory=list)
methods: list[Method] = field(default_factory=list)
commands: list[CommandAttribute] = field(default_factory=list)
description: str = ""
enums: list[EnumClass] = field(default_factory=list) # Enum classes for this module
[docs]
@dataclass
class NodeClass:
"""Represents a node class to be generated."""
name: str
bases: list[str]
properties: list[PropertyAttribute] = field(default_factory=list)
modules: list[ModuleAttribute] = field(default_factory=list)
description: str = ""
[docs]
class GenNodeCode:
"""Generates annotated Python classes for a single SECoP node.
Each instance describes exactly one SEC node. ``write_gen_node_class_file``
writes it to its own file, named after the generated node class, inside
the configured output directory -- overwriting any previous file of that
name.
The generated code uses Jinja2 templates and is formatted with Black.
"""
default_output_dir: str = "secop_ophyd_devs"
def __init__(self, path: str | None = None, log=None):
"""Initialize the code generator.
Args:
path: Optional path to the output folder. Defaults to
``secop_ophyd_devs`` relative to the current working
directory when not given.
log: Optional logger instance
"""
self.log: Logger | None = log
self.module_folder_path: Path = Path(path or self.default_output_dir)
# Data structures for classes and imports
self.imports: dict[str, set[str] | None] = {}
self.module_classes: list[ModuleClass] = []
self.node_classes: list[NodeClass] = []
self.enum_classes: list[EnumClass] = []
self.inline_comment_threshold: int = 120
self.comment_wrap_width: int = 100
# Required imports for generated classes
self.add_import("typing", "Annotated as A")
self.add_import("ophyd_async.core", "SignalR")
self.add_import("ophyd_async.core", "SignalRW")
self.add_import("ophyd_async.core", "Command")
self.add_import("ophyd_async.core", "TriggerableCommand")
self.add_import("ophyd_async.core", "StandardReadableFormat as Format")
self.add_import("ophyd_async.core", "StrictEnum")
self.add_import("ophyd_async.core", "SupersetEnum")
self.add_import("typing", "Any")
self.add_import("numpy", "ndarray")
self.add_import("secop_ophyd.SECoPDevices", "ParameterType as ParamT")
self.add_import("secop_ophyd.SECoPDevices", "PropertyType as PropT")
# Add necessary Device imports
self.add_import("secop_ophyd.SECoPDevices", "SECoPDevice")
self.add_import("secop_ophyd.SECoPDevices", "SECoPCommunicatorDevice")
self.add_import("secop_ophyd.SECoPDevices", "SECoPReadableDevice")
self.add_import("secop_ophyd.SECoPDevices", "SECoPTriggerableDevice")
self.add_import("secop_ophyd.SECoPDevices", "SECoPWritableDevice")
self.add_import("secop_ophyd.SECoPDevices", "SECoPMoveableDevice")
self.add_import("secop_ophyd.SECoPDevices", "SECoPNodeDevice")
# Setup Jinja2 environment
self.jinja_env = Environment(
loader=PackageLoader("secop_ophyd", "templates"),
autoescape=select_autoescape(),
trim_blocks=False,
lstrip_blocks=False,
keep_trailing_newline=True,
)
def _normalize_description(self, description: str | None) -> str:
"""Normalize description text for generated comments.
- Trim trailing whitespace/newlines
- Preserve intentional internal newlines
"""
if description is None:
return ""
normalized = description.rstrip()
return normalized if normalized else ""
[docs]
def add_import(self, module: str, class_str: str | None = None):
"""Add an import to the import dictionary.
Args:
module: Python module to import from
class_str: Class/symbol to import. If None or empty, imports the module
directly.
"""
if class_str is None or class_str == "":
# For module-only imports (import module), use None as value
if module not in self.imports:
self.imports[module] = None
else:
existing = self.imports.get(module)
if existing is None:
# Convert from module-only import to specific imports
self.imports[module] = {class_str}
elif isinstance(existing, set):
existing.add(class_str)
else:
self.imports[module] = {class_str}
[docs]
def add_mod_class(
self,
module_cls: str,
bases: list[str],
parameters: list[ParameterAttribute],
properties: list[PropertyAttribute],
cmd_plans: list[Method],
description: str = "",
enum_classes: list[EnumClass] | None = None,
commands: list[CommandAttribute] | None = None,
):
"""Add a module class to be generated.
Args:
module_cls: Name of the module class
bases: Base classes
parameters: List of parameter attributes
properties: List of property attributes
cmd_plans: List of method definitions
description: Optional class description
commands: List of command annotations (Command/TriggerableCommand)
"""
# Check if this module class was already added, e.g. two module
# instances within the same node share this implementation class
existing_class = next(
(cls for cls in self.module_classes if cls.name == module_cls), None
)
if existing_class:
# Class already exists - merge enums if provided
if enum_classes:
existing_class.enums.extend(enum_classes)
if self.log:
self.log.info(
f"Module class {module_cls} already exists, "
f"merged {len(enum_classes)} enum(s)"
)
return
mod_cls = ModuleClass(
name=module_cls,
bases=bases,
parameters=parameters,
properties=properties,
methods=cmd_plans,
commands=commands or [],
description=description,
enums=enum_classes or [],
)
self.module_classes.append(mod_cls)
[docs]
def add_node_class(
self,
node_cls: str,
bases: list[str],
properties: list[PropertyAttribute],
modules: list[ModuleAttribute],
description: str = "",
):
"""Add a node class to be generated.
Args:
node_cls: Name of the node class
bases: Base classes
attrs: List of attribute tuples. Supported formats:
- (name, type)
- (name, type, type_param)
- (name, type, type_param, description, category)
"""
node_class = NodeClass(
name=node_cls,
bases=bases,
properties=properties,
modules=modules,
description=description,
)
self.node_classes.append(node_class)
def _parse_command_signature(
self, cmd_name: str, datainfo: dict, description: str
) -> Method:
"""Parse command datainfo to create Method signature.
Args:
cmd_name: Name of the command
datainfo: Command datainfo with argument/result types
description: Command description
Returns:
Method object with signature
"""
# Extract argument and result types
arg_type = datainfo.get("argument")
# Create a basic signature object
sig = Signature.from_callable(lambda self, wait_for_idle=False: None)
if arg_type is not None:
sig = Signature.from_callable(lambda self, arg, wait_for_idle=False: None)
return Method(cmd_name=cmd_name, description=description, cmd_sign=sig)
[docs]
def from_json_describe(self, json_data: str | dict):
"""Generate classes from a SECoP JSON describe message.
Args:
json_data: JSON string or dict containing SECoP describe data
"""
# Parse JSON if string
if isinstance(json_data, str):
describe_data = json.loads(json_data)
else:
describe_data = json_data
modules: dict[str, dict] = describe_data.get("modules", {})
node_properties = {k: v for k, v in describe_data.items() if k != "modules"}
# Parse modules
node_module_attrs: list[ModuleAttribute] = []
node_property_attrs: list[PropertyAttribute] = []
for modname, moddescr in modules.items():
# separate accessibles into command and parameters
parameters = {}
commands = {}
accessibles = moddescr["accessibles"]
for aname, aentry in accessibles.items():
iname = internalize_name(aname)
datatype = get_datatype(aentry["datainfo"], iname)
aentry = dict(aentry, datatype=datatype)
if datatype.IS_COMMAND:
commands[iname] = aentry
else:
parameters[iname] = aentry
properties = {k: v for k, v in moddescr.items() if k != "accessibles"}
# Add module class (highest secop interface class) that the actual
# module class is derived from
secop_ophyd_modclass = class_from_interface(properties)
module_bases = [secop_ophyd_modclass.__name__]
# Add the module class, use self reported "implementation" module property,
# if not present use the module name
module_class = modname
if properties.get("implementation", ""):
module_class = properties.get("implementation", "").split(".").pop()
module_class_list = (
module_class.replace(" ", "_").replace("-", "_").split("_")
)
module_class = "".join(word.capitalize() for word in module_class_list)
# Module enum classes
module_enum_classes = []
# Prepare attributes
# Module Commands
command_plans: list[Method] = []
mod_commands: list[CommandAttribute] = []
for command, command_data in commands.items():
# "stop" is skipped: SECoPMoveableDevice already implements
# Stoppable.stop() (a real bound method) via StandardMovable,
# so the generated annotation must not claim an attribute
# that will never actually be created (see connect_real's
# matching skip in SECoPDevices.py).
if command == "stop":
continue
cmd_datatype: CommandType = command_data["datatype"]
arg_dt, res_dt = cmd_datatype.argument, cmd_datatype.result
command_name_list = (
command.replace(" ", "_").replace("-", "_").split("_")
)
command_camel = "".join(word.capitalize() for word in command_name_list)
arg_type = (
command_dtype_to_annotation_str(arg_dt)
if arg_dt is not None
else None
)
if isinstance(arg_dt, EnumType):
enum_cls = _build_enum_class(
f"{module_class}_{command_camel}_Arg_Enum",
arg_dt.export_datatype().get("members", {}),
f"{command} argument enum for `{module_class}`.",
)
if enum_cls:
module_enum_classes.append(enum_cls)
arg_type = enum_cls.name
return_type = (
command_dtype_to_annotation_str(res_dt)
if res_dt is not None
else None
)
if isinstance(res_dt, EnumType):
enum_cls = _build_enum_class(
f"{module_class}_{command_camel}_Result_Enum",
res_dt.export_datatype().get("members", {}),
f"{command} result enum for `{module_class}`.",
)
if enum_cls:
module_enum_classes.append(enum_cls)
return_type = enum_cls.name
is_triggerable = arg_dt is None and res_dt is None
description = self._normalize_description(
command_data.get("description", "")
)
# Struct args unravel into one keyword-only parameter per member --
# a shape the Command[[ArgT], ResT] annotation above can't express,
# so spell it out here. build_command_signature() is the same
# function used to build the real runtime signature, so the
# rendered text matches `.signature` exactly (including the
# dynamically-named member enum classes, e.g. PresetEnum).
# For a bare (non-struct) arg/result, the already-resolved
# arg_type/return_type (which reuse the concrete generated enum
# class name, unlike a fresh build_command_signature() call) are
# used directly instead.
call_signature: str | None = None
if not is_triggerable:
if isinstance(arg_dt, StructOf):
sig = build_command_signature(cmd_datatype)
call_signature = "execute" + str(sig).replace("typing.", "")
else:
arg_part = f"arg: {arg_type}, " if arg_type else ""
call_signature = (
f"execute({arg_part}wait_for_idle: bool = False) -> "
f"{return_type or 'None'}"
)
mod_commands.append(
CommandAttribute(
name=command,
command_type=(
"TriggerableCommand" if is_triggerable else "Command"
),
arg_type=arg_type,
return_type=return_type,
description=description,
call_signature=call_signature,
)
)
mod_parameters: list[ParameterAttribute] = []
for param_name, param_data in parameters.items():
descr = self._normalize_description(param_data.get("description", ""))
unit = param_data["datainfo"].get("unit")
if unit:
param_descr = (
f"{descr}; Unit: ({unit})" if descr else f"Unit: ({unit})"
)
else:
param_descr = descr
signal_base = SignalR if param_data["readonly"] else SignalRW
format = None
# infer format from parameter property
match param_data.get("_signal_format", None):
case "HINTED_SIGNAL":
format = Format.HINTED_SIGNAL
case "HINTED_UNCACHED_SIGNAL":
format = Format.HINTED_UNCACHED_SIGNAL
case "UNCACHED_SIGNAL":
format = Format.UNCACHED_SIGNAL
case _:
format = None
# depending on the Interface class other parameter need to be declared
# as readsignals as well
if param_name in secop_ophyd_modclass.hinted_signals:
format = format or Format.HINTED_SIGNAL
# Remove "StandardReadable" prefix from format for cleaner annotation
format = (
str(format).removeprefix("StandardReadable") if format else None
)
datainfo = param_data.get("datainfo", {})
# infer the ophyd type from secop datatype
type_param = get_type_param(param_data["datatype"])
# Handle StrictEnum types - generate enum class
if type_param and "StrictEnum" in type_param:
# Generate unique enum class name:
# ModuleClass + ParamName + Enum
param_name_list = (
param_name.replace(" ", "_").replace("-", "_").split("_")
)
param_name_camel = "".join(
word.capitalize() for word in param_name_list
)
enum_class_name = f"{module_class}_{param_name_camel}_Enum"
enum_cls = _build_enum_class(
enum_class_name,
datainfo.get("members", {}),
f"{param_name} enum for `{module_class}`.",
)
if enum_cls:
module_enum_classes.append(enum_cls)
# Use the specific enum class name instead of generic
# StrictEnum
type_param = enum_class_name
# Default format for parameters is CONFIG_SIGNAL
mod_parameters.append(
ParameterAttribute(
name=param_name,
type=signal_base.__name__,
type_param=type_param,
description=param_descr,
path_annotation=str(ParameterType()),
format_annotation=format,
)
)
# Module properties
module_properties: list[PropertyAttribute] = []
# Process module properties
for prop_name, property_value in properties.items():
if prop_name in IGNORED_PROPS:
continue
type_param = get_type_param(secop_dtype_obj_from_json(property_value))
module_properties.append(
PropertyAttribute(
name=prop_name,
type=SignalR.__name__,
type_param=type_param,
path_annotation=str(PropertyType()),
)
)
self.add_mod_class(
module_cls=module_class,
bases=module_bases,
parameters=mod_parameters,
properties=module_properties,
cmd_plans=command_plans,
description=properties.get("description", ""),
enum_classes=module_enum_classes,
commands=mod_commands,
)
# Add to node attributes
# Type the None explicitly as str | None to match other entries
node_module_attrs.append(ModuleAttribute(name=modname, type=module_class))
# Process module properties
for prop_name, property_value in node_properties.items():
type_param = get_type_param(secop_dtype_obj_from_json(property_value))
# Generate PropPath annotation for node-level properties
node_property_attrs.append(
PropertyAttribute(
name=prop_name,
type=SignalR.__name__,
type_param=type_param,
path_annotation=str(PropertyType()),
)
)
# Add node class
node_bases = ["SECoPNodeDevice"]
equipment_id: str = node_properties["equipment_id"]
# format node class accordingly
node_class_name = equipment_id.replace(".", "_").replace("-", "_").capitalize()
self.add_node_class(
node_cls=node_class_name,
bases=node_bases,
modules=node_module_attrs,
properties=node_property_attrs,
description=node_properties.get("description", ""),
)
# Add required imports
self.add_import("secop_ophyd.SECoPDevices", "SECoPNodeDevice")
self.add_import("secop_ophyd.SECoPDevices", "SECoPBaseDevice")
self.add_import("secop_ophyd.SECoPDevices", "SECoPCommunicatorDevice")
self.add_import("secop_ophyd.SECoPDevices", "SECoPReadableDevice")
self.add_import("secop_ophyd.SECoPDevices", "SECoPWritableDevice")
self.add_import("secop_ophyd.SECoPDevices", "SECoPMoveableDevice")
self.add_import("secop_ophyd.SECoPDevices", "SECoPTriggerableDevice")
[docs]
def generate_code(self) -> str:
"""Generate Python code using Jinja2 template.
Returns:
Generated Python code as string
"""
template = self.jinja_env.get_template("generated_classes.py.jinja2")
# Prepare template context
context = {
"imports": self.imports,
"module_classes": self.module_classes,
"node_classes": self.node_classes,
"enum_classes": self._collect_all_enums(),
"inline_comment_threshold": self.inline_comment_threshold,
"comment_wrap_width": self.comment_wrap_width,
}
# Render template
code = template.render(**context)
# Remove unused imports with autoflake
try:
code = autoflake.fix_code(
code,
remove_all_unused_imports=True,
remove_unused_variables=False,
remove_duplicate_keys=True,
)
except Exception as e:
if self.log:
self.log.warning(f"Autoflake processing failed: {e}")
else:
print(f"Warning: Autoflake processing failed: {e}")
# Format with Black
try:
code = black.format_str(code, mode=black.Mode(line_length=200))
except Exception as e:
if self.log:
self.log.warning(f"Black formatting failed: {e}")
else:
print(f"Warning: Black formatting failed: {e}")
return code
def _collect_all_enums(self) -> list[EnumClass]:
"""Collect and merge enum definitions from all module classes.
When multiple module classes have enums with the same base name but different
members, they are merged into a single SupersetEnum containing the union of all
members.
Returns:
List of deduplicated EnumClass definitions
"""
from collections import defaultdict
# Group enums by their base name (ModuleClass + ParamName + Enum)
# We need to track which module classes use each enum
enum_groups = defaultdict(list) # base_name -> [(module_class, enum)]
for mod_cls in self.module_classes:
for enum in mod_cls.enums:
# Extract base enum name by removing module class prefix
# e.g., "MassflowController1Gastype_selectEnum" -> need module class
# name
enum_groups[enum.name].append((mod_cls.name, enum))
# Process each enum group
merged_enums = []
for enum_name, enum_list in enum_groups.items():
if len(enum_list) == 1:
# Single enum definition - use StrictEnum
_, enum = enum_list[0]
# an already merged enum read in from a file
if enum.base_enum_class != "StrictEnum":
merged_enums.append(enum)
continue
enum.base_enum_class = "StrictEnum"
merged_enums.append(enum)
else:
# Multiple definitions - need to check if members are identical
member_sets = [
frozenset((m.name, m.value) for m in enum.members)
for _, enum in enum_list
]
if len(set(member_sets)) == 1:
# All enums have identical members - use StrictEnum
_, enum = enum_list[0]
enum.base_enum_class = "StrictEnum"
merged_enums.append(enum)
else:
# Different members - merge into SupersetEnum
all_members_dict = {} # (name, value) -> EnumMember
for _, enum in enum_list:
for member in enum.members:
key = (member.name, member.value)
if key not in all_members_dict:
all_members_dict[key] = member
# Create merged enum with all unique members
_, base_enum = enum_list[0]
merged_enum = EnumClass(
name=enum_name,
members=list(all_members_dict.values()),
description=base_enum.description,
base_enum_class="SupersetEnum",
)
merged_enums.append(merged_enum)
return merged_enums
[docs]
def write_gen_node_class_file(self):
"""Generate and write one file for the generated SEC node, named after
its node class, overwriting any existing file of that name."""
if not self.node_classes:
raise ValueError(
"No node class has been added yet (call from_json_describe() "
"or add_node_class() first) -- cannot determine output filename."
)
code = self.generate_code()
self.module_folder_path.mkdir(parents=True, exist_ok=True)
filep = self.module_folder_path / f"{self.node_classes[0].name}.py"
with open(filep, "w") as file:
file.write(code)
self._write_package_init()
if self.log:
self.log.info(f"Generated class file: {filep}")
else:
print(f"Generated class file: {filep}")
def _write_package_init(self):
"""(Re)generate __init__.py so the output directory is an importable
package that re-exports every generated node class currently on disk,
e.g. ``from <output_dir> import <NodeClassName>``."""
node_names = sorted(
p.stem
for p in self.module_folder_path.glob("*.py")
if p.stem != "__init__" and p.stem.isidentifier()
)
lines = [f"from .{name} import {name}\n" for name in node_names]
init_file = self.module_folder_path / "__init__.py"
with open(init_file, "w") as file:
file.writelines(lines)
[docs]
def get_type_param(secop_dtype: DataType) -> str | None:
sig_type = SECoPdtype(secop_dtype).np_datatype
# Get the module name
module = sig_type.__module__
# For builtins, just return the name without module prefix
if module == "builtins":
return sig_type.__name__
return sig_type.__name__
[docs]
def get_type_prop(prop_value) -> str | None:
secop_dtype: DataType = secop_dtype_obj_from_json(prop_value)
return get_type_param(secop_dtype)
