Add table append method

Author: benjeffery

docs/conf.py

@@ -311,6 +311,7 @@ def handle_item(fieldarg, content):
     # TODO these have been triaged here to make the docs compile, but we should
     # sort them out properly. https://github.com/tskit-dev/tskit/issues/336
     ("py:class", "array_like"),
+    ("py:class", "row-like"),
     ("py:class", "array-like"),
     ("py:class", "dtype=np.uint32"),
     ("py:class", "dtype=np.uint32."),

python/CHANGELOG.rst

@@ -9,6 +9,9 @@
 
 **Features**
 
+- Add `Table.append` method for adding rows from classes such as `SiteTableRow` and
+  `Site` (:user:`benjeffery`, :issue:`1111`, :pr:`1254`).
+
 - SVG visualization of a single tree allows all mutations on an edge to be plotted
   via the ``all_edge_mutations`` param (:user:`hyanwong`,:issue:`1253`, :pr:`1258`).
 

python/tests/conftest.py

@@ -128,7 +128,7 @@ def ts_fixture():
     for name, table in tables.name_map.items():
         if name != "provenances":
             table.metadata_schema = tskit.MetadataSchema({"codec": "json"})
-            metadatas = [f"n_{name}_{u}" for u in range(len(table))]
+            metadatas = [f'{{"foo":"n_{name}_{u}"}}' for u in range(len(table))]
             metadata, metadata_offset = tskit.pack_strings(metadatas)
             table.set_columns(
                 **{

python/tests/simplify.py

@@ -24,6 +24,7 @@
 Python implementation of the simplify algorithm.
 """
 import sys
+from dataclasses import replace
 
 import numpy as np
 import portion
@@ -158,12 +159,11 @@ def record_node(self, input_id, is_sample=False):
         flags &= ~tskit.NODE_IS_SAMPLE
         if is_sample:
             flags |= tskit.NODE_IS_SAMPLE
-        output_id = self.tables.nodes.add_row(
-            flags=flags,
-            time=node.time,
-            population=node.population,
-            metadata=node.metadata,
-            individual=node.individual,
+        output_id = self.tables.nodes.append(
+            replace(
+                node,
+                flags=flags,
+            )
         )
         self.node_id_map[input_id] = output_id
         return output_id
@@ -186,9 +186,7 @@ def flush_edges(self):
         num_edges = 0
         for child in sorted(self.edge_buffer.keys()):
             for edge in self.edge_buffer[child]:
-                self.tables.edges.add_row(
-                    edge.left, edge.right, edge.parent, edge.child
-                )
+                self.tables.edges.append(edge)
                 num_edges += 1
         self.edge_buffer.clear()
         return num_edges
@@ -413,19 +411,15 @@ def finalise_sites(self):
                         mapped_parent = -1
                         if mut.parent != -1:
                             mapped_parent = mutation_id_map[mut.parent]
-                        self.tables.mutations.add_row(
-                            site=len(self.tables.sites),
-                            node=self.mutation_node_map[mut.id],
-                            time=mut.time,
-                            parent=mapped_parent,
-                            derived_state=mut.derived_state,
-                            metadata=mut.metadata,
+                        self.tables.mutations.append(
+                            replace(
+                                mut,
+                                site=len(self.tables.sites),
+                                node=self.mutation_node_map[mut.id],
+                                parent=mapped_parent,
+                            )
                         )
-                self.tables.sites.add_row(
-                    position=site.position,
-                    ancestral_state=site.ancestral_state,
-                    metadata=site.metadata,
-                )
+                self.tables.sites.append(site)
 
     def finalise_references(self):
         input_populations = self.ts.tables.populations
@@ -455,17 +449,12 @@ def finalise_references(self):
         for input_id, count in enumerate(population_ref_count):
             if count > 0:
                 row = input_populations[input_id]
-                output_id = self.tables.populations.add_row(metadata=row.metadata)
+                output_id = self.tables.populations.append(row)
                 population_id_map[input_id] = output_id
         for input_id, count in enumerate(individual_ref_count):
             if count > 0:
                 row = input_individuals[input_id]
-                output_id = self.tables.individuals.add_row(
-                    flags=row.flags,
-                    location=row.location,
-                    parents=row.parents,
-                    metadata=row.metadata,
-                )
+                output_id = self.tables.individuals.append(row)
                 individual_id_map[input_id] = output_id
 
         # Remap the population ID references for nodes.
@@ -489,11 +478,11 @@ def finalise_references(self):
                     mapped_parents.append(-1)
                 else:
                     mapped_parents.append(individual_id_map[p])
-            self.tables.individuals.add_row(
-                flags=row.flags,
-                location=row.location,
-                parents=mapped_parents,
-                metadata=row.metadata,
+            self.tables.individuals.append(
+                replace(
+                    row,
+                    parents=mapped_parents,
+                )
             )
 
         # We don't support migrations for now. We'll need to remap these as well.
@@ -710,7 +699,7 @@ def flush_edges(self):
         num_edges = 0
         for child in sorted(self.edge_buffer.keys()):
             for edge in self.edge_buffer[child]:
-                self.table.add_row(edge.left, edge.right, edge.parent, edge.child)
+                self.table.append(edge)
                 num_edges += 1
         self.edge_buffer.clear()
         return num_edges

python/tests/test_highlevel.py

@@ -40,6 +40,7 @@
 import unittest
 import uuid as _uuid
 import warnings
+from dataclasses import replace
 
 import kastore
 import msprime
@@ -2854,15 +2855,14 @@ def verify_random_permutation(self, ts):
         inv_node_map = {v: k for k, v in node_map.items()}
         for j in range(ts.num_nodes):
             node = ts.node(inv_node_map[j])
-            other_tables.nodes.add_row(
-                flags=node.flags, time=node.time, population=node.population
-            )
+            other_tables.nodes.append(node)
         for e in ts.edges():
-            other_tables.edges.add_row(
-                left=e.left,
-                right=e.right,
-                parent=node_map[e.parent],
-                child=node_map[e.child],
+            other_tables.edges.append(
+                replace(
+                    e,
+                    parent=node_map[e.parent],
+                    child=node_map[e.child],
+                )
             )
         for _ in range(ts.num_populations):
             other_tables.populations.add_row()
@@ -3168,6 +3168,30 @@ def get_instances(self, n):
         ]
 
 
+class TestContainersAppend:
+    def test_containers_append(self, ts_fixture):
+        """
+        Test that the containers work with `Table.append`
+        """
+        tables = ts_fixture.dump_tables()
+        tables.clear(clear_provenance=True)
+        for table_name in [
+            "individuals",
+            "nodes",
+            "edges",
+            "migrations",
+            "sites",
+            "mutations",
+            "populations",
+            "provenances",
+        ]:
+            table = getattr(tables, table_name)
+            for i in range(len(getattr(ts_fixture.tables, table_name))):
+                table.append(getattr(ts_fixture, table_name[:-1])(i))
+        print(ts_fixture.tables, tables)
+        assert ts_fixture.tables == tables
+
+
 class TestTskitConversionOutput(unittest.TestCase):
     """
     Tests conversion output to ensure it is correct.

python/tests/test_parsimony.py

@@ -24,6 +24,7 @@
 """
 import io
 import itertools
+from dataclasses import replace
 
 import attr
 import Bio.Phylo.TreeConstruction
@@ -549,18 +550,20 @@ def verify(self, ts):
                 ancestral_state, mutations = self.do_map_mutations(
                     tree, G[site.id], alleles[site.id]
                 )
-                site_id = tables.sites.add_row(site.position, ancestral_state)
+                site_id = tables.sites.append(
+                    replace(site, ancestral_state=ancestral_state)
+                )
                 parent_offset = len(tables.mutations)
                 for mutation in mutations:
                     parent = mutation.parent
                     if parent != tskit.NULL:
                         parent += parent_offset
-                    tables.mutations.add_row(
-                        site_id,
-                        node=mutation.node,
-                        time=mutation.time,
-                        parent=parent,
-                        derived_state=mutation.derived_state,
+                    tables.mutations.append(
+                        replace(
+                            mutation,
+                            site=site_id,
+                            parent=parent,
+                        )
                     )
         other_ts = tables.tree_sequence()
         for h1, h2 in zip(
@@ -715,18 +718,20 @@ def verify(self, ts):
                 ancestral_state, mutations = self.do_map_mutations(
                     tree, G[site.id], alleles[site.id]
                 )
-                site_id = tables.sites.add_row(site.position, ancestral_state)
+                site_id = tables.sites.append(
+                    replace(site, ancestral_state=ancestral_state)
+                )
                 parent_offset = len(tables.mutations)
                 for m in mutations:
                     parent = m.parent
                     if m.parent != tskit.NULL:
                         parent = m.parent + parent_offset
-                    tables.mutations.add_row(
-                        site_id,
-                        node=m.node,
-                        time=m.time,
-                        parent=parent,
-                        derived_state=m.derived_state,
+                    tables.mutations.append(
+                        replace(
+                            m,
+                            site=site_id,
+                            parent=parent,
+                        )
                     )
         other_ts = tables.tree_sequence()
         assert ts.num_samples == other_ts.num_samples
@@ -1206,12 +1211,12 @@ def verify(self, ts, k):
                 parent = mutation.parent
                 if parent != tskit.NULL:
                     parent += parent_offset
-                tables.mutations.add_row(
-                    j,
-                    node=mutation.node,
-                    time=mutation.time,
-                    parent=parent,
-                    derived_state=mutation.derived_state,
+                tables.mutations.append(
+                    replace(
+                        mutation,
+                        site=j,
+                        parent=parent,
+                    )
                 )
 
         ts2 = tables.tree_sequence()

python/tests/test_tables.py

@@ -386,6 +386,42 @@ def test_add_row_round_trip(self):
                 t2.add_row(**dataclasses.asdict(row))
             assert t1 == t2
 
+    def test_append_row(self):
+        for num_rows in [0, 10, 100]:
+            table = self.table_class()
+            for j, row in enumerate(self.make_transposed_input_data(num_rows)):
+                k = table.append(table.row_class(**row))
+                assert k == j
+            for colname, input_array in self.make_input_data(num_rows).items():
+                output_array = getattr(table, colname)
+                assert input_array.shape == output_array.shape
+                assert np.all(input_array == output_array)
+            table.clear()
+            assert table.num_rows == 0
+            assert len(table) == 0
+
+    def test_append_duck_type(self):
+        class Duck:
+            pass
+
+        table = self.table_class()
+        for j, row in enumerate(self.make_transposed_input_data(20)):
+            duck = Duck()
+            for k, v in row.items():
+                setattr(duck, k, v)
+            k = table.append(duck)
+            assert k == j
+        for colname, input_array in self.make_input_data(20).items():
+            output_array = getattr(table, colname)
+            assert np.array_equal(input_array, output_array)
+
+    def test_append_error(self):
+        class NotADuck:
+            pass
+
+        with pytest.raises(AttributeError, match="'NotADuck' object has no attribute"):
+            self.table_class().append(NotADuck())
+
     def test_set_columns_data(self):
         for num_rows in [0, 10, 100, 1000]:
             input_data = {col.name: col.get_input(num_rows) for col in self.columns}
@@ -1627,7 +1663,7 @@ def verify_edge_sort_offset(self, ts):
             all_edges = keep + reversed_edges
             tables.edges.clear()
             for e in all_edges:
-                tables.edges.add_row(e.left, e.right, e.parent, e.child)
+                tables.edges.append(e)
             # Verify that import fails for randomised edges
             with pytest.raises(_tskit.LibraryError):
                 tables.tree_sequence()
@@ -1638,7 +1674,7 @@ def verify_edge_sort_offset(self, ts):
             # Sorting from the correct index should give us back the original table.
             tables.edges.clear()
             for e in all_edges:
-                tables.edges.add_row(e.left, e.right, e.parent, e.child)
+                tables.edges.append(e)
             tables.sort(edge_start=start)
             # Verify the new and old edges are equal.
             assert edges == tables.edges