|
| 1 | +"""Helper classes and functions for MeasurementLink examples.""" |
| 2 | + |
| 3 | +import logging |
| 4 | +import pathlib |
| 5 | +import types |
| 6 | +from typing import Any, Callable, List, NamedTuple, Optional, Tuple, TypeVar, Union |
| 7 | + |
| 8 | +import click |
| 9 | +import grpc |
| 10 | + |
| 11 | +import ni_measurementlink_service as nims |
| 12 | +from ni_measurementlink_service import session_management |
| 13 | +from ni_measurementlink_service._internal.discovery_client import DiscoveryClient |
| 14 | +from ni_measurementlink_service._internal.stubs.ni.measurementlink.pinmap.v1 import ( |
| 15 | + pin_map_service_pb2, |
| 16 | + pin_map_service_pb2_grpc, |
| 17 | +) |
| 18 | +from ni_measurementlink_service.measurement.service import ( |
| 19 | + GrpcChannelPool, |
| 20 | + MeasurementService, |
| 21 | +) |
| 22 | + |
| 23 | + |
| 24 | +class ServiceOptions(NamedTuple): |
| 25 | + """Service options specified on the command line.""" |
| 26 | + |
| 27 | + use_grpc_device: bool = False |
| 28 | + grpc_device_address: str = "" |
| 29 | + |
| 30 | + use_simulation: bool = False |
| 31 | + |
| 32 | + |
| 33 | +def get_service_options(**kwargs: Any) -> ServiceOptions: |
| 34 | + """Get service options from keyword arguments.""" |
| 35 | + return ServiceOptions( |
| 36 | + use_grpc_device=kwargs.get("use_grpc_device", False), |
| 37 | + grpc_device_address=kwargs.get("grpc_device_address", ""), |
| 38 | + use_simulation=kwargs.get("use_simulation", False), |
| 39 | + ) |
| 40 | + |
| 41 | + |
| 42 | +class PinMapClient(object): |
| 43 | + """Class that communicates with the pin map service.""" |
| 44 | + |
| 45 | + def __init__(self, *, grpc_channel: grpc.Channel): |
| 46 | + """Initialize pin map client.""" |
| 47 | + self._client: pin_map_service_pb2_grpc.PinMapServiceStub = ( |
| 48 | + pin_map_service_pb2_grpc.PinMapServiceStub(grpc_channel) |
| 49 | + ) |
| 50 | + |
| 51 | + def update_pin_map(self, pin_map_path: str) -> str: |
| 52 | + """Update registered pin map contents. |
| 53 | +
|
| 54 | + Create and register a pin map if a pin map resource for the specified pin map id is not |
| 55 | + found. |
| 56 | +
|
| 57 | + Args: |
| 58 | + pin_map_path: The file path of the pin map to register as a pin map resource. |
| 59 | +
|
| 60 | + Returns: |
| 61 | + The resource id of the pin map that is registered to the pin map service. |
| 62 | + """ |
| 63 | + pin_map_path_obj = pathlib.Path(pin_map_path) |
| 64 | + # By convention, the pin map id is the .pinmap file path. |
| 65 | + request = pin_map_service_pb2.UpdatePinMapFromXmlRequest( |
| 66 | + pin_map_id=pin_map_path, pin_map_xml=pin_map_path_obj.read_text(encoding="utf-8") |
| 67 | + ) |
| 68 | + response: pin_map_service_pb2.PinMap = self._client.UpdatePinMapFromXml(request) |
| 69 | + return response.pin_map_id |
| 70 | + |
| 71 | + |
| 72 | +class GrpcChannelPoolHelper(GrpcChannelPool): |
| 73 | + """Class that manages gRPC channel lifetimes.""" |
| 74 | + |
| 75 | + def __init__(self) -> None: |
| 76 | + """Initialize the GrpcChannelPool object.""" |
| 77 | + super().__init__() |
| 78 | + self._discovery_client = DiscoveryClient() |
| 79 | + |
| 80 | + @property |
| 81 | + def pin_map_channel(self) -> grpc.Channel: |
| 82 | + """Return gRPC channel to pin map service.""" |
| 83 | + return self.get_channel( |
| 84 | + self._discovery_client.resolve_service( |
| 85 | + provided_interface="ni.measurementlink.pinmap.v1.PinMapService", |
| 86 | + service_class="ni.measurementlink.pinmap.v1.PinMapService", |
| 87 | + ).insecure_address |
| 88 | + ) |
| 89 | + |
| 90 | + @property |
| 91 | + def session_management_channel(self) -> grpc.Channel: |
| 92 | + """Return gRPC channel to session management service.""" |
| 93 | + return self.get_channel( |
| 94 | + self._discovery_client.resolve_service( |
| 95 | + provided_interface=session_management.GRPC_SERVICE_INTERFACE_NAME, |
| 96 | + service_class=session_management.GRPC_SERVICE_CLASS, |
| 97 | + ).insecure_address |
| 98 | + ) |
| 99 | + |
| 100 | + def get_grpc_device_channel(self, provided_interface: str) -> grpc.Channel: |
| 101 | + """Return gRPC channel to specified NI gRPC Device service. |
| 102 | +
|
| 103 | + Args: |
| 104 | + provided_interface (str): The gRPC Full Name of the service. |
| 105 | +
|
| 106 | + """ |
| 107 | + return self.get_channel( |
| 108 | + self._discovery_client.resolve_service( |
| 109 | + provided_interface=provided_interface, |
| 110 | + service_class="ni.measurementlink.v1.grpcdeviceserver", |
| 111 | + ).insecure_address |
| 112 | + ) |
| 113 | + |
| 114 | + |
| 115 | +class TestStandSupport(object): |
| 116 | + """Class that communicates with TestStand.""" |
| 117 | + |
| 118 | + def __init__(self, sequence_context: Any) -> None: |
| 119 | + """Initialize the TestStandSupport object. |
| 120 | +
|
| 121 | + Args: |
| 122 | + sequence_context: |
| 123 | + The SequenceContext COM object from the TestStand sequence execution. |
| 124 | + (Dynamically typed.) |
| 125 | + """ |
| 126 | + self._sequence_context = sequence_context |
| 127 | + |
| 128 | + def get_active_pin_map_id(self) -> str: |
| 129 | + """Get the active pin map id from the NI.MeasurementLink.PinMapId temporary global variable. |
| 130 | +
|
| 131 | + Returns: |
| 132 | + The resource id of the pin map that is registered to the pin map service. |
| 133 | + """ |
| 134 | + return self._sequence_context.Engine.TemporaryGlobals.GetValString( |
| 135 | + "NI.MeasurementLink.PinMapId", 0x0 |
| 136 | + ) |
| 137 | + |
| 138 | + def set_active_pin_map_id(self, pin_map_id: str) -> None: |
| 139 | + """Set the NI.MeasurementLink.PinMapId temporary global variable to the specified id. |
| 140 | +
|
| 141 | + Args: |
| 142 | + pin_map_id: |
| 143 | + The resource id of the pin map that is registered to the pin map service. |
| 144 | + """ |
| 145 | + self._sequence_context.Engine.TemporaryGlobals.SetValString( |
| 146 | + "NI.MeasurementLink.PinMapId", 0x1, pin_map_id |
| 147 | + ) |
| 148 | + |
| 149 | + def resolve_file_path(self, file_path: str) -> str: |
| 150 | + """Resolve the absolute path to a file using the TestStand search directories. |
| 151 | +
|
| 152 | + Args: |
| 153 | + file_path: |
| 154 | + An absolute or relative path to the file. If this is a relative path, this function |
| 155 | + searches the TestStand search directories for it. |
| 156 | +
|
| 157 | + Returns: |
| 158 | + The absolute path to the file. |
| 159 | + """ |
| 160 | + if pathlib.Path(file_path).is_absolute(): |
| 161 | + return file_path |
| 162 | + (_, absolute_path, _, _, user_canceled) = self._sequence_context.Engine.FindFileEx( |
| 163 | + fileToFind=file_path, |
| 164 | + absolutePath=None, |
| 165 | + srchDirType=None, |
| 166 | + searchDirectoryIndex=None, |
| 167 | + userCancelled=None, # Must match spelling used by TestStand |
| 168 | + searchContext=self._sequence_context.SequenceFile, |
| 169 | + ) |
| 170 | + if user_canceled: |
| 171 | + raise RuntimeError("File lookup canceled by user.") |
| 172 | + return absolute_path |
| 173 | + |
| 174 | + |
| 175 | +def configure_logging(verbosity: int) -> None: |
| 176 | + """Configure logging for this process.""" |
| 177 | + if verbosity > 1: |
| 178 | + level = logging.DEBUG |
| 179 | + elif verbosity == 1: |
| 180 | + level = logging.INFO |
| 181 | + else: |
| 182 | + level = logging.WARNING |
| 183 | + logging.basicConfig(format="%(asctime)s %(levelname)s: %(message)s", level=level) |
| 184 | + |
| 185 | + |
| 186 | +F = TypeVar("F", bound=Callable) |
| 187 | + |
| 188 | + |
| 189 | +def verbosity_option(func: F) -> F: |
| 190 | + """Decorator for --verbose command line option.""" |
| 191 | + return click.option( |
| 192 | + "-v", |
| 193 | + "--verbose", |
| 194 | + "verbosity", |
| 195 | + count=True, |
| 196 | + help="Enable verbose logging. Repeat to increase verbosity.", |
| 197 | + )(func) |
| 198 | + |
| 199 | + |
| 200 | +def grpc_device_options(func: F) -> F: |
| 201 | + """Decorator for NI gRPC Device Server command line options.""" |
| 202 | + use_grpc_device_option = click.option( |
| 203 | + "--use-grpc-device/--no-use-grpc-device", |
| 204 | + default=True, |
| 205 | + is_flag=True, |
| 206 | + help="Use the NI gRPC Device Server.", |
| 207 | + ) |
| 208 | + grpc_device_address_option = click.option( |
| 209 | + "--grpc-device-address", |
| 210 | + default="", |
| 211 | + help="NI gRPC Device Server address (e.g. localhost:31763). If unspecified, use the discovery service to resolve the address.", |
| 212 | + ) |
| 213 | + return grpc_device_address_option(use_grpc_device_option(func)) |
| 214 | + |
| 215 | + |
| 216 | +def use_simulation_option(default: bool) -> Callable[[F], F]: |
| 217 | + """Decorator for --use-simulation command line option.""" |
| 218 | + return click.option( |
| 219 | + "--use-simulation/--no-use-simulation", |
| 220 | + default=default, |
| 221 | + is_flag=True, |
| 222 | + help="Use simulated instruments.", |
| 223 | + ) |
| 224 | + |
| 225 | + |
| 226 | +def get_grpc_device_channel( |
| 227 | + measurement_service: MeasurementService, |
| 228 | + driver_module: types.ModuleType, |
| 229 | + service_options: ServiceOptions, |
| 230 | +) -> Optional[grpc.Channel]: |
| 231 | + """Returns driver specific grpc device channel.""" |
| 232 | + if service_options.use_grpc_device: |
| 233 | + if service_options.grpc_device_address: |
| 234 | + return measurement_service.channel_pool.get_channel(service_options.grpc_device_address) |
| 235 | + |
| 236 | + return measurement_service.get_channel( |
| 237 | + provided_interface=getattr(driver_module, "GRPC_SERVICE_INTERFACE_NAME"), |
| 238 | + service_class="ni.measurementlink.v1.grpcdeviceserver", |
| 239 | + ) |
| 240 | + return None |
| 241 | + |
| 242 | + |
| 243 | +def create_session_management_client( |
| 244 | + measurement_service: MeasurementService, |
| 245 | +) -> nims.session_management.Client: |
| 246 | + """Return created session management client.""" |
| 247 | + return nims.session_management.Client( |
| 248 | + grpc_channel=measurement_service.get_channel( |
| 249 | + provided_interface=nims.session_management.GRPC_SERVICE_INTERFACE_NAME, |
| 250 | + service_class=nims.session_management.GRPC_SERVICE_CLASS, |
| 251 | + ) |
| 252 | + ) |
| 253 | + |
| 254 | + |
| 255 | +def get_session_and_channel_for_pin( |
| 256 | + session_info: List[nims.session_management.SessionInformation], |
| 257 | + pin: str, |
| 258 | + site: Optional[int] = None, |
| 259 | +) -> Tuple[int, List[str]]: |
| 260 | + """Returns the session information based on the given pin names.""" |
| 261 | + session_and_channel_info = get_sessions_and_channels_for_pins( |
| 262 | + session_info=session_info, pins=[pin], site=site |
| 263 | + ) |
| 264 | + |
| 265 | + if len(session_and_channel_info) != 1: |
| 266 | + raise ValueError(f"Unsupported number of sessions for {pin}: {len(session_info)}") |
| 267 | + return session_and_channel_info[0] |
| 268 | + |
| 269 | + |
| 270 | +def get_sessions_and_channels_for_pins( |
| 271 | + session_info: List[nims.session_management.SessionInformation], |
| 272 | + pins: Union[str, List[str]], |
| 273 | + site: Optional[int] = None, |
| 274 | +) -> List[Tuple[int, List[str]]]: |
| 275 | + """Returns the session information based on the given pin names.""" |
| 276 | + pin_names = [pins] if isinstance(pins, str) else pins |
| 277 | + session_and_channel_info = [] |
| 278 | + for session_index, session_details in enumerate(session_info): |
| 279 | + channel_list = [ |
| 280 | + mapping.channel |
| 281 | + for mapping in session_details.channel_mappings |
| 282 | + if mapping.pin_or_relay_name in pin_names and (site is None or mapping.site == site) |
| 283 | + ] |
| 284 | + if len(channel_list) != 0: |
| 285 | + session_and_channel_info.append((session_index, channel_list)) |
| 286 | + |
| 287 | + if len(session_and_channel_info) == 0: |
| 288 | + raise KeyError(f"Pin(s) {pins} and site {site} not found") |
| 289 | + |
| 290 | + return session_and_channel_info |
0 commit comments