More type fixes

SpiNNakerManchester · Nov 3, 2023 · 83799a2 · 83799a2
1 parent 8d9aa2a
commit 83799a2
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 7 deletions.
diff --git a/pacman/model/routing_info/__init__.py b/pacman/model/routing_info/__init__.py
@@ -16,6 +16,7 @@
 from .routing_info import RoutingInfo
 from .machine_vertex_routing_info import MachineVertexRoutingInfo
 from .app_vertex_routing_info import AppVertexRoutingInfo
+from .vertex_routing_info import VertexRoutingInfo
 
 __all__ = ["BaseKeyAndMask", "MachineVertexRoutingInfo",
-           "RoutingInfo", "AppVertexRoutingInfo"]
+           "RoutingInfo", "AppVertexRoutingInfo", "VertexRoutingInfo"]
diff --git a/pacman/model/routing_info/app_vertex_routing_info.py b/pacman/model/routing_info/app_vertex_routing_info.py
@@ -21,6 +21,8 @@
 if TYPE_CHECKING:
     from pacman.model.graphs.application import ApplicationVertex
     from pacman.model.routing_info import BaseKeyAndMask
+    from pacman.model.routing_table_by_partition import (
+        MulticastRoutingTableByPartitionEntry)
     from .machine_vertex_routing_info import MachineVertexRoutingInfo
 
 logger = logging.getLogger(__name__)
@@ -56,7 +58,7 @@ def __init__(
         self.__max_machine_index = max_machine_index
 
     def merge_machine_entries(self, entries: List[Tuple[
-            MulticastRoutingEntry,
+            MulticastRoutingTableByPartitionEntry,
             MachineVertexRoutingInfo]]) -> Iterable[MulticastRoutingEntry]:
         """
         Merge the machine entries.
@@ -65,7 +67,7 @@ def merge_machine_entries(self, entries: List[Tuple[
             The entries to merge
         :type entries:
             list(tuple(
-                ~spinn_machine.MulticastRoutingEntry, VertexRoutingInfo))
+                MulticastRoutingTableByPartitionEntry, VertexRoutingInfo))
         :rtype: iterable(~spinn_machine.MulticastRoutingEntry)
         """
         n_entries = len(entries)

diff --git a/pacman/operations/routing_table_generators/merged_routing_table_generator.py b/pacman/operations/routing_table_generators/merged_routing_table_generator.py
@@ -20,7 +20,8 @@
 from pacman.model.routing_tables import (
     UnCompressedMulticastRoutingTable, MulticastRoutingTables)
 from pacman.model.graphs.application import ApplicationVertex
-from pacman.model.routing_info import RoutingInfo, AppVertexRoutingInfo
+from pacman.model.routing_info import (
+    RoutingInfo, AppVertexRoutingInfo, MachineVertexRoutingInfo)
 from pacman.model.graphs import AbstractVertex
 from pacman.model.routing_table_by_partition import (
     MulticastRoutingTableByPartitionEntry)
@@ -83,19 +84,32 @@ def __create_routing_table(
             continue
         # Otherwise it has to be a machine vertex...
         assert isinstance(vertex, MachineVertex)
+        assert isinstance(r_info, MachineVertexRoutingInfo)
+
+        # If there is no application vertex, just add the entry
+        if vertex.app_vertex is None:
+            table.add_multicast_routing_entry(
+                MulticastRoutingEntry(
+                    r_info.key, r_info.mask, defaultable=entry.defaultable,
+                    spinnaker_route=entry.spinnaker_route))
+            continue
+
+        # This has to be AppVertexRoutingInfo!
         app_r_info = routing_info.get_routing_info_from_pre_vertex(
             vertex.app_vertex, part_id)
-        # And this has to be AppVertexRoutingInfo!
         assert isinstance(app_r_info, AppVertexRoutingInfo)
+
+        # Get the entries to merge
         entries: List[Tuple[
             MulticastRoutingTableByPartitionEntry,
-            AppVertexRoutingInfo]] = [(entry, r_info)]
+            MachineVertexRoutingInfo]] = [(entry, r_info)]
         while __match(iterator, vertex, part_id, r_info, entry, routing_info,
                       app_r_info):
             (vertex, part_id), entry = iterator.pop()
             r_info = routing_info.get_routing_info_from_pre_vertex(
                 vertex, part_id)
             if r_info is not None:
+                assert isinstance(r_info, MachineVertexRoutingInfo)
                 entries.append((entry, r_info))
 
         # Now attempt to merge sources together as much as possible
@@ -149,7 +163,7 @@ def __merged_keys_and_masks(
         app_r_info: AppVertexRoutingInfo,
         entries: List[Tuple[
             MulticastRoutingTableByPartitionEntry,
-            AppVertexRoutingInfo]]) -> Iterable[MulticastRoutingEntry]:
+            MachineVertexRoutingInfo]]) -> Iterable[MulticastRoutingEntry]:
     if not entries:
         return
     (entry, r_info) = entries[0]