Add low-level fromdict and use in high-level code

Author: jeromekelleher

python/_tskitmodule.c

@@ -6054,6 +6054,28 @@ TableCollection_asdict(TableCollection *self, PyObject *args, PyObject *kwds)
     return ret;
 }
 
+static PyObject *
+TableCollection_fromdict(TableCollection *self, PyObject *args)
+{
+    int err;
+    PyObject *ret = NULL;
+    PyObject *dict = NULL;
+
+    if (TableCollection_check_state(self) != 0) {
+        goto out;
+    }
+    if (!PyArg_ParseTuple(args, "O!", &PyDict_Type, &dict)) {
+        goto out;
+    }
+    err = parse_table_collection_dict(self->tables, dict);
+    if (err != 0) {
+        goto out;
+    }
+    ret = Py_BuildValue("");
+out:
+    return ret;
+}
+
 static PyGetSetDef TableCollection_getsetters[] = {
     { .name = "individuals",
         .get = (getter) TableCollection_get_individuals,
@@ -6178,6 +6200,10 @@ static PyMethodDef TableCollection_methods[] = {
         .ml_meth = (PyCFunction) TableCollection_asdict,
         .ml_flags = METH_VARARGS | METH_KEYWORDS,
         .ml_doc = "Returns the table collection in dictionary encoding. " },
+    { .ml_name = "fromdict",
+        .ml_meth = (PyCFunction) TableCollection_fromdict,
+        .ml_flags = METH_VARARGS | METH_KEYWORDS,
+        .ml_doc = "Sets the state of this table collection from the specified dict" },
     { NULL } /* Sentinel */
 };
 

python/tests/test_lowlevel.py

@@ -417,13 +417,28 @@ def test_asdict(self):
             assert isinstance(d, dict)
             assert len(d) > 0
 
+    def test_fromdict(self):
+        for ts in self.get_example_tree_sequences():
+            tc1 = _tskit.TableCollection(sequence_length=ts.get_sequence_length())
+            ts.dump_tables(tc1)
+            d = tc1.asdict()
+            tc2 = _tskit.TableCollection(sequence_length=0)
+            tc2.fromdict(d)
+            assert tc1.equals(tc2)
+
     def test_asdict_bad_args(self):
         ts = msprime.simulate(10, random_seed=1242)
         tc = ts.tables._ll_tables
         for bad_type in [None, 0.1, "str"]:
             with pytest.raises(TypeError):
                 tc.asdict(force_offset_64=bad_type)
 
+    def test_fromdict_bad_args(self):
+        tc = _tskit.TableCollection(0)
+        for bad_type in [None, 0.1, "str"]:
+            with pytest.raises(TypeError):
+                tc.fromdict(bad_type)
+
 
 class TestTableMethods:
     """

python/tskit/tables.py

@@ -2649,45 +2649,13 @@ def dump(self, file_or_path):
 
     # Unpickle support
     def __setstate__(self, state):
-        self.__init__(state["sequence_length"])
-        self.metadata_schema = tskit.parse_metadata_schema(state["metadata_schema"])
-        self.metadata = self.metadata_schema.decode_row(state["metadata"])
-        self.individuals.set_columns(**state["individuals"])
-        self.nodes.set_columns(**state["nodes"])
-        self.edges.set_columns(**state["edges"])
-        self.migrations.set_columns(**state["migrations"])
-        self.sites.set_columns(**state["sites"])
-        self.mutations.set_columns(**state["mutations"])
-        self.populations.set_columns(**state["populations"])
-        self.provenances.set_columns(**state["provenances"])
+        self.__init__()
+        self._ll_tables.fromdict(state)
 
     @classmethod
     def fromdict(self, tables_dict):
-        tables = TableCollection(tables_dict["sequence_length"])
-        try:
-            tables.metadata_schema = tskit.parse_metadata_schema(
-                tables_dict["metadata_schema"]
-            )
-        except KeyError:
-            pass
-        try:
-            tables.metadata = tables.metadata_schema.decode_row(tables_dict["metadata"])
-        except KeyError:
-            pass
-        tables.individuals.set_columns(**tables_dict["individuals"])
-        tables.nodes.set_columns(**tables_dict["nodes"])
-        tables.edges.set_columns(**tables_dict["edges"])
-        tables.migrations.set_columns(**tables_dict["migrations"])
-        tables.sites.set_columns(**tables_dict["sites"])
-        tables.mutations.set_columns(**tables_dict["mutations"])
-        tables.populations.set_columns(**tables_dict["populations"])
-        tables.provenances.set_columns(**tables_dict["provenances"])
-
-        # Indexes must be last as other wise the check for their consistency will fail
-        try:
-            tables.indexes = TableCollectionIndexes(**tables_dict["indexes"])
-        except KeyError:
-            pass
+        tables = TableCollection()
+        tables._ll_tables.fromdict(tables_dict)
         return tables
 
     def copy(self):