from __future__ import annotations
import networkx as nx
from collections.abc import Iterable, Iterator, Collection
from typing import Any, TypeVar
from collections import deque
from qcodes.instrument import InstrumentBase
from semi_cr.core.lab.station.routing_v1.errors import (
AmbiguousRouteError,
NoRouteError,
PinNotInGraphError,
)
from semi_cr.core.lab.station.routing_v1.graph import (
build_routing_graph,
)
from semi_cr.core.lab.devices.pinned import PinnedDevice
from semi_cr.core.lab.station.routing_v1.models import (
Route,
RouteRequest,
)
T = TypeVar("T")
[docs]
def node_supports_role(
data: dict[str, Any],
role: str,
quantity: str | None = None,
) -> bool:
"""
Determine whether a graph node supplies the requested capability.
The new ``capabilities`` representation is authoritative whenever
present. The legacy ``role`` attribute is only consulted when the
node has no capabilities attribute.
"""
capabilities = data.get("capabilities")
# =====================================================
# NEW MODEL
#
# If capabilities exist, they are authoritative.
# NEVER fall back to legacy role matching.
# =====================================================
if capabilities is not None:
capabilities = set(capabilities)
# Direct capability query:
#
# role="voltage_source"
# role="current_meter"
if role in capabilities:
return True
# Specific query:
#
# role="source", quantity="voltage"
# -> voltage_source
#
# role="source", quantity="current"
# -> current_source
if quantity is not None:
required_capability = f"{quantity}_{role}"
return required_capability in capabilities
# Generic query:
#
# role="source" -> any *_source
# role="meter" -> any *_meter
if role in {"source", "meter"}:
suffix = f"_{role}"
return any(
capability.endswith(suffix)
for capability in capabilities
)
return False
# =====================================================
# LEGACY MODEL
#
# Only use this if the node has NO capabilities field.
# =====================================================
legacy_role = data.get("role")
if isinstance(legacy_role, str):
return legacy_role == role
if isinstance(
legacy_role,
(list, tuple, set, frozenset),
):
return role in legacy_role
return False
[docs]
def node_matches_filters(
data: dict[str, Any],
filters: dict[str, Any],
) -> bool:
"""
Check node attributes against routing filters.
Singular query ``quantity`` is matched against the newer plural
node attribute ``quantities`` when available.
"""
for key, expected in filters.items():
# -----------------------------------------
# New plural collection representation
# -----------------------------------------
plural_key = {
"quantity": "quantities",
}.get(key)
if plural_key is not None and plural_key in data:
actual = data[plural_key]
if isinstance(
actual,
(set, frozenset, list, tuple),
):
if expected not in actual:
return False
elif actual != expected:
return False
continue
# -----------------------------------------
# Legacy/current scalar representation
# -----------------------------------------
actual = data.get(key)
if actual is None:
return False
if isinstance(
actual,
(set, frozenset, list, tuple),
):
if expected not in actual:
return False
elif actual != expected:
return False
return True
def _find_graph_node_for_object(
graph: nx.MultiDiGraph,
obj: object,
) -> Any:
"""Find the graph node whose `object` attribute is `obj`."""
for node, data in graph.nodes(data=True):
if data.get("obj") is obj:
return node
raise ValueError(
f"Object {obj!r} is not present in the context graph."
)
def _reachable_neighbors(
graph: nx.MultiDiGraph,
node: Any,
) -> Iterable[tuple[Any, dict]]:
"""
Yield neighboring nodes and edge attributes in both directions.
This makes traversal independent of whether a relation such as
`contains` or `represents` happens to be stored parent->child or
child->parent.
"""
# Outgoing edges
for _, neighbor, _, data in graph.out_edges(
node,
keys=True,
data=True,
):
yield neighbor, data
# Incoming edges
for neighbor, _, _, data in graph.in_edges(
node,
keys=True,
data=True,
):
yield neighbor, data
[docs]
def find_device_pin_node(
graph: nx.MultiDiGraph,
device,
pin_name: str,
) -> str:
"""
Find the graph node representing one of a device's pins.
Object identity is used deliberately. Two different pin objects that
happen to compare equal must not be treated as the same physical pin.
"""
pin = getattr(device, pin_name)
for node, data in graph.nodes(data=True):
if data.get("obj") is pin:
return node
raise PinNotInGraphError(
f"No graph node found for "
f"{getattr(device, 'short_name', device)!r}.{pin_name}."
)
[docs]
def find_pin_node(
graph: nx.MultiDiGraph,
pin: object,
) -> str:
"""Find the graph node representing a specific pin object."""
for node_id, data in graph.nodes(data=True):
if data.get("obj") is pin:
return node_id
pin_name = getattr(
pin,
"short_name",
repr(pin),
)
raise PinNotInGraphError(
f"No graph node found for pin {pin_name!r}."
)
# %%
[docs]
def electrical_net_for_pin(
graph: nx.MultiDiGraph,
device: object,
pin_name: str,
modality: str = "dc",
) -> set[str]:
pin_id = find_device_pin_node(
graph,
device,
pin_name,
)
routing_graph = build_routing_graph(
graph,
modality=modality,
)
if pin_id not in routing_graph:
raise PinNotInGraphError(
f"Pin node {pin_id!r} is not present "
"in the routing graph."
)
return set(
nx.node_connected_component(
routing_graph,
pin_id,
)
)
# %%
[docs]
def find_pin_route_to_role(
graph: nx.MultiDiGraph,
device: object,
pin_name: str,
role: str,
**filters: Any,
) -> list[str]:
"""
Compatibility wrapper around :func:`find_route`.
New code should construct a RouteRequest directly.
"""
pin = getattr(device, pin_name)
route = find_route(
graph,
RouteRequest(
pin=pin,
role=role,
quantity=filters.get("quantity"),
modality=filters.get("modality"),
),
)
return list(route.path)
# %%
def _format_request(
request: RouteRequest,
) -> str:
pin_name = getattr(
request.pin,
"short_name",
repr(request.pin),
)
parts = [
f"pin={pin_name!r}",
f"role={request.role!r}",
]
if request.quantity is not None:
parts.append(
f"quantity={request.quantity!r}"
)
if request.modality is not None:
parts.append(
f"modality={request.modality!r}"
)
return ", ".join(parts)
def _format_no_route_message(
request: RouteRequest,
) -> str:
return (
"No route found for "
f"{_format_request(request)}."
)
[docs]
def find_route(
graph: nx.MultiDiGraph,
request: RouteRequest,
) -> Route:
"""
Find an unambiguous route satisfying a routing request.
Raises
------
PinNotInGraphError
If the requested device pin is absent.
NoRouteError
If no matching reachable backend exists.
AmbiguousRouteError
If multiple different backend targets satisfy the request.
"""
pin_id = find_pin_node(
graph,
request.pin,
)
routing_graph = build_routing_graph(
graph,
modality=request.modality,
)
filters: dict[str, Any] = {}
if request.quantity is not None:
filters["quantity"] = request.quantity
if request.modality is not None:
filters["modality"] = request.modality
candidate_targets = [
node_id
for node_id, data in graph.nodes(data=True)
if node_supports_role(
data,
role=request.role,
quantity=request.quantity,
)
and node_matches_filters(
data,
filters,
)
]
reachable_routes: list[
tuple[str, list[str]]
] = []
for target_id in candidate_targets:
if target_id not in routing_graph:
continue
if not nx.has_path(
routing_graph,
pin_id,
target_id,
):
continue
path = nx.shortest_path(
routing_graph,
source=pin_id,
target=target_id,
)
reachable_routes.append(
(target_id, path)
)
if not reachable_routes:
raise NoRouteError(
_format_no_route_message(
request,
)
)
if len(reachable_routes) > 1:
targets = [
target
for target, _path in reachable_routes
]
raise AmbiguousRouteError(
"Multiple routing targets satisfy "
f"{_format_request(request)}: "
f"{targets!r}."
)
target_id, path = reachable_routes[0]
return Route(
request=request,
target_node=target_id,
path=tuple(path),
)
# %%
[docs]
def represented_objects_for_node(
graph: nx.MultiDiGraph,
node: str,
) -> Iterator[Any]:
"""
Yield objects attached through incoming or outgoing
``represents`` edges.
"""
seen_object_ids: set[int] = set()
for source, target, _key, data in graph.edges(
node,
keys=True,
data=True,
):
if data.get("kind") != "represents":
continue
other = (
target
if source == node
else source
)
obj = graph.nodes[other].get("obj")
if (
obj is not None
and id(obj) not in seen_object_ids
):
seen_object_ids.add(id(obj))
yield obj
for source, target, _key, data in graph.in_edges(
node,
keys=True,
data=True,
):
if data.get("kind") != "represents":
continue
other = (
source
if target == node
else target
)
obj = graph.nodes[other].get("obj")
if (
obj is not None
and id(obj) not in seen_object_ids
):
seen_object_ids.add(id(obj))
yield obj
# %%
# def objects_on_electrical_net(
# graph: nx.MultiDiGraph,
# device,
# pin_name: str,
# ) -> list[Any]:
# objects: list[Any] = []
# for node in electrical_net_for_pin(graph, device, pin_name):
# obj = graph.nodes[node].get("obj")
# if obj is not None:
# objects.append(obj)
# objects.extend(represented_objects_for_node(graph, node))
# return objects
[docs]
def objects_on_electrical_net(
graph: nx.MultiDiGraph,
device: object,
pin_name: str,
modality: str = "dc",
) -> list[Any]:
objects: list[Any] = []
seen_object_ids: set[int] = set()
def add_object(obj: Any) -> None:
if obj is None:
return
object_id = id(obj)
if object_id in seen_object_ids:
return
seen_object_ids.add(object_id)
objects.append(obj)
for node_id in electrical_net_for_pin(
graph,
device,
pin_name,
modality=modality,
):
add_object(
graph.nodes[node_id].get("obj")
)
for obj in represented_objects_for_node(
graph,
node_id,
):
add_object(obj)
return objects
# %%
[docs]
def objects_reachable_from_device(
graph: nx.MultiDiGraph,
device: PinnedDevice,
modality: str | None = None,
include_configurable: bool = False,
) -> list[Any]:
routing_graph = build_routing_graph(
graph,
modality=modality,
include_configurable=include_configurable,
)
pin_ids: set[str] = set()
reachable_node_ids: set[str] = set()
for pin in device.pins:
pin_id = find_pin_node(
graph,
pin,
)
pin_ids.add(pin_id)
if pin_id not in routing_graph:
continue
reachable_node_ids.update(
nx.node_connected_component(
routing_graph,
pin_id,
)
)
# The pins are traversal starting points, not objects
# reached from the device.
reachable_node_ids.difference_update(
pin_ids
)
objects: list[Any] = []
seen_object_ids: set[int] = set()
for node_id in reachable_node_ids:
obj = routing_graph.nodes[node_id].get("obj")
if obj is None:
continue
if id(obj) in seen_object_ids:
continue
seen_object_ids.add(id(obj))
objects.append(obj)
return objects
[docs]
def find_instruments(
graph: nx.MultiDiGraph,
instrument_type: type[InstrumentBase] | None = None,
) -> list[InstrumentBase]:
instruments = []
for _, data in graph.nodes(data=True):
obj = data.get("obj")
if not isinstance(obj, InstrumentBase):
continue
if instrument_type is not None and not isinstance(obj, instrument_type):
continue
instruments.append(obj)
return instruments
[docs]
def find_instruments_for_device(
graph: nx.MultiDiGraph,
device: PinnedDevice,
instrument_type: type[InstrumentBase] | None = None,
modality: str | None = None,
include_configurable: bool = False,
) -> list[InstrumentBase]:
reachable_objects = objects_reachable_from_device(
graph,
device,
modality=modality,
include_configurable=include_configurable,
)
instruments: list[InstrumentBase] = []
seen_instrument_ids: set[int] = set()
for obj in reachable_objects:
root = getattr(
obj,
"root_instrument",
None,
)
if isinstance(root, InstrumentBase):
instrument = root
elif isinstance(obj, InstrumentBase):
instrument = obj
else:
continue
if (
instrument_type is not None
and not isinstance(
instrument,
instrument_type,
)
):
continue
instrument_id = id(instrument)
if instrument_id in seen_instrument_ids:
continue
seen_instrument_ids.add(instrument_id)
instruments.append(instrument)
return instruments
[docs]
def resolve_synchronization_config(
device: Any,
method: str,
) -> Any | None:
"""
Find synchronization configuration belonging to a measurement method.
"""
profiles = getattr(device, "measurement_profiles", None)
if profiles is None:
raise AttributeError(
f"{device!r} has no measurement_profiles."
)
for profile in _iter_values(profiles):
methods = getattr(profile, "methods", None)
if methods is None:
continue
method_config = _mapping_get(methods, method)
if method_config is None:
continue
return getattr(
method_config,
"synchronization",
None,
)
raise LookupError(
f"Could not find measurement method {method!r} "
f"on device {device!r}."
)
def _mapping_get(container: Any, key: str) -> Any | None:
if isinstance(container, dict):
return container.get(key)
try:
return container[key]
except (KeyError, TypeError):
pass
return getattr(container, key, None)
def _iter_values(container: Any):
if isinstance(container, dict):
return container.values()
values = getattr(container, "values", None)
if callable(values):
return values()
return container
[docs]
def objects_in_pin_route_with_attr(
graph: nx.MultiDiGraph,
device,
pin_name: str,
attr_name: str,
) -> list[Any]:
found: list[Any] = []
for obj in objects_on_electrical_net(graph, device, pin_name):
if hasattr(obj, attr_name):
found.append(obj)
continue
if hasattr(obj, "parameters") and attr_name in obj.parameters:
found.append(obj)
return found