make Iterable(s) Sized

SpiNNakerManchester · Oct 11, 2024 · fadcbe9 · fadcbe9
1 parent 0383253
commit fadcbe9
Show file tree

Hide file tree

Showing 6 changed files with 20 additions and 0 deletions.
diff --git a/pacman/model/routing_info/routing_info.py b/pacman/model/routing_info/routing_info.py
@@ -84,3 +84,6 @@ def __iter__(self) -> Iterator[VertexRoutingInfo]:
         :return: a iterator of routing information
         """
         return iter(self._info.values())
+
+    def __len__(self) -> int:
+        return len(self._info)
diff --git a/pacman/model/routing_tables/multicast_routing_tables.py b/pacman/model/routing_tables/multicast_routing_tables.py
@@ -135,6 +135,9 @@ def __iter__(self) -> Iterator[AbstractMulticastRoutingTable]:
         """
         return iter(self._routing_tables_by_chip.values())
 
+    def __len__(self) -> int:
+        return len(self._routing_tables_by_chip)
+
 
 def to_json(router_table: MulticastRoutingTables) -> JsonObjectArray:
     """

diff --git a/pacman/utilities/algorithm_utilities/routing_tree.py b/pacman/utilities/algorithm_utilities/routing_tree.py
@@ -146,6 +146,15 @@ def __iter__(self) -> Iterator[Union[RoutingTree, MachineVertex]]:
             else:
                 yield obj
 
+    def __len__(self) -> int:
+        count = 1  # self
+        for _route, obj in self._children:
+            if isinstance(obj, RoutingTree):
+                count += len(obj)
+            else:
+                count += 1
+        return count
+
     def __repr__(self) -> str:
         return f"<RoutingTree at {self.chip} with {len(self._children)}" \
                f" {'child' if len(self._children) == 1 else 'children'}>"

diff --git a/unittests/model_tests/routing_info_tests/test_routing_info.py b/unittests/model_tests/routing_info_tests/test_routing_info.py
@@ -73,6 +73,8 @@ def test_routing_info(self):
         assert routing_info.get_routing_info_from_pre_vertex(
             pre_vertex, "Test2").get_keys().tolist() == [key]
 
+        self.assertEqual(len(routing_info), len(list(routing_info)))
+
     def test_base_key_and_mask(self):
         with self.assertRaises(PacmanConfigurationException):
             BaseKeyAndMask(0xF0, 0x40)

diff --git a/unittests/model_tests/routing_table_tests/test_routing_tables_model.py b/unittests/model_tests/routing_table_tests/test_routing_tables_model.py
@@ -146,6 +146,7 @@ def test_new_multicast_routing_tables(self):
         tables = MulticastRoutingTables(mrt)
         retrieved_tables = tables.routing_tables
         self.assertEqual(len(retrieved_tables), len(mrt))
+        self.assertEqual(len(tables), len(mrt))
         for tab in retrieved_tables:
             self.assertIn(tab, mrt)
 

diff --git a/unittests/utilities_tests/test_routing_tree.py b/unittests/utilities_tests/test_routing_tree.py
@@ -37,6 +37,7 @@ def test_call(self):
                              list(rt1.children))
         self.assertListEqual([rt1, "m_vertexA", "m_vertexB"],
                              list(iter(rt1)))
+        self.assertEqual(len(rt1), len(list(rt1)))
 
         rt2 = RoutingTree((4, 5), "foo")
         self.assertEqual("foo", rt2.label)
@@ -56,6 +57,7 @@ def test_call(self):
             [(None, (4, 5), {2, 3}), (3, (3, 4), {1, 2})],
             list(rt2.traverse())
         )
+        self.assertEqual(len(rt2), len(list(rt2)))
 
 
 if __name__ == '__main__':