have claude sonnet 4.5 add edge attributes

examples/link_attributes_example.py

@@ -0,0 +1,106 @@
+#!/usr/bin/env python3
+"""
+Example demonstrating arbitrary key/value attributes on links.
+
+This example shows how to add custom attributes to links in addition to
+the standard latency attribute. This provides flexibility for modeling
+various link properties such as bandwidth, protocol, reliability, etc.
+"""
+
+import sys
+sys.path.insert(0, '../src')
+
+from ahp_graph.Device import *
+from ahp_graph.DeviceGraph import *
+
+
+class NetworkSwitch(Device):
+    """A simple network switch device."""
+    library = 'network.Switch'
+    portinfo = PortInfo()
+    portinfo.add('port', limit=None, required=False)  # Multi-port
+
+    def __init__(self, name: str):
+        super().__init__(name, attr={'switch_type': '10G'})
+
+
+class Server(Device):
+    """A server device."""
+    library = 'compute.Server'
+    portinfo = PortInfo()
+    portinfo.add('nic')  # Network interface
+
+    def __init__(self, name: str):
+        super().__init__(name, attr={'cores': 32, 'memory': '256GB'})
+
+
+def main():
+    """Create a simple network topology with diverse link attributes."""
+    print("Creating network topology with link attributes...\n")
+
+    graph = DeviceGraph()
+
+    # Create devices
+    switch1 = NetworkSwitch('switch1')
+    switch2 = NetworkSwitch('switch2')
+    server1 = Server('server1')
+    server2 = Server('server2')
+
+    # Link servers to switches with different link properties
+    # Server 1 has a high-speed, low-latency connection
+    graph.link(
+        server1.nic,
+        switch1.port('port', 0),
+        attr={
+            'latency': '1ns',
+            'bandwidth': '100GB/s',
+            'protocol': 'InfiniBand',
+            'mtu': 9000,
+            'reliability': 0.99999
+        }
+    )
+
+    # Server 2 has a standard Ethernet connection
+    graph.link(
+        server2.nic,
+        switch1.port('port', 1),
+        attr={
+            'latency': '10ns',
+            'bandwidth': '10GB/s',
+            'protocol': 'Ethernet',
+            'mtu': 1500,
+            'reliability': 0.9999
+        }
+    )
+
+    # Inter-switch link with high bandwidth
+    graph.link(
+        switch1.port('port', 2),
+        switch2.port('port', 0),
+        attr={
+            'latency': '5ns',
+            'bandwidth': '400GB/s',
+            'protocol': 'Ethernet',
+            'trunk': True,
+            'vlan_ids': [100, 200, 300]
+        }
+    )
+
+    # Display the graph
+    print(graph)
+    print("\n" + "="*80 + "\n")
+
+    # Demonstrate accessing link attributes
+    print("Link Attributes Summary:")
+    print("-" * 80)
+    for (p0, p1), attr in graph.links.items():
+        print(f"\nLink: {p0} <--> {p1}")
+        for key, value in sorted(attr.items()):
+            print(f"  {key}: {value}")
+
+    print("\n" + "="*80)
+    print("Example complete! Link attributes provide rich metadata for modeling.")
+
+
+if __name__ == '__main__':
+    main()

pyproject.toml

@@ -22,3 +22,8 @@ classifiers = [
 
 [project.urls]
 "Homepage" = "https://github.com/lpsmodsim/ahp_graph"
+
+[tool.pytest.ini_options]
+pythonpath = [
+  ".", "src"
+]

src/ahp_graph/DeviceGraph.py

@@ -76,7 +76,9 @@ def __repr__(self) -> str:
         for device in self.devices.values():
             lines.append(str(device))
         for p0, p1 in self.links:
-            lines.append(f"{p0} <--{self.links[(p0, p1)]}--> {p1}")
+            attr = self.links[(p0, p1)]
+            attr_str = ', '.join(f"{k}={v}" for k, v in sorted(attr.items()))
+            lines.append(f"{p0} <--[{attr_str}]--> {p1}")
         return "\n".join(lines)
 
     def _link_other_port(self, p0: DevicePort, p1: DevicePort) -> None:
@@ -91,27 +93,47 @@ def _link_other_port(self, p0: DevicePort, p1: DevicePort) -> None:
             p1.link = p2
             self.ports.remove(p0)
             self.ports.add(p1)
-            latency = self.links.pop(_orderedtuple(p0, p2))
+            attr = self.links.pop(_orderedtuple(p0, p2))
             # add the other device to the graph
             if p1.device.name not in self.devices:
                 self.add(p1.device)
-            self.links[_orderedtuple(p1, p2)] = latency
+            self.links[_orderedtuple(p1, p2)] = attr
             if self.expand_new_links is not None:
                 self.expand_new_links.append((p1,p2))
 
     def link(self, p0: DevicePort, p1: DevicePort,
-             latency: str = '0s') -> None:
+             latency: str = None, attr: dict = None) -> None:
         """
-        Link two DevicePorts with latency if provided.
+        Link two DevicePorts with optional attributes.
 
         Links are bidirectional and the key is a frozenset of the two
         DevicePorts. Duplicate links (links between the same DevicePorts)
         are not permitted. Keep in mind that a unique DevicePort is created
         for each port number in a multi-port style port. If the link
         types to not match, then throw an exception. Devices that are linked
-        to will be added to the graph automatically.  Latency is expressed
-        as a string with time units (ps, ns, us...)
+        to will be added to the graph automatically.
+
+        Args:
+            p0: First DevicePort to link
+            p1: Second DevicePort to link
+            latency: Optional latency string (e.g., '0s', '10ns'). If provided,
+                     added to attr as 'latency' key. For backward compatibility.
+            attr: Optional dictionary of link attributes. If not provided,
+                  defaults to {'latency': '0s'} or {'latency': latency} if
+                  latency is specified.
         """
+        # Process attributes
+        if attr is None:
+            attr = {}
+        else:
+            attr = dict(attr)  # Make a copy to avoid modifying caller's dict
+
+        # Handle backward compatibility: if latency is provided, add it to attr
+        if latency is not None:
+            attr['latency'] = latency
+        elif 'latency' not in attr:
+            attr['latency'] = '0s'
+
         if callable(p0) or callable(p1):
             raise RuntimeError(f"{p0} or {p1} is callable. This probably means"
                                f" you have a multi port and didn't pick a port"
@@ -150,7 +172,7 @@ def link(self, p0: DevicePort, p1: DevicePort,
             p0.link = p1
             p1.link = p0
         key = _orderedtuple(p0, p1)
-        self.links[key] = latency
+        self.links[key] = attr
         if self.expand_new_links is not None:
             self.expand_new_links.append(key)
 

src/ahp_graph/SSTGraph.py

@@ -198,18 +198,19 @@ def recurseSubcomponents(dev: Device, comp: 'sst.Component') -> None:
                 recurseSubcomponents(d0, c0)
 
         # Second, link the component ports using graph links
-        for ((p0,p1),t) in self.links.items():
+        for ((p0,p1),attr) in self.links.items():
             if p0.device.library is not None \
                     and p1.device.library is not None:
                 c0 = n2c[p0.device.name]
                 c1 = n2c[p1.device.name]
                 s0 = p0.get_name()
                 s1 = p1.get_name()
+                latency_str = attr.get('latency', '0s')
                 if str(p0) < str(p1):
-                    link = sst.Link(f'{p0}__{t}__{p1}')
+                    link = sst.Link(f'{p0}__{latency_str}__{p1}')
                 else:
-                    link = sst.Link(f'{p1}__{t}__{p0}')
-                latency = t if t != '0s' else '1ps'
+                    link = sst.Link(f'{p1}__{latency_str}__{p0}')
+                latency = latency_str if latency_str != '0s' else '1ps'
                 link.connect((c0, s0, latency), (c1, s1, latency))
 
     def __write_model(self,
@@ -307,17 +308,18 @@ def recurseSubcomponents(dev: Device) -> list:
         # Now define the links between components.
         #
         links = list()
-        for ((p0,p1),t) in self.links.items():
+        for ((p0,p1),attr) in self.links.items():
             if p0.device.library is None:
                 raise RuntimeError(f"No SST library: {p0.device.name}")
             if p1.device.library is None:
                 raise RuntimeError(f"No SST library: {p1.device.name}")
 
-            latency = t if t != '0s' else '1ps'
+            latency_str = attr.get('latency', '0s')
+            latency = latency_str if latency_str != '0s' else '1ps'
             if str(p0) < str(p1):
-                name = f'{p0}__{t}__{p1}'
+                name = f'{p0}__{latency_str}__{p1}'
             else:
-                name = f'{p1}__{t}__{p0}'
+                name = f'{p1}__{latency_str}__{p0}'
 
             d0 = p0.device
             d1 = p1.device

tests/test_DeviceGraph.py

@@ -106,7 +106,11 @@ def test_link() -> None:
     graph.link(lptd.input, ptd0.optional)  # type: ignore[arg-type]
     assert len(graph.devices) == 3, 'linking add submodule parent'
     assert ltd in graph.devices, 'submodule parent included'
-    assert graph.links[frozenset({lptd.input, ptd0.optional})] == '0s', 'default latency'  # type: ignore[arg-type]
+    # Find the link key (links are keyed by ordered tuples)
+    link_key = [k for k in graph.links.keys() if lptd.input in k and ptd0.optional in k][0]
+    link_attr = graph.links[link_key]
+    assert isinstance(link_attr, dict), 'link attributes are dict'
+    assert link_attr['latency'] == '0s', 'default latency'
     linkAgain = None
     try:
         graph.link(ptd0.optional, lptd.input)  # type: ignore[arg-type]
@@ -136,7 +140,42 @@ def test_link() -> None:
         changeSinglePortLink = True
     assert changeSinglePortLink, 'linking from a single port again'
     graph.link(ptd0.limit(0), ptd1.limit(0), '123ns')  # type: ignore[operator]
-    assert graph.links[frozenset({ptd0.limit(0), ptd1.limit(0)})] == '123ns', 'latency'  # type: ignore[operator]
+    # Find the link key (links are keyed by ordered tuples)
+    link_key = [k for k in graph.links.keys() if ptd0.limit(0) in k and ptd1.limit(0) in k][0]
+    link_attr = graph.links[link_key]
+    assert link_attr['latency'] == '123ns', 'latency from positional param'
+
+
+def test_link_attributes() -> None:
+    """Test of arbitrary link attributes in a DeviceGraph."""
+    graph = DeviceGraph()
+
+    ptd0 = PortTestDevice('0')
+    ptd1 = PortTestDevice('1')
+
+    # Test link with attr dict including latency
+    graph.link(ptd0.default, ptd1.default, attr={'latency': '10ns', 'bandwidth': '100GB/s'})  # type: ignore[arg-type]
+    # Find the link key (links are keyed by ordered tuples of ports)
+    link_key = [k for k in graph.links.keys() if ptd0.default in k and ptd1.default in k][0]
+    link_attr = graph.links[link_key]
+    assert link_attr['latency'] == '10ns', 'latency in attr dict'
+    assert link_attr['bandwidth'] == '100GB/s', 'custom attribute'
+
+    # Test link with latency param and additional attr
+    graph.link(ptd0.ptype, ptd1.ptype, latency='5ns', attr={'bandwidth': '50GB/s', 'protocol': 'PCIe'})  # type: ignore[arg-type]
+    link_key = [k for k in graph.links.keys() if ptd0.ptype in k and ptd1.ptype in k][0]
+    link_attr = graph.links[link_key]
+    assert link_attr['latency'] == '5ns', 'latency param overrides attr dict'
+    assert link_attr['bandwidth'] == '50GB/s', 'custom attribute with latency param'
+    assert link_attr['protocol'] == 'PCIe', 'multiple custom attributes'
+
+    # Test link with only custom attributes (no latency specified)
+    graph.link(ptd0.no_limit(0), ptd1.no_limit(0), attr={'bandwidth': '200GB/s', 'protocol': 'NVLink'})  # type: ignore[operator]
+    link_key = [k for k in graph.links.keys() if ptd0.no_limit(0) in k and ptd1.no_limit(0) in k][0]
+    link_attr = graph.links[link_key]
+    assert link_attr['latency'] == '0s', 'default latency when not in attr'
+    assert link_attr['bandwidth'] == '200GB/s', 'custom attribute without latency'
+    assert link_attr['protocol'] == 'NVLink', 'protocol attribute'
 
 
 def test_verifyLinks() -> None: