Make TableCollections iterable

python/tests/test_tables.py

@@ -30,6 +30,7 @@
 import string
 import unittest
 import warnings
+import itertools
 
 import numpy as np
 
@@ -1497,6 +1498,17 @@ def test_asdict(self):
         self.assertEqual(set(d1.keys()), set(d2.keys()))
         # TODO test the fromdict constructor
 
+    def test_iter(self):
+        def test_iter(table_collection):
+            table_names = [
+                attr_name for attr_name in sorted(dir(table_collection))
+                if isinstance(getattr(table_collection, attr_name), tskit.BaseTable)]
+            for n in table_names:
+                yield n, getattr(table_collection, n)
+        ts = msprime.simulate(10, mutation_rate=1, random_seed=1)
+        for t1, t2 in itertools.zip_longest(test_iter(ts.tables), ts.tables):
+            self.assertEquals(t1, t2)
+
     def test_equals_empty(self):
         self.assertEqual(tskit.TableCollection(), tskit.TableCollection())
 

python/tskit/tables.py

@@ -1505,6 +1505,19 @@ def __banner(self, title):
         title_line += "#"
         return line + "\n" + title_line + "\n" + line + "\n"
 
+    def __iter__(self):
+        """
+        Iterate over all the tables in this TableCollection, ordered by table name
+        (i.e. deterministically), returning a tuple of (table_name, table_object)
+        """
+        yield 'edges', self.edges
+        yield 'individuals', self.individuals
+        yield 'migrations', self.migrations
+        yield 'mutations', self.mutations
+        yield 'nodes', self.nodes
+        yield 'populations', self.populations
+        yield 'provenances', self.provenances
+
     def __str__(self):
         s = self.__banner("Individuals")
         s += str(self.individuals) + "\n"