refactor(contact_player.py): part2: improve type hints, add docstrings

install_symlinks.sh

@@ -25,4 +25,6 @@ for script in tools/helpers/*; do
     ln -sf $(pwd)/$script $TARGET/$(basename $script)
 done
 
+echo "export PYTHONPATH=\"$(pwd):\$PYTHONPATH\"" >> .venv/bin/activate
+
 echo "Done"

tools/contact_player/ccp.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from enum import Enum
-from typing import TYPE_CHECKING, Dict, List, Tuple, Optional
+from typing import override
 
 
 class ContactState(Enum):
@@ -15,30 +15,31 @@ class ContactState(Enum):
 class CoreContact(object):
     def __init__(
         self,
-        timespan: Tuple[int, int],
-        nodes: Tuple[str, str],
+        timespan: tuple[int, int],
+        nodes: tuple[str, str],
         bw: str,
         loss: float,
         delay: float,
         jitter: float,
         symmetric: bool = False,
     ) -> None:
-        self.timespan = timespan
-        self.nodes = nodes
-        self.bw = bw
-        self.loss = loss
-        self.delay = delay
-        self.jitter = jitter
-        self.symmetric = symmetric
-
+        self.timespan: tuple[int, int] = timespan
+        self.nodes: tuple[str, str] = nodes
+        self.bw: str = bw
+        self.loss: float = loss
+        self.delay: float = delay
+        self.jitter: float = jitter
+        self.symmetric: bool = symmetric
+
+    @override
     def __str__(self) -> str:
         return (
             f"CoreContact(timespan={self.timespan}, nodes={self.nodes}, bw={self.bw}, loss={self.loss}, "
             + f"delay={self.delay}, jitter={self.jitter}, symmetric={self.symmetric})"
         )
 
     @classmethod
-    def from_string(cls, line: str, mapping: dict[str, str] = {}) -> "CoreContact":
+    def from_string(cls, line: str, mapping: dict[str, str]) -> "CoreContact":
         line = line.strip()
         fixed_link = False
         if line.startswith("a contact"):
@@ -88,31 +89,29 @@ class CoreContactPlan(object):
 
     def __init__(
         self,
-        filename: str = None,
-        contacts: Dict[CoreContact, ContactState] = {},
-        fixed: List[CoreContact] = [],
-        mapping: Dict[int, str] = {},
+        filename: str | None = None,
+        contacts: dict[CoreContact, ContactState] = {},
+        fixed: list[CoreContact] = [],
+        mapping: dict[str, str] = {},
     ) -> None:
-        self.loop = False
-        self.contacts = contacts
-        self.fixed = fixed
+        self.loop: bool = False
+        self.contacts: dict[CoreContact, ContactState] = contacts
+        self.fixed: list[CoreContact] = fixed
         if filename:
             self.load(filename, mapping=mapping)
 
     @classmethod
-    def from_file(cls, filename, mapping: Dict[int, str] = {}) -> CoreContactPlan:
+    def from_file(cls, filename: str, mapping: dict[str, str] = {}) -> CoreContactPlan:
         plan = cls(filename, mapping=mapping)
         return plan
 
+    @override
     def __str__(self) -> str:
-        return "CoreContactPlan(loop=%r, #contacts=%d)" % (
-            self.loop,
-            len(self.contacts),
-        )
+        return f"CoreContactPlan(loop={self.loop}, #contacts={len(self.contacts)})"
 
-    def load(self, filename: str, mapping: Dict[int, str] = {}) -> None:
-        contacts = {}
-        fixed = []
+    def load(self, filename: str, mapping: dict[str, str] = {}) -> None:
+        contacts: dict[CoreContact, ContactState] = {}
+        fixed: list[CoreContact] = []
         with open(filename, "r") as f:
             for line in f:
                 line = line.strip()
@@ -136,9 +135,8 @@ def load(self, filename: str, mapping: Dict[int, str] = {}) -> None:
         self.contacts = contacts
         self.fixed = fixed
 
-    def at(self, time: int) -> List[Tuple[CoreContact, ContactState]]:
+    def at(self, time: int) -> list[tuple[CoreContact, ContactState]]:
         """Returns the list of contacts at the given time."""
-        orig = time
         if self.loop:
             time = time % self.get_max_time()
         return [
@@ -147,20 +145,20 @@ def at(self, time: int) -> List[Tuple[CoreContact, ContactState]]:
             if c.timespan[0] <= time and c.timespan[1] >= time
         ]
 
-    def need_activation(self, time: int) -> List[Tuple[CoreContact, ContactState]]:
+    def need_activation(self, time: int) -> list[tuple[CoreContact, ContactState]]:
         """Returns the list of contacts at the given time that need to be activated."""
         all = self.at(time)
         return [(c, s) for (c, s) in all if s == ContactState.PRE]
 
-    def need_deactivation(self, time: int) -> List[Tuple[CoreContact, ContactState]]:
+    def need_deactivation(self, time: int) -> list[tuple[CoreContact, ContactState]]:
         """Returns the list of contacts at the given time that need to be deactivated."""
         return [
             (c, s)
             for c, s in self.contacts.items()
             if time >= c.timespan[1] and s == ContactState.LIVE
         ]
 
-    def next_activation(self, time: int) -> Optional[int]:
+    def next_activation(self, time: int) -> int | None:
         """Returns the next activation time."""
         activations = [
             c.timespan[0]
@@ -171,7 +169,7 @@ def next_activation(self, time: int) -> Optional[int]:
             return None
         return min(activations)
 
-    def next_deactivation(self, time: int) -> Optional[int]:
+    def next_deactivation(self, time: int) -> int | None:
         """Returns the next deactivation time."""
         deactivations = [
             c.timespan[1]
@@ -191,13 +189,14 @@ def get_max_time(self) -> int:
         """Returns the maximum time in the contact plan."""
         return max([c.timespan[1] for c in self.contacts])
 
-    def has_contact(self, simtime: float, node1: str, node2: str) -> bool:
+    # TODO: unused.. remove..?
+    def has_contact(self, simtime: int, node1: str, node2: str) -> bool:
         current_contacts = self.at(simtime)
         # print("[ %f ] has_contact: %d %d | %s" % (simtime, node1, node2, current_contacts[0]))
         for c in current_contacts:
-            if c.nodes[0] == node1 and c.nodes[1] == node2:
+            if c[0].nodes[0] == node1 and c[0].nodes[1] == node2:
                 return True
-            if c.nodes[0] == node2 and c.nodes[1] == node1:
+            if c[0].nodes[0] == node2 and c[0].nodes[1] == node1:
                 return True
         return False
 

tools/contact_player/contact_player.py

@@ -1,35 +1,52 @@
 #!/usr/bin/env python3
 
 
-from tc_netem import *
-from ccp import *
 import argparse
-import time
+import os
 import signal
+import socket
 import sys
-import os
+import time
+from itertools import combinations
+from pathlib import Path
+from typing import Any, TypedDict, cast
+
 import yaml
-import socket
+
+from tools.contact_player.ccp import ContactState, CoreContact, CoreContactPlan
+from tools.contact_player.tc_netem import run_in_container, set_on_interface
+
+
+class Service(TypedDict):
+    environment: list[str]
+    networks: dict[str, dict[str, str]]
 
 
-def load_scenario(path):
+class Node(TypedDict):
+    eid: str
+    name: str
+    networks: dict[str, bool]
+    IPs: dict[str, str]
+
+
+def load_scenario(path: str | Path) -> dict[str, Node]:
     """
     Loads the docker compose scenario from the passed filepath.
     """
     print(f"Loading scenario from {path}.")
-    nodes: dict[str, dict] = {}
+    nodes: dict[str, Node] = {}
 
     with open(path) as f:
-        config = yaml.load(f, Loader=yaml.FullLoader)
+        config: dict[str, Any] = yaml.load(f, Loader=yaml.FullLoader)
         if "x-description" in config:
             print(f"Description: {config['x-description']}")
 
-        services: dict[str, dict[str, list[str]]] = config["services"]
+        services = cast(dict[str, Service], config["services"])
         for name, item in services.items():
             env_vars: list[str] = item["environment"]
             node_id = next(var for var in env_vars if var.startswith("NODE_ID"))
             node_eID = f"ipn:{node_id.split('=')[1]}.0"
-            new_node: dict[str, str] = {
+            new_node: Node = {
                 "eid": node_eID,
                 "name": name,
                 "networks": {},
@@ -50,22 +67,20 @@ def load_scenario(path):
 
 
 def find_common_subnet_between_nodes(
-    node1: str, node2: str, nodes: dict[str, dict[str, dict[str, str]]]
+    node1: str, node2: str, nodes: dict[str, Node]
 ) -> str | None:
     for k in nodes[node1].keys():
         if k in nodes[node2]:
             return k
     return None
 
 
-def get_dev_for_subnet(
-    node: str, subnet: str, nodes: dict[str, dict[str, dict[str, str]]]
-) -> str:
+def get_dev_for_subnet(node: str, subnet: str, nodes: dict[str, Node]) -> str:
     return nodes[node][subnet]["dev"]
 
 
 def get_network_for_interface(
-    node: str, interface: str, nodes: dict[str, dict[str, dict[str, str]]]
+    node: str, interface: str, nodes: dict[str, Node]
 ) -> str | None:
     for network, net_conf in nodes[node].items():
         if net_conf["dev"] == interface:
@@ -77,7 +92,7 @@ def get_network_for_interface(
 
 
 def contact_to_node_iface(
-    contact: CoreContact, nodes: dict[str, dict[str, dict[str, str]]]
+    contact: CoreContact, nodes: dict[str, Node]
 ) -> list[tuple[str, str]]:
     """
     Resolve a contact into the list of (node, interface) tuples, taking (a)symmetry of the contact into account.
@@ -130,10 +145,10 @@ def contact_to_node_iface(
 
 
 def set_link(
-    nodes: dict[str, dict[str, dict[str, str]]],
+    nodes: dict[str, Node],
     contact: CoreContact,
-    deactivate=False,
-    command="change",
+    deactivate: bool = False,
+    command: str = "change",
 ):
     loss = contact.loss
     if deactivate:
@@ -153,16 +168,29 @@ def set_link(
         )
 
 
-def get_pure_node_links(links: list) -> set:
-    pure_node_links = set()
+def get_pure_node_links(links: list[tuple[str, str, str]]) -> set[tuple[str, str, str]]:
+    """Resolves interface identifiers to node names and deduplicates links.
+
+    Converts links that reference specific network interfaces (e.g. ``dev:pcc_gs1``)
+    into plain node-to-node links, then deduplicates by sorting each pair
+    so that ``(pcc, gs1)`` and ``(gs1, pcc)`` are treated as the same link.
+
+    Args:
+        links: Raw link tuples of the form ``(endpoint_a, endpoint_b, link_type)``,
+            where either endpoint may be a ``dev:<node>_<node>`` interface reference.
+
+    Returns:
+        A set of deduplicated 3-tuples ``(node_a, node_b, link_type)`` with
+        node names sorted alphabetically and all interface references resolved.
+    """
+    pure_node_links: set[tuple[str, str, str]] = set()
     for l in links:
         # print("Link: ", l)
         nodes = [l[0], l[1]]
         link_type = l[2]
 
         if l[0].startswith("dev:"):
             dev_str = l[0].split(":")[1]
-            other_node = l[1]
             components = dev_str.split("_")
             if len(components) >= 2:
                 if components[0] == l[1]:
@@ -175,7 +203,6 @@ def get_pure_node_links(links: list) -> set:
                 )
         if l[1].startswith("dev:"):
             dev_str = l[1].split(":")[1]
-            other_node = l[0]
             components = dev_str.split("_")
             if len(components) >= 2:
                 if components[0] == l[0]:
@@ -192,7 +219,16 @@ def get_pure_node_links(links: list) -> set:
     return pure_node_links
 
 
-def update_netmap(netmap: bool, scenario_name: str, links: list):
+def update_netmap(
+    netmap: bool, scenario_name: str, links: list[tuple[str, str, str]]
+) -> None:
+    """Writes or updates resolved node links to a .netmap file in the tmp/ directory.
+
+    Args:
+        netmap: When False, this function does nothing.
+        scenario_name: Used as the output filename (tmp/<scenario_name>.netmap).
+        links: Raw link tuples to resolve and write.
+    """
     if netmap:
         # check if tmp directory exists
         if not os.path.exists("tmp"):
@@ -224,8 +260,8 @@ def main() -> None:
     netmap = args.map_network
 
     mapping = {}
-    nodes: dict[str, dict[str, dict[str, str]]] = {}
-    links = []
+    nodes: dict[str, Node] = {}
+    links: list[tuple[str, str, str]] = []
 
     for k, v in scenario.items():
         # extract node number from key
@@ -242,21 +278,11 @@ def main() -> None:
 
     # check all node combinations for common subnets/links
 
-    for n1 in nodes.keys():
-        for n2 in nodes.keys():
-            if n1 == n2:
-                continue
-            link = find_common_subnet_between_nodes(n1, n2, nodes)
-            if link is not None:
-                # sort n1 and n2 to avoid duplicates
-                l = sorted([n1, n2])
-                l.append("-")
-                links.append(l)
-    links = list(set([tuple(l) for l in links]))
-
-    # print(mapping)
-    # print(nodes)
-    # print(links)
+    for n1, n2 in combinations(nodes, 2):
+        if find_common_subnet_between_nodes(n1, n2, nodes) is None:
+            continue
+        a, b = sorted((n1, n2))
+        links.append((a, b, "-"))
 
     scenario_name = os.path.basename(args.scenario)
     scenario_name = os.path.splitext(scenario_name)[0]